Source code for nsaph.util.query_builder

#  Copyright (c) 2024.  Harvard University
#
#   Developed by Research Software Engineering,
#   Harvard University Research Computing and Data (RCD) Services.
#
#   Author: Michael A Bouzinier
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#          http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#
#
from typing import List, Tuple

SQL_TABLE_COLUMNS = """
SELECT
    COLUMN_NAME
FROM 
    information_schema.columns
WHERE 
    TABLE_SCHEMA = '{schema}' 
    AND TABLE_NAME = '{table}'

"""


SQL_SCHEMA_TABLES = """
SELECT 
    table_schema||'.'||TABLE_NAME 
FROM 
    information_schema.tables
WHERE 
    table_schema = '{schema}'

"""

[docs]class QueryBuilder: def __init__(self, connection): self.connection = connection self.tables: List[str] = [] self.select: List[str] = []
[docs] def add_table(self, table: str): self.tables.append(table) self.select.extend(self.get_columns(table)) return self
[docs] @staticmethod def split(t:str) -> Tuple[str]: return tuple(t.split('.'))
[docs] def get_columns(self, table_fqn: str) -> List[str]: schema, table = self.split(table_fqn) sql = SQL_TABLE_COLUMNS.format(schema=schema, table=table) with (self.connection.cursor()) as cursor: cursor.execute(sql) columns = [table + '.' + r[0] for r in cursor] return columns
[docs] def query(self) -> str: sql = "SELECT\n" sql += ",\n\t".join(self.select) sql += "\nFROM\n\t" sql += ", ".join(self.tables) return sql
[docs] def clear(self): self.tables.clear() self.select.clear() return self
[docs] @classmethod def get_tables(cls, connection, schema: str): with (connection.cursor()) as cursor: cursor.execute(SQL_SCHEMA_TABLES.format(schema=schema)) tables = [r[0] for r in cursor] return tables