[KM-437][DB] add mysql, sqlserver, bigquery, snowflake connections

src/api/v1/db_client.py

@@ -340,12 +340,7 @@ async def ingest_database_client(
     try:
         with db_pipeline_service.engine_scope(
             db_type=client.db_type,
-            host=creds["host"],
-            port=creds["port"],
-            database=creds["database"],
-            username=creds["username"],
-            password=creds["password"],
-            ssl_mode=creds.get("ssl_mode"),
+            credentials=creds,
         ) as engine:
             total = await db_pipeline_service.run(user_id=user_id, engine=engine)
     except NotImplementedError as e:

src/pipeline/db_pipeline/db_pipeline_service.py

@@ -10,7 +10,7 @@ async vector writes stay on the event loop.
 
 import asyncio
 from contextlib import contextmanager
-from typing import Iterator, Optional
+from typing import Any, Iterator, Optional
 
 from langchain_core.documents import Document as LangChainDocument
 from sqlalchemy import URL, create_engine
@@ -27,70 +27,106 @@ logger = get_logger("db_pipeline")
 class DbPipelineService:
     """End-to-end DB ingestion: connect -> introspect -> profile -> embed -> store."""
 
-    def connect(
-        self,
-        db_type: DbType,
-        host: str,
-        port: int,
-        database: str,
-        username: str,
-        password: str,
-        ssl_mode: Optional[str] = None,
-    ) -> Engine:
+    def connect(self, db_type: DbType, credentials: dict[str, Any]) -> Engine:
         """Build a SQLAlchemy engine for the user's database.
 
-        Supabase aliases to the Postgres driver (same URL shape). Other
-        engines raise NotImplementedError until their connector is added.
+        `credentials` is the plaintext dict matching the per-type schema in
+        `src/models/credentials.py`. BigQuery/Snowflake auth models differ
+        from host/port/user/pass, so every shape flows through one dict.
 
-        `ssl_mode` maps to libpq's `sslmode` query param for postgres/supabase
-        (required for managed DBs like Neon/Supabase: "require", "verify-ca",
-        "verify-full"). Ignored for other db_types until those connectors land.
+        Optional driver imports (snowflake-sqlalchemy, json for BigQuery) are
+        done lazily so an env missing one driver doesn't break module import.
         """
-        logger.info(
-            "connecting to user db", db_type=db_type, host=host, port=port, database=database
-        )
+        logger.info("connecting to user db", db_type=db_type)
+
         if db_type in ("postgres", "supabase"):
-            query = {"sslmode": ssl_mode} if ssl_mode else {}
+            query = (
+                {"sslmode": credentials["ssl_mode"]} if credentials.get("ssl_mode") else {}
+            )
             url = URL.create(
                 drivername="postgresql+psycopg2",
-                username=username,
-                password=password,
-                host=host,
-                port=port,
-                database=database,
+                username=credentials["username"],
+                password=credentials["password"],
+                host=credentials["host"],
+                port=credentials["port"],
+                database=credentials["database"],
                 query=query,
             )
             return create_engine(url)
-        elif db_type == "mysql":
-            raise NotImplementedError("MySQL support coming soon")
-        elif db_type == "sqlserver":
-            raise NotImplementedError("SQL Server support coming soon")
-        elif db_type == "bigquery":
-            raise NotImplementedError("BigQuery support coming soon")
-        elif db_type == "snowflake":
-            raise NotImplementedError("Snowflake support coming soon")
-        else:
-            raise ValueError(f"Unsupported db_type: {db_type}")
+
+        if db_type == "mysql":
+            url = URL.create(
+                drivername="mysql+pymysql",
+                username=credentials["username"],
+                password=credentials["password"],
+                host=credentials["host"],
+                port=credentials["port"],
+                database=credentials["database"],
+            )
+            # pymysql: empty-dict ssl arg flips SSL on with defaults.
+            connect_args = {"ssl": {}} if credentials.get("ssl", True) else {}
+            return create_engine(url, connect_args=connect_args)
+
+        if db_type == "sqlserver":
+            # `driver` applies to pyodbc only; we ship pymssql. Accept-and-ignore
+            # keeps the credential schema stable.
+            if credentials.get("driver"):
+                logger.info(
+                    "sqlserver driver hint ignored (using pymssql)",
+                    driver=credentials["driver"],
+                )
+            url = URL.create(
+                drivername="mssql+pymssql",
+                username=credentials["username"],
+                password=credentials["password"],
+                host=credentials["host"],
+                port=credentials["port"],
+                database=credentials["database"],
+            )
+            return create_engine(url)
+
+        if db_type == "bigquery":
+            import json
+
+            sa_info = json.loads(credentials["service_account_json"])
+            # sqlalchemy-bigquery URL shape: bigquery://<project>/<dataset>
+            url = f"bigquery://{credentials['project_id']}/{credentials['dataset_id']}"
+            return create_engine(
+                url,
+                credentials_info=sa_info,
+                location=credentials.get("location", "US"),
+            )
+
+        if db_type == "snowflake":
+            from snowflake.sqlalchemy import URL as SnowflakeURL
+
+            url = SnowflakeURL(
+                account=credentials["account"],
+                user=credentials["username"],
+                password=credentials["password"],
+                database=credentials["database"],
+                schema=(
+                    credentials.get("db_schema")
+                    or credentials.get("schema")
+                    or "PUBLIC"
+                ),
+                warehouse=credentials["warehouse"],
+                role=credentials.get("role") or "",
+            )
+            return create_engine(url)
+
+        raise ValueError(f"Unsupported db_type: {db_type}")
 
     @contextmanager
     def engine_scope(
-        self,
-        db_type: DbType,
-        host: str,
-        port: int,
-        database: str,
-        username: str,
-        password: str,
-        ssl_mode: Optional[str] = None,
+        self, db_type: DbType, credentials: dict[str, Any]
     ) -> Iterator[Engine]:
         """Yield a connected Engine and dispose its pool on exit.
 
         API callers should prefer this over raw `connect(...)` so user DB
         connection pools do not leak between pipeline runs.
         """
-        engine = self.connect(
-            db_type, host, port, database, username, password, ssl_mode
-        )
+        engine = self.connect(db_type, credentials)
         try:
             yield engine
         finally:

src/pipeline/db_pipeline/extractor.py

@@ -18,6 +18,28 @@ logger = get_logger("db_extractor")
 
 TOP_VALUES_THRESHOLD = 0.05  # show top values if distinct_ratio <= 5%
 
+# Dialects where PERCENTILE_CONT(...) WITHIN GROUP is supported as an aggregate.
+# MySQL has no percentile aggregate; BigQuery has PERCENTILE_CONT only as an
+# analytic (window) function — both drop median and keep min/max/mean.
+_MEDIAN_DIALECTS = frozenset({"postgresql", "mssql", "snowflake"})
+
+
+def _supports_median(engine: Engine) -> bool:
+    return engine.dialect.name in _MEDIAN_DIALECTS
+
+
+def _head_query(
+    engine: Engine,
+    select_clause: str,
+    from_clause: str,
+    n: int,
+    order_by: str = "",
+) -> str:
+    """LIMIT/TOP-equivalent head query for the engine's dialect."""
+    if engine.dialect.name == "mssql":
+        return f"SELECT TOP {n} {select_clause} FROM {from_clause} {order_by}".strip()
+    return f"SELECT {select_clause} FROM {from_clause} {order_by} LIMIT {n}".strip()
+
 
 def _qi(engine: Engine, name: str) -> str:
     """Dialect-correct identifier quoting (schema.table also handled if dotted)."""
@@ -95,11 +117,10 @@ def profile_column(
         select_cols.append(f"MIN({qc}) AS min_val")
         select_cols.append(f"MAX({qc}) AS max_val")
         select_cols.append(f"AVG({qc}) AS mean_val")
-        # PERCENTILE_CONT is supported by Postgres and SQL Server; MySQL would need
-        # a dialect-specific fallback when that connector is added.
-        select_cols.append(
-            f"PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY {qc}) AS median_val"
-        )
+        if _supports_median(engine):
+            select_cols.append(
+                f"PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY {qc}) AS median_val"
+            )
     stats = pd.read_sql(f"SELECT {', '.join(select_cols)} FROM {qt}", engine)
 
     null_count = int(stats.iloc[0]["nulls"])
@@ -116,17 +137,21 @@ def profile_column(
         profile["min"] = stats.iloc[0]["min_val"]
         profile["max"] = stats.iloc[0]["max_val"]
         profile["mean"] = stats.iloc[0]["mean_val"]
-        profile["median"] = stats.iloc[0]["median_val"]
+        if _supports_median(engine):
+            profile["median"] = stats.iloc[0]["median_val"]
 
     if 0 < distinct_ratio <= TOP_VALUES_THRESHOLD:
-        top = pd.read_sql(
-            f"SELECT {qc}, COUNT(*) AS cnt FROM {qt} "
-            f"GROUP BY {qc} ORDER BY cnt DESC LIMIT 10",
+        top_sql = _head_query(
             engine,
+            select_clause=f"{qc}, COUNT(*) AS cnt",
+            from_clause=f"{qt} GROUP BY {qc}",
+            n=10,
+            order_by="ORDER BY cnt DESC",
         )
+        top = pd.read_sql(top_sql, engine)
         profile["top_values"] = list(zip(top[col_name].tolist(), top["cnt"].tolist()))
 
-    sample = pd.read_sql(f"SELECT {qc} FROM {qt} LIMIT 5", engine)
+    sample = pd.read_sql(_head_query(engine, qc, qt, 5), engine)
     profile["sample_values"] = sample[col_name].tolist()
 
     return profile
@@ -178,7 +203,9 @@ def build_text(table_name: str, row_count: int, col: dict, profile: dict) -> str
     text += f"Distinct count: {profile['distinct_count']} ({profile['distinct_ratio']:.1%})\n"
     if "min" in profile:
         text += f"Min: {profile['min']}, Max: {profile['max']}\n"
-        text += f"Mean: {profile['mean']}, Median: {profile['median']}\n"
+        text += f"Mean: {profile['mean']}\n"
+        if profile.get("median") is not None:
+            text += f"Median: {profile['median']}\n"
     if "top_values" in profile:
         top_str = ", ".join(f"{v} ({c})" for v, c in profile["top_values"])
         text += f"Top values: {top_str}\n"