Add postgres support

Author: BenConstable9

deploy_ai_search/src/deploy_ai_search/text_2_sql_schema_store.py

@@ -65,6 +65,8 @@ def excluded_fields_for_database_engine(self):
             engine_specific_fields = ["Database"]
         elif self.database_engine == DatabaseEngine.DATABRICKS:
             engine_specific_fields = ["Catalog"]
+        elif self.database_engine == DatabaseEngine.POSTGRESQL:
+            engine_specific_fields = ["Database"]
 
         return [
             field

text_2_sql/autogen/pyproject.toml

@@ -11,7 +11,7 @@ dependencies = [
     "autogen-ext[azure,openai]==0.4.0.dev11",
     "grpcio>=1.68.1",
     "pyyaml>=6.0.2",
-    "text_2_sql_core[snowflake,databricks]",
+    "text_2_sql_core[snowflake,databricks,postgresql]",
 ]
 
 [dependency-groups]

text_2_sql/text_2_sql_core/pyproject.toml

@@ -46,6 +46,10 @@ databricks = [
     "databricks-sql-connector>=3.0.1",
     "pyarrow>=14.0.2,<17",
 ]
+postgresql = [
+    "psycopg>=3.2.3",
+]
+
 
 [build-system]
 requires = ["hatchling"]

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/postgresql_sql.py

@@ -0,0 +1,99 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from text_2_sql_core.connectors.sql import SqlConnector
+import psycopg
+from typing import Annotated
+import os
+import logging
+import json
+
+from text_2_sql_core.utils.database import DatabaseEngine
+
+
+class PostgresqlSqlConnector(SqlConnector):
+    def __init__(self):
+        super().__init__()
+
+        self.database_engine = DatabaseEngine.POSTGRESQL
+
+    async def query_execution(
+        self,
+        sql_query: Annotated[str, "The SQL query to run against the database."],
+        cast_to: any = None,
+        limit=None,
+    ) -> list[dict]:
+        """Run the SQL query against the PostgreSQL database asynchronously.
+
+        Args:
+        ----
+            sql_query (str): The SQL query to run against the database.
+
+        Returns:
+        -------
+            list[dict]: The results of the SQL query.
+        """
+        logging.info(f"Running query: {sql_query}")
+        results = []
+        connection_string = os.environ["Text2Sql__DatabaseConnectionString"]
+
+        # Establish an asynchronous connection to the PostgreSQL database
+        async with psycopg.AsyncConnection.connect(connection_string) as conn:
+            # Create an asynchronous cursor
+            async with conn.cursor() as cursor:
+                await cursor.execute(sql_query)
+
+                # Fetch column names
+                columns = [column[0] for column in cursor.description]
+
+                # Fetch rows based on the limit
+                if limit is not None:
+                    rows = await cursor.fetchmany(limit)
+                else:
+                    rows = await cursor.fetchall()
+
+                # Process the rows
+                for row in rows:
+                    if cast_to:
+                        results.append(cast_to.from_sql_row(row, columns))
+                    else:
+                        results.append(dict(zip(columns, row)))
+
+        logging.debug("Results: %s", results)
+        return results
+
+    async def get_entity_schemas(
+        self,
+        text: Annotated[
+            str,
+            "The text to run a semantic search against. Relevant entities will be returned.",
+        ],
+        excluded_entities: Annotated[
+            list[str],
+            "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
+        ] = [],
+        as_json: bool = True,
+    ) -> str:
+        """Gets the schema of a view or table in the SQL Database by selecting the most relevant entity based on the search term. Several entities may be returned.
+
+        Args:
+        ----
+            text (str): The text to run the search against.
+
+        Returns:
+            str: The schema of the views or tables in JSON format.
+        """
+
+        schemas = await self.ai_search_connector.get_entity_schemas(
+            text, excluded_entities
+        )
+
+        for schema in schemas:
+            schema["SelectFromEntity"] = ".".join([schema["Schema"], schema["Entity"]])
+
+            del schema["Entity"]
+            del schema["Schema"]
+
+        if as_json:
+            return json.dumps(schemas, default=str)
+        else:
+            return schemas

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/data_dictionary_creator.py

@@ -758,6 +758,8 @@ def excluded_fields_for_database_engine(self):
             engine_specific_fields = ["Database"]
         elif self.database_engine == DatabaseEngine.DATABRICKS:
             engine_specific_fields = ["Catalog"]
+        elif self.database_engine == DatabaseEngine.POSTGRESQL:
+            engine_specific_fields = ["Database"]
         else:
             engine_specific_fields = []
 

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/postgresql_data_dictionary_creator.py

@@ -0,0 +1,87 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from text_2_sql_core.data_dictionary.data_dictionary_creator import (
+    DataDictionaryCreator,
+    EntityItem,
+)
+import os
+
+from text_2_sql_core.utils.database import DatabaseEngine
+from text_2_sql_core.connectors.postgresql_sql import PostgresSqlConnector
+
+
+class PostgresqlDataDictionaryCreator(DataDictionaryCreator):
+    def __init__(self, **kwargs):
+        """A method to initialize the DataDictionaryCreator class."""
+        super().__init__(**kwargs)
+
+        self.database = os.environ["Text2Sql__DatabaseName"]
+        self.database_engine = DatabaseEngine.POSTGRESQL
+
+        self.sql_connector = PostgresSqlConnector()
+
+    @property
+    def extract_table_entities_sql_query(self) -> str:
+        """A property to extract table entities from a PostgreSQL database."""
+        return """SELECT
+            t.table_name AS entity,
+            t.table_schema AS entity_schema,
+            pg_catalog.col_description(c.oid, 0) AS definition
+        FROM
+            information_schema.tables t
+        JOIN
+            pg_catalog.pg_class c ON c.relname = t.table_name
+            AND c.relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = t.table_schema)
+        WHERE
+            t.table_type = 'BASE TABLE'
+        ORDER BY entity_schema, entity;"""
+
+    @property
+    def extract_view_entities_sql_query(self) -> str:
+        """A property to extract view entities from a PostgreSQL database."""
+        return """SELECT
+            v.view_name AS entity,
+            v.table_schema AS entity_schema,
+            pg_catalog.col_description(c.oid, 0) AS definition
+        FROM
+            information_schema.views v
+        JOIN
+            pg_catalog.pg_class c ON c.relname = v.view_name
+            AND c.relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = v.table_schema)
+        ORDER BY entity_schema, entity;"""
+
+    def extract_columns_sql_query(self, entity: EntityItem) -> str:
+        """A property to extract column information from a PostgreSQL database."""
+        return f"""SELECT
+            c.column_name AS name,
+            c.data_type AS data_type,
+            pg_catalog.col_description(t.oid, c.ordinal_position) AS definition
+        FROM
+            information_schema.columns c
+        JOIN
+            pg_catalog.pg_class t ON t.relname = c.table_name
+            AND t.relnamespace = (SELECT oid FROM pg_catalog.pg_namespace WHERE nspname = c.table_schema)
+        WHERE
+            c.table_schema = '{entity.entity_schema}'
+            AND c.table_name = '{entity.name}'
+        ORDER BY c.ordinal_position;"""
+
+    @property
+    def extract_entity_relationships_sql_query(self) -> str:
+        """A property to extract entity relationships from a PostgreSQL database."""
+        return """SELECT
+            tc.table_schema AS entity_schema,
+            tc.table_name AS entity,
+            rc.unique_constraint_schema AS foreign_entity_schema,
+            rc.unique_constraint_name AS foreign_entity_constraint,
+            rc.constraint_name AS foreign_key_constraint
+        FROM
+            information_schema.referential_constraints rc
+        JOIN
+            information_schema.table_constraints tc
+            ON rc.constraint_schema = tc.constraint_schema
+            AND rc.constraint_name = tc.constraint_name
+        WHERE
+            tc.constraint_type = 'FOREIGN KEY'
+        ORDER BY
+            entity_schema, entity, foreign_entity_schema, foreign_entity_constraint;"""

text_2_sql/text_2_sql_core/src/text_2_sql_core/utils/database.py

@@ -7,3 +7,4 @@ class DatabaseEngine(StrEnum):
     DATABRICKS = "DATABRICKS"
     SNOWFLAKE = "SNOWFLAKE"
     TSQL = "TSQL"
+    POSTGRESQL = "POSTGRESQL"