Function create_query_executor

web/pgadmin/llm/reports/generator.py:24–77

Create a query executor function for the pipeline. Args: conn: Database connection object. Returns: A callable that executes queries by ID.

(conn)

def create_query_executor(conn) -> callable:
    """Create a query executor function for the pipeline.

    Args:
        conn: Database connection object.

    Returns:
        A callable that executes queries by ID.
    """
    def executor(query_id: str, context: dict) -> dict[str, Any]:
        """Execute a query by ID.

        Args:
            query_id: The query identifier from QUERIES registry.
            context: Execution context (may contain schema_id for filtering).

        Returns:
            Dictionary with query results.
        """
        query_def = QUERIES.get(query_id)
        if not query_def:
            return {'error': f'Unknown query: {query_id}', 'rows': []}

        sql = query_def['sql']

        # Check if query requires an extension
        required_ext = query_def.get('requires_extension')
        if required_ext:
            check_sql = """
                SELECT EXISTS (
                    SELECT 1 FROM pg_catalog.pg_extension WHERE extname = %s
                ) as available
            """
            status, result = conn.execute_dict(check_sql, [required_ext])
            if not (status and result and
                    result.get('rows', [{}])[0].get('available', False)):
                return {
                    'note': f"Extension '{required_ext}' not installed",
                    'rows': []
                }

        # Handle schema-scoped queries
        schema_id = context.get('schema_id')
        if schema_id and '%s' in sql:
            status, result = conn.execute_dict(sql, [schema_id])
        else:
            status, result = conn.execute_dict(sql)

        if status and result:
            return {'rows': result.get('rows', [])}
        else:
            return {'error': 'Query failed', 'rows': []}

    return executor


def generate_report_streaming(