Communicate easily with a sqp database

interface_sql.py

import psycopg2                         # pip install psycopg2-binary
from psycopg2 import sql
from typing import Optional, Any, Tuple, Union, List


class InterfaceSQL:
    """
    An interface for building and executing SQL queries against a PostgreSQL database.
    Query building and examination are handled by static methods that are connection-independent.
    """

    # --- Query Building & Examination (Static Methods - No Connection Needed) ---

    @staticmethod
    def create_table(table_name: str,
                     column: list[tuple[str, str]],
                     namespace: str = 'public',
                     if_not_exists: bool = False) -> sql.Composed:
        column_definitions = [sql.SQL("{} {}").format(sql.Identifier(n), sql.SQL(t)) for n, t in column]
        formatted_columns = sql.SQL(',\n\t').join(column_definitions)
        query_template = sql.SQL("CREATE TABLE {if_not_exists}{schema}.{table} (\n\t{columns}\n);")
        return query_template.format(
            if_not_exists=sql.SQL("IF NOT EXISTS ") if if_not_exists else sql.SQL(""),
            schema=sql.Identifier(namespace),
            table=sql.Identifier(table_name),
            columns=formatted_columns
        )


    @staticmethod
    def drop_table(namespace: str, table_name: str) -> sql.Composed:
        return sql.SQL("DROP TABLE {}.{};").format(sql.Identifier(namespace), sql.Identifier(table_name))


    @staticmethod
    def select_from(table_name: str,
                    namespace: str = 'public',
                    values: list[str] | None = None,
                    where: dict[str, Any] | None = None,
                    limit: int | None = None,
                    count: bool = False) -> Tuple[sql.Composed, tuple]:
        if count:
            select_clause = sql.SQL("COUNT(*)")
        elif values:
            select_clause = sql.SQL(', ').join(map(sql.Identifier, values))
        else:
            select_clause = sql.SQL('*')
        query_parts = [
            sql.SQL("SELECT {select} FROM {schema}.{table}").format(
                select=select_clause,
                schema=sql.Identifier(namespace),
                table=sql.Identifier(table_name)
            )
        ]
        params = []
        if where:
            where_clauses = [sql.SQL("{} = %s").format(sql.Identifier(k)) for k in where.keys()]
            params.extend(where.values())
            query_parts.append(sql.SQL("WHERE ") + sql.SQL(" AND ").join(where_clauses))
        if limit is not None:
            if not isinstance(limit, int) or limit < 0:
                raise ValueError("Limit must be a non-negative integer.")
            query_parts.append(sql.SQL("LIMIT %s"))
            params.append(limit)
        return (sql.SQL(' ').join(query_parts) + sql.SQL(';'), tuple(params))


    @staticmethod
    def examine_query_offline(query_object: sql.Composed, variables: tuple = None) -> str:
        """
        Recursively renders a human-readable string from a query object WITHOUT a connection.
        This is for inspection only and is not guaranteed to be executable.
        """
        def _flatten_recursively(obj) -> List[str]:
            parts = []
            if isinstance(obj, sql.Composed):
                for part in obj:
                    parts.extend(_flatten_recursively(part))
            elif isinstance(obj, sql.SQL):
                parts.append(obj.string)
            elif isinstance(obj, sql.Identifier):
                # THE FIX: Access the string inside the Identifier object's `strings` list.
                # This correctly extracts "id" instead of showing "Identifier('id')".
                string_inside = ".".join([f'"{s}"' for s in obj.strings])
                parts.append(string_inside)
            else:
                parts.append(str(obj))
            return parts

        string_parts = _flatten_recursively(query_object)
        rendered_sql = "".join(string_parts)
        disclaimer = "-- [OFFLINE PREVIEW] FOR HUMAN INSPECTION ONLY. NOT GUARANTEED TO BE EXECUTABLE.\n"
        preview = f"{disclaimer}SQL Template: {rendered_sql}"
        if variables:
            preview += f"\nParameters:   {variables}"
        return preview


    # --- Connection and Execution (Instance Methods) ---

    def __init__(self, database: str, host: str, port: str, user: str, password: str):
        self.conn = psycopg2.connect(database=database, host=host, port=port, user=user, password=password)

    def get_query_string(self, query_object: sql.Composed, variables: tuple | None = None) -> str:
        """Renders a query object into a final, executable string using a live connection."""
        with self.conn.cursor() as cur:
            return cur.mogrify(query_object, variables).decode(self.conn.encoding)

    def send_query(self, query: Union[str, sql.Composed], variables: tuple | None = None, fetch: bool = False) -> Optional[list[tuple]]:
        """Sends a command to the database. Accepts a string or a query object."""
        with self.conn.cursor() as cur:
            cur.execute(query, variables)
            self.conn.commit()
            if fetch:
                return cur.fetchall()
        return None

    def close(self):
        """Closes the database connection."""
        self.conn.close()


if __name__ == "__main__":

    print("--- Running User-Provided Test Cases (Offline Preview) ---")

    # Test 1: Select specific columns
    print("\n--- Test 1: SELECT specific columns ---")
    query_obj, params = InterfaceSQL.select_from('test1', values=['id', 'class_vector'])
    print(InterfaceSQL.examine_query_offline(query_obj, params))

    # Test 2: Select with a limit
    print("\n--- Test 2: SELECT with a LIMIT ---")
    query_obj, params = InterfaceSQL.select_from('test1', limit=2)
    print(InterfaceSQL.examine_query_offline(query_obj, params))

    # Test 3: Select with WHERE IS NULL
    print("\n--- Test 3: SELECT with WHERE IS NULL ---")
    query_obj, params = InterfaceSQL.select_from('test1', where={'filepath': None})
    print(InterfaceSQL.examine_query_offline(query_obj, params))

    # Test 4: Select with COUNT
    print("\n--- Test 4: SELECT with COUNT ---")
    # Note that `values` is correctly ignored by the function when `count=True`.
    query_obj, params = InterfaceSQL.select_from('test1', count=True, values=['id'])
    print(InterfaceSQL.examine_query_offline(query_obj, params))

    # Test 5: Select with COUNT, WHERE, and LIMIT
    print("\n--- Test 5: SELECT with COUNT, WHERE, and LIMIT ---")
    query_obj, params = InterfaceSQL.select_from('test1', where={'filepath': None}, limit=2, count=True)
    print(InterfaceSQL.examine_query_offline(query_obj, params))

    # Test 6: Select with specific columns, WHERE, and LIMIT
    print("\n--- Test 6: SELECT with specific columns, WHERE, and LIMIT ---")
    query_obj, params = InterfaceSQL.select_from(
        'test1',
        values=['id', 'class_vector'],
        where={'filepath': None},
        limit=2
    )
    print(InterfaceSQL.examine_query_offline(query_obj, params))

    print("\n\n--- Offline tests complete. You can now use these objects with a live connection. ---")


    # --- Part 2: Online Rendering and Execution ---
    print("\n\n--- 3. Connecting to database to render and execute ---")
    try:
        interface = InterfaceSQL(
            database='vectorexample', host='localhost', port='5432', user='postgres', password='password'
        )
        print("Database connection successful.")

        print("\n--- 4. Getting the true, executable query string ---")
        executable_string = interface.get_query_string(query_obj, params)
        print(executable_string)

        interface.close()

    except psycopg2.OperationalError as e:
        print("\n--- DATABASE CONNECTION FAILED ---")
        print(f"Error: {e}")