diff --git a/private_gpt/components/database/connection_factory.py b/private_gpt/components/database/connection_factory.py new file mode 100644 index 00000000..c980a39f --- /dev/null +++ b/private_gpt/components/database/connection_factory.py @@ -0,0 +1,118 @@ +import enum +import functools +import re +from typing import Any +from urllib.parse import parse_qs, unquote, urlparse + +from sqlalchemy import Engine, create_engine + +from private_gpt.utils.dependencies import format_missing_dependency_message + + +class DatabaseDialect(enum.StrEnum): + """Database family, independent of the exact SQLAlchemy driver name. + + Extend this (and _DIALECT_NAME_MAP below) when adding support for a new + database family -- consumers should switch on this enum rather than + comparing raw dialect-name strings, so a new family is a single place to + wire up instead of a set of ad hoc string-list checks scattered around. + """ + + POSTGRES = "postgres" + MYSQL = "mysql" + MSSQL = "mssql" + DB2 = "db2" + UNKNOWN = "unknown" + + +_DIALECT_NAME_MAP: dict[str, DatabaseDialect] = { + "postgresql": DatabaseDialect.POSTGRES, + "postgres": DatabaseDialect.POSTGRES, + "mysql": DatabaseDialect.MYSQL, + "mssql": DatabaseDialect.MSSQL, + "microsoft": DatabaseDialect.MSSQL, + "db2": DatabaseDialect.DB2, + "ibm_db_sa": DatabaseDialect.DB2, +} + +_URL_PASSWORD_RE = re.compile(r"(://[^:/?#@]+:)([^@/?#]+)(@)") +_DB2_BRACED_PWD_RE = re.compile(r"(PWD=)\{((?:[^}]|}})*)\}(;)", re.IGNORECASE) +_DB2_PLAIN_PWD_RE = re.compile(r"(PWD=)(?!\{)([^;]*)(;)", re.IGNORECASE) +_DB2_VALUE_NEEDS_ESCAPING_RE = re.compile(r"[;{}]|^\s|\s$") + + +@functools.cache +def _load_ibm_db_dbi() -> Any: + try: + import ibm_db_dbi # type: ignore[import-not-found,import-untyped] + except ImportError as e: + raise ImportError( + format_missing_dependency_message( + "DB2 database query", + extras=("database-db2", "database"), + ) + ) from e + + return ibm_db_dbi + + +def classify_dialect(dialect_name: str | None) -> DatabaseDialect: + """Classify a dialect/scheme identifier into a DatabaseDialect.""" + return _DIALECT_NAME_MAP.get((dialect_name or "").lower(), DatabaseDialect.UNKNOWN) + + +def is_db2_connection_string(connection_string: str) -> bool: + """Whether a SQLAlchemy-style connection string targets DB2.""" + scheme = connection_string.split("://", 1)[0].split("+", 1)[0] + return classify_dialect(scheme) is DatabaseDialect.DB2 + + +def escape_db2_cli_value(value: str) -> str: + """Escape a value for a DB2 CLI/db2dsdriver keyword=value; connection string.""" + if not _DB2_VALUE_NEEDS_ESCAPING_RE.search(value): + return value + return "{" + value.replace("}", "}}") + "}" + + +def build_db2_native_connection_string(connection_string: str) -> str: + """Build a native DB2 CLI connection string from a SQLAlchemy-style URL.""" + remainder = connection_string.split("://", 1)[-1] + parsed = urlparse(f"//{remainder}") + params = parse_qs(parsed.query, keep_blank_values=True) + + # urlparse does not percent-decode username/password/path, unlike parse_qs + database = unquote(parsed.path.lstrip("/")) + username = unquote(parsed.username or "") + password = unquote(parsed.password or "") + + parts = [f"DATABASE={database}", f"HOSTNAME={parsed.hostname or ''}"] + if parsed.port: + parts.append(f"PORT={parsed.port}") + parts.append("PROTOCOL=TCPIP") + parts.append(f"UID={escape_db2_cli_value(username)}") + parts.append(f"PWD={escape_db2_cli_value(password)}") + + for key, values in params.items(): + parts.append(f"{key.upper()}={escape_db2_cli_value(values[-1])}") + + return ";".join(parts) + ";" + + +def mask_connection_secrets(text: str) -> str: + """Mask any password embedded in a connection string or error message.""" + masked = _URL_PASSWORD_RE.sub(r"\1***\3", text) + masked = _DB2_BRACED_PWD_RE.sub(r"\1{***}\3", masked) + masked = _DB2_PLAIN_PWD_RE.sub(r"\1***\3", masked) + return masked + + +def create_engine_for_connection_string(connection_string: str) -> Engine: + """Create a SQLAlchemy engine for a connection string.""" + if is_db2_connection_string(connection_string): + return create_engine( + "db2+ibm_db://", + creator=lambda: _load_ibm_db_dbi().connect( + build_db2_native_connection_string(connection_string), "", "" + ), + ) + return create_engine(connection_string) diff --git a/private_gpt/components/database/function_inspector.py b/private_gpt/components/database/function_inspector.py index 7a35d3c3..c336c164 100644 --- a/private_gpt/components/database/function_inspector.py +++ b/private_gpt/components/database/function_inspector.py @@ -1,6 +1,7 @@ from sqlalchemy import text from sqlalchemy.exc import SQLAlchemyError +from private_gpt.components.database.connection_factory import DatabaseDialect from private_gpt.components.database.inspected_schema import ( InspectedFunction, InspectedFunctionParams, @@ -17,13 +18,13 @@ class DatabaseFunctionsInspector(DatabaseObjectInspector): def get_objects(self, schema: str) -> list[InspectedFunction]: # type: ignore[override] try: - if self._db_type in ["mssql", "microsoft"]: + if self.dialect == DatabaseDialect.MSSQL: return self._get_sqlserver_functions(schema) - if self._db_type in ["db2", "ibm_db_sa"]: + if self.dialect == DatabaseDialect.DB2: return self._get_db2_functions(schema) # TODO: Postgres support is not complete - # elif self._db_type == "postgresql": + # elif self.dialect == DatabaseDialect.POSTGRES: # return self._get_postgresql_functions(schema) else: return [] diff --git a/private_gpt/components/database/inspector_interface.py b/private_gpt/components/database/inspector_interface.py index cc31cd46..e9770865 100644 --- a/private_gpt/components/database/inspector_interface.py +++ b/private_gpt/components/database/inspector_interface.py @@ -1,7 +1,13 @@ import enum from abc import ABC, abstractmethod -from sqlalchemy import Connection, Engine, create_engine +from sqlalchemy import Connection, Engine + +from private_gpt.components.database.connection_factory import ( + DatabaseDialect, + classify_dialect, + create_engine_for_connection_string, +) class DatabaseObjectType(enum.StrEnum): @@ -30,6 +36,7 @@ class DatabaseObjectInspector(ABC): _is_readonly: bool _connection: Connection | None = None connection_string: str + dialect: DatabaseDialect def __init__( self, @@ -40,6 +47,7 @@ class DatabaseObjectInspector(ABC): ): self._engine = engine self._db_type = engine.dialect.name.lower() if engine else "unknown" + self.dialect = classify_dialect(self._db_type) self._is_readonly = is_readonly self.connection_string = connection_string self._connection = connection @@ -56,7 +64,7 @@ class DatabaseObjectInspector(ABC): if self._connection and not self._connection.closed: return self._connection - self._engine = create_engine( + self._engine = create_engine_for_connection_string( self.connection_string # TODO: SSL Certificates ) diff --git a/private_gpt/components/database/procedure_inspector.py b/private_gpt/components/database/procedure_inspector.py index 314711c1..9767b421 100644 --- a/private_gpt/components/database/procedure_inspector.py +++ b/private_gpt/components/database/procedure_inspector.py @@ -1,6 +1,7 @@ from sqlalchemy import text from sqlalchemy.exc import SQLAlchemyError +from private_gpt.components.database.connection_factory import DatabaseDialect from private_gpt.components.database.inspected_schema import ( InspectedProcedure, InspectedProcedureParams, @@ -17,12 +18,12 @@ class DatabaseProcedureInspector(DatabaseObjectInspector): def get_objects(self, schema: str) -> list[InspectedProcedure]: # type: ignore[override] try: - if self._db_type in ["mssql", "microsoft"]: + if self.dialect == DatabaseDialect.MSSQL: return self._get_sqlserver_procedures(schema) - if self._db_type in ["db2", "ibm_db_sa"]: + if self.dialect == DatabaseDialect.DB2: return self._get_db2_procedures(schema) # TODO: postgres support is not complete - # elif self._db_type == "postgresql": + # elif self.dialect == DatabaseDialect.POSTGRES: # return self._get_postgresql_procedures(schema) else: return [] diff --git a/private_gpt/components/tabular/database_query_generator.py b/private_gpt/components/tabular/database_query_generator.py index 4a8e7a72..ec67d16f 100644 --- a/private_gpt/components/tabular/database_query_generator.py +++ b/private_gpt/components/tabular/database_query_generator.py @@ -16,7 +16,7 @@ from Levenshtein import distance from llama_index.core.base.llms.types import ChatMessage, MessageRole from pandas import DataFrame from pydantic import BaseModel, Field -from sqlalchemy import Connection, Engine, create_engine, inspect, text +from sqlalchemy import Connection, Engine, inspect, text from sqlalchemy.exc import ProgrammingError, SQLAlchemyError from private_gpt.components.chat.models.chat_config_models import ( @@ -27,6 +27,13 @@ from private_gpt.components.chat.models.chat_config_models import ( from private_gpt.components.chat.processors.chat_history.memory.utils.content import ( messages_to_history_str, ) +from private_gpt.components.database.connection_factory import ( + DatabaseDialect, + build_db2_native_connection_string, + classify_dialect, + create_engine_for_connection_string, + mask_connection_secrets, +) from private_gpt.components.database.function_inspector import ( DatabaseFunctionsInspector, ) @@ -428,17 +435,17 @@ class DatabaseQueryGenerator: """ try: conn = self._ensure_connected() - match self._engine.dialect.name.lower() if self._engine else "unknown": - case "db2" | "ibm_db_sa": - conn.execute(text("SELECT 1 FROM SYSIBM.SYSDUMMY1")) - case _: - conn.execute(text("SELECT 1")) + dialect = classify_dialect(self._engine.dialect.name if self._engine else None) + if dialect == DatabaseDialect.DB2: + conn.execute(text("SELECT 1 FROM SYSIBM.SYSDUMMY1")) + else: + conn.execute(text("SELECT 1")) return None except (SQLAlchemyError, ProgrammingError) as e: - logging.error(e) + logging.error(mask_connection_secrets(str(e))) return f"Failed to connect to database '{self._extract_database_name()}'" except Exception as e: - logging.error(e) + logging.error(mask_connection_secrets(str(e))) return f"Failed to connect to database '{self._extract_database_name()}'" finally: self._disconnect() @@ -704,7 +711,7 @@ class DatabaseQueryGenerator: if self._connection and not self._connection.closed: return self._connection - self._engine = create_engine( + self._engine = create_engine_for_connection_string( self.connection_string # TODO: SSL Certificates ) @@ -724,12 +731,11 @@ class DatabaseQueryGenerator: value = self._engine.dialect.name.lower() if self._engine else self._dialect system_prompt = f"The database SQL dialect is: {value}." - match value: - case "db2" | "ibm_db_sa": - system_prompt += ( - " For DB2 procedures: input parameters (Params) use literal values, output parameters (Returns) use '?' placeholders." - " Example with 1 input + 3 outputs: CALL schema.procedure_name('value1', ?, ?, ?);" - ) + if classify_dialect(value) == DatabaseDialect.DB2: + system_prompt += ( + " For DB2 procedures: input parameters (Params) use literal values, output parameters (Returns) use '?' placeholders." + " Example with 1 input + 3 outputs: CALL schema.procedure_name('value1', ?, ?, ?);" + ) return system_prompt async def generate_sql_query( @@ -1200,7 +1206,7 @@ class DatabaseQueryGenerator: return QueryResult( query=call_statement, error=ErrorQueryResult( - description=f"DB2 procedure execution failed: {e!s}", + description=f"DB2 procedure execution failed: {mask_connection_secrets(str(e))}", type=ErrorType.UNKNOWN, ), row_count=-1, @@ -1232,21 +1238,7 @@ class DatabaseQueryGenerator: return proc_name, param_values, out_indices def _create_db2_connection(self) -> Any: - conn_str = self.connection_string.split("://")[1] - user_pass, host_db = conn_str.split("@") - user, password = user_pass.split(":") - host_port, database = host_db.split("/") - host, port = host_port.split(":") - - db2_conn_str = ( - f"DATABASE={database};" - f"HOSTNAME={host};" - f"PORT={port};" - f"PROTOCOL=TCPIP;" - f"UID={user};" - f"PWD={password};" - ) - + db2_conn_str = build_db2_native_connection_string(self.connection_string) return _load_ibm_db().connect(db2_conn_str, "", "") def _execute_procedure( diff --git a/private_gpt/components/tools/builders/database_query_builder.py b/private_gpt/components/tools/builders/database_query_builder.py index 43fcf5ba..828ecc36 100644 --- a/private_gpt/components/tools/builders/database_query_builder.py +++ b/private_gpt/components/tools/builders/database_query_builder.py @@ -12,6 +12,9 @@ from llama_index.core.base.llms.types import ( from private_gpt.chat.input_models import BlobVisibilityMode from private_gpt.components.chat.models.chat_config_models import ToolSpec +from private_gpt.components.database.connection_factory import ( + mask_connection_secrets, +) from private_gpt.components.llm.llm_component import LLMComponent from private_gpt.components.tools.binary_block_decorators import ( auto_resolve_media_blocks, @@ -303,7 +306,7 @@ class DatabaseQueryToolBuilder: result_as_block_list: list[list[ResultContentBlockType]] = [] for sql_artifact, db_query_result in results: prefix = ( - f"Database: {sql_artifact.connection_string}\n" + f"Database: {mask_connection_secrets(sql_artifact.connection_string)}\n" if len(results) > 1 else "" )