github/privateGPT
1
0
Fork 1
mirror of https://github.com/imartinez/privateGPT.git synced 2026-07-16 17:00:12 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add db2 native connection

Browse Source
This commit is contained in:
Alfonso Lozana
2026-07-08 17:20:35 +02:00
parent bcadb418f3
commit 41f130527d
6 changed files with 163 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files

118
private_gpt/components/database/connection_factory.py Normal file
View File

@@ -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)

7
private_gpt/components/database/function_inspector.py
View File

@@ -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 []

12
private_gpt/components/database/inspector_interface.py
View File

@@ -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
)

7
private_gpt/components/database/procedure_inspector.py
View File

@@ -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 []

54
private_gpt/components/tabular/database_query_generator.py
View File

@@ -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(

5
private_gpt/components/tools/builders/database_query_builder.py
View File

@@ -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 ""
)