fischer-agentkit/src/agentkit/bitable/ingestion/database.py

"""Database ingestion — reflect external DB tables into bitable-ready data.

Uses SQLAlchemy reflection to read table structure and rows. The caller
(BitableTool) then creates a bitable table + fields and upserts the rows
via the bitable REST API.

Type mapping (KTD: DB → bitable):
    INTEGER / BIGINT / SMALLINT / NUMERIC / FLOAT / DECIMAL → number
    VARCHAR / TEXT / CHAR / UUID                            → text
    TIMESTAMP / DATETIME / DATE                             → date
    BOOLEAN                                                  → text (v1: no bool type)
    JSON / JSONB                                             → text
    fallback                                                 → text
"""

from __future__ import annotations

import logging

from sqlalchemy import (
    BigInteger,
    Boolean,
    Date,
    DateTime,
    Float,
    Integer,
    Numeric,
    SmallInteger,
    String,
    Text,
    create_engine,
    inspect,
    select,
)
from sqlalchemy.engine import Engine

logger = logging.getLogger(__name__)

# ponytail: Static mapping covers all common SQL types. Unknown types fall
# back to text — safe but lossy. Upgrade path: add entries as needed.
DB_TYPE_MAP: dict[type, str] = {
    Integer: "number",
    BigInteger: "number",
    SmallInteger: "number",
    Numeric: "number",
    Float: "number",
    String: "text",
    Text: "text",
    DateTime: "date",
    Date: "date",
    Boolean: "text",
}

# Batch size for reading rows from the source DB
READ_BATCH = 1000


def infer_field_type(sqla_type: object) -> str:
    """Map a SQLAlchemy column type instance or class to a bitable field type.

    Handles both type instances (``Integer()``) and type classes (``Integer``).
    Falls back to ``"text"`` for unknown types.
    """
    for sqla_cls, bitable_type in DB_TYPE_MAP.items():
        if isinstance(sqla_type, sqla_cls):
            return bitable_type
    # If sqla_type is a class (not instance), check subclass relationship
    if isinstance(sqla_type, type):
        for sqla_cls, bitable_type in DB_TYPE_MAP.items():
            if issubclass(sqla_type, sqla_cls):
                return bitable_type
    return "text"


def import_table(
    connection_string: str,
    table_name: str,
    *,
    max_rows: int = 50_000,
) -> dict[str, object]:
    """Reflect a single table from an external DB.

    Returns ``{"table_name": str, "fields": [...], "records": [...],
    "primary_key": str | None, "row_count": int}``.

    Raises ``ConnectionError`` if the DB is unreachable.
    """
    try:
        engine = create_engine(connection_string)
    except Exception as e:
        raise ConnectionError(f"Failed to create engine for connection string: {e}") from e

    try:
        return _reflect_and_read(engine, table_name, max_rows)
    finally:
        engine.dispose()


def _reflect_and_read(engine: Engine, table_name: str, max_rows: int) -> dict[str, object]:
    """Reflect one table and read its rows."""
    insp = inspect(engine)

    # Validate table exists
    if table_name not in insp.get_table_names():
        raise ValueError(f"Table {table_name!r} not found in database")

    from sqlalchemy import Table, MetaData

    metadata = MetaData()
    table = Table(table_name, metadata, autoload_with=engine)

    # Build field definitions
    fields: list[dict[str, object]] = []
    pk_columns = list(table.primary_key.columns)
    pk_name = pk_columns[0].name if pk_columns else None

    for col in table.columns:
        field_type = infer_field_type(col.type)
        fields.append(
            {
                "name": col.name,
                "field_type": field_type,
                "is_primary_key": col.name == pk_name,
            }
        )

    # If no PK, auto-generate one
    if pk_name is None:
        fields.insert(0, {"name": "id", "field_type": "text", "is_primary_key": True})
        pk_name = "id"

    # Read rows
    records: list[dict[str, object]] = []
    with engine.connect() as conn:
        result = conn.execute(select(table))
        for i, row in enumerate(result):
            if i >= max_rows:
                logger.warning("Table %r truncated at %d rows during import", table_name, max_rows)
                break
            rec: dict[str, object] = {}
            for col in table.columns:
                val = getattr(row, col.name, None)
                if val is not None:
                    val = _serialize(val)
                rec[col.name] = val
            records.append(rec)

    return {
        "table_name": table_name,
        "fields": fields,
        "records": records,
        "primary_key": pk_name,
        "row_count": len(records),
    }


def _serialize(val: object) -> object:
    """Serialize a DB value to JSON-safe form."""
    from datetime import date, datetime
    from decimal import Decimal

    if isinstance(val, datetime):
        return val.isoformat()
    if isinstance(val, date):
        return val.isoformat()
    if isinstance(val, Decimal):
        return float(val)
    if isinstance(val, bytes):
        return val.decode("utf-8", errors="replace")
    return val