Function get_table_columns

web/pgadmin/llm/tools/database.py:519–626

Get column information for a specific table. Uses pgAdmin's SQL templates for version-aware column listing. Args: sid: Server ID did: Database ID schema_name: Schema name table_name: Table name Returns: Dictionary containing column informat

(
    sid: int,
    did: int,
    schema_name: str,
    table_name: str
)

def get_table_columns(
    sid: int,
    did: int,
    schema_name: str,
    table_name: str
) -> dict:
    """
    Get column information for a specific table.

    Uses pgAdmin's SQL templates for version-aware column listing.

    Args:
        sid: Server ID
        did: Database ID
        schema_name: Schema name
        table_name: Table name

    Returns:
        Dictionary containing column information
    """
    conn_id = f"llm_{secrets.choice(range(1, 9999999))}"
    manager = None

    try:
        manager, conn = _get_connection(sid, did, conn_id)
        status, error = _connect_readonly(manager, conn, conn_id)
        if not status:
            raise DatabaseToolError(f"Connection failed: {error}",
                                    code="CONNECTION_ERROR")

        sversion = manager.sversion or 0
        driver = get_driver(config.PG_DEFAULT_DRIVER)

        # Use qtLiteral for safe SQL escaping
        schema_lit = driver.qtLiteral(schema_name, conn)
        table_lit = driver.qtLiteral(table_name, conn)

        # Get table OID first
        oid_sql = f"""
            SELECT c.oid
            FROM pg_catalog.pg_class c
            JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            WHERE c.relname = {table_lit}
              AND n.nspname = {schema_lit}
        """

        status, _ = conn.execute_void("BEGIN TRANSACTION READ ONLY")
        if not status:
            raise DatabaseToolError("Failed to start transaction",
                                    code="TRANSACTION_ERROR")

        try:
            status, oid_res = conn.execute_dict(oid_sql)
            if not status or not oid_res.get('rows'):
                raise DatabaseToolError(
                    f"Table {schema_name}.{table_name} not found",
                    code="NOT_FOUND"
                )