[DH-5733] Support schemas column to add a db connection

Author: jcjc712

Diff for: dataherald/api/fastapi.py

@@ -70,6 +70,9 @@
     SQLInjectionError,
 )
 from dataherald.sql_database.models.types import DatabaseConnection
+from dataherald.sql_database.services.database_connection import (
+    DatabaseConnectionService,
+)
 from dataherald.types import (
     BaseLLM,
     CancelFineTuningRequest,
@@ -173,27 +176,9 @@ def create_database_connection(
         self, database_connection_request: DatabaseConnectionRequest
     ) -> DatabaseConnectionResponse:
         try:
-            db_connection = DatabaseConnection(
-                alias=database_connection_request.alias,
-                connection_uri=database_connection_request.connection_uri.strip(),
-                path_to_credentials_file=database_connection_request.path_to_credentials_file,
-                llm_api_key=database_connection_request.llm_api_key,
-                use_ssh=database_connection_request.use_ssh,
-                ssh_settings=database_connection_request.ssh_settings,
-                file_storage=database_connection_request.file_storage,
-                metadata=database_connection_request.metadata,
-            )
-            sql_database = SQLDatabase.get_sql_engine(db_connection, True)
-
-            # Get tables and views and create table-descriptions as NOT_SCANNED
-            db_connection_repository = DatabaseConnectionRepository(self.storage)
-
-            scanner_repository = TableDescriptionRepository(self.storage)
             scanner = self.system.instance(Scanner)
-
-            tables = sql_database.get_tables_and_views()
-            db_connection = db_connection_repository.insert(db_connection)
-            scanner.create_tables(tables, str(db_connection.id), scanner_repository)
+            db_connection_service = DatabaseConnectionService(scanner, self.storage)
+            db_connection = db_connection_service.create(database_connection_request)
         except Exception as e:
             # Encrypt sensible values
             fernet_encrypt = FernetEncrypt()
@@ -209,7 +194,6 @@ def create_database_connection(
             return error_response(
                 e, database_connection_request.dict(), "invalid_database_connection"
             )
-
         return DatabaseConnectionResponse(**db_connection.dict())
 
     @override

Diff for: dataherald/db_scanner/__init__.py

@@ -34,6 +34,7 @@ def create_tables(
         self,
         tables: list[str],
         db_connection_id: str,
+        schema: str,
         repository: TableDescriptionRepository,
         metadata: dict = None,
     ) -> None:

Diff for: dataherald/db_scanner/models/types.py

@@ -31,6 +31,7 @@ class TableDescriptionStatus(Enum):
 class TableDescription(BaseModel):
     id: str | None
     db_connection_id: str
+    schema_name: str | None
     table_name: str
     description: str | None
     table_schema: str | None

Diff for: dataherald/db_scanner/sqlalchemy.py

@@ -44,13 +44,15 @@ def create_tables(
         self,
         tables: list[str],
         db_connection_id: str,
+        schema: str,
         repository: TableDescriptionRepository,
         metadata: dict = None,
     ) -> None:
         for table in tables:
             repository.save_table_info(
                 TableDescription(
                     db_connection_id=db_connection_id,
+                    schema_name=schema,
                     table_name=table,
                     status=TableDescriptionStatus.NOT_SCANNED.value,
                     metadata=metadata,

Diff for: dataherald/sql_database/models/types.py

@@ -96,6 +96,7 @@ class DatabaseConnection(BaseModel):
     dialect: SupportedDialects | None
     use_ssh: bool = False
     connection_uri: str | None
+    schemas: list[str] | None
     path_to_credentials_file: str | None
     llm_api_key: str | None = None
     ssh_settings: SSHSettings | None = None

Diff for: dataherald/sql_database/services/__init__.py

@@ -0,0 +1,119 @@
+import re
+
+from sqlalchemy import inspect
+
+from dataherald.db import DB
+from dataherald.db_scanner import Scanner
+from dataherald.db_scanner.repository.base import TableDescriptionRepository
+from dataherald.repositories.database_connections import DatabaseConnectionRepository
+from dataherald.sql_database.base import SQLDatabase
+from dataherald.sql_database.models.types import DatabaseConnection
+from dataherald.types import DatabaseConnectionRequest
+from dataherald.utils.encrypt import FernetEncrypt
+
+
+class DatabaseConnectionService:
+    def __init__(self, scanner: Scanner, storage: DB):
+        self.scanner = scanner
+        self.storage = storage
+
+    def get_current_schema(self, database_connection: DatabaseConnection) -> list[str]:
+        sql_database = SQLDatabase.get_sql_engine(database_connection, True)
+        inspector = inspect(sql_database.engine)
+        if inspector.default_schema_name and database_connection.dialect not in [
+            "mssql",
+            "mysql",
+            "clickhouse",
+            "duckdb",
+        ]:
+            return [inspector.default_schema_name]
+        if database_connection.dialect == "bigquery":
+            pattern = r"([^:/]+)://([^/]+)/([^/]+)(\?[^/]+)"
+            match = re.match(pattern, str(sql_database.engine.url))
+            if match:
+                return [match.group(3)]
+        elif database_connection.dialect == "databricks":
+            pattern = r"&schema=([^&]*)"
+            match = re.search(pattern, str(sql_database.engine.url))
+            if match:
+                return [match.group(1)]
+        return ["default"]
+
+    def remove_schema_in_uri(self, connection_uri: str, dialect: str) -> str:
+        if dialect in ["snowflake"]:
+            pattern = r"([^:/]+)://([^:]+):([^@]+)@([^:/]+)(?::(\d+))?/([^/]+)"
+            match = re.match(pattern, connection_uri)
+            if match:
+                return match.group(0)
+        if dialect in ["bigquery"]:
+            pattern = r"([^:/]+)://([^/]+)"
+            match = re.match(pattern, connection_uri)
+            if match:
+                return match.group(0)
+        elif dialect in ["databricks"]:
+            pattern = r"&schema=[^&]*"
+            return re.sub(pattern, "", connection_uri)
+        elif dialect in ["postgresql"]:
+            pattern = r"\?options=-csearch_path" r"=[^&]*"
+            return re.sub(pattern, "", connection_uri)
+        return connection_uri
+
+    def add_schema_in_uri(self, connection_uri: str, schema: str, dialect: str) -> str:
+        connection_uri = self.remove_schema_in_uri(connection_uri, dialect)
+        if dialect in ["snowflake", "bigquery"]:
+            return f"{connection_uri}/{schema}"
+        if dialect in ["databricks"]:
+            return f"{connection_uri}&schema={schema}"
+        if dialect in ["postgresql"]:
+            return f"{connection_uri}?options=-csearch_path={schema}"
+        return connection_uri
+
+    def create(
+        self, database_connection_request: DatabaseConnectionRequest
+    ) -> DatabaseConnection:
+        database_connection = DatabaseConnection(
+            alias=database_connection_request.alias,
+            connection_uri=database_connection_request.connection_uri.strip(),
+            schemas=database_connection_request.schemas,
+            path_to_credentials_file=database_connection_request.path_to_credentials_file,
+            llm_api_key=database_connection_request.llm_api_key,
+            use_ssh=database_connection_request.use_ssh,
+            ssh_settings=database_connection_request.ssh_settings,
+            file_storage=database_connection_request.file_storage,
+            metadata=database_connection_request.metadata,
+        )
+        if not database_connection.schemas:
+            database_connection.schemas = self.get_current_schema(database_connection)
+
+        schemas_and_tables = {}
+        fernet_encrypt = FernetEncrypt()
+
+        if database_connection.schemas:
+            for schema in database_connection.schemas:
+                database_connection.connection_uri = fernet_encrypt.encrypt(
+                    self.add_schema_in_uri(
+                        database_connection_request.connection_uri.strip(),
+                        schema,
+                        str(database_connection.dialect),
+                    )
+                )
+                sql_database = SQLDatabase.get_sql_engine(database_connection, True)
+                schemas_and_tables[schema] = sql_database.get_tables_and_views()
+
+        # Connect db
+        db_connection_repository = DatabaseConnectionRepository(self.storage)
+        database_connection.connection_uri = fernet_encrypt.encrypt(
+            self.remove_schema_in_uri(
+                database_connection_request.connection_uri.strip(),
+                str(database_connection.dialect),
+            )
+        )
+        db_connection = db_connection_repository.insert(database_connection)
+
+        scanner_repository = TableDescriptionRepository(self.storage)
+        # Add created tables
+        for schema, tables in schemas_and_tables.items():
+            self.scanner.create_tables(
+                tables, str(db_connection.id), schema, scanner_repository
+            )
+        return DatabaseConnection(**db_connection.dict())

Diff for: dataherald/sql_database/services/database_connection.py

@@ -87,6 +87,7 @@ class DatabaseConnectionRequest(BaseModel):
     alias: str
     use_ssh: bool = False
     connection_uri: str
+    schemas: list[str] | None
     path_to_credentials_file: str | None
     llm_api_key: str | None
     ssh_settings: SSHSettings | None