Agent/db_uri.py

"""
Database-agnostic URI normalization for SQLAlchemy.

QueryLens is architected to interface with PostgreSQL, MySQL, SQLite,
MS SQL Server and Oracle. Users habitually type bare scheme URLs such as
`mysql://user:pass@host/db`, but SQLAlchemy resolves bare schemes to
dialects whose drivers may not be installed (`mysqldb`, `pyodbc`,
`cx_Oracle`). This module maps every commonly typed scheme onto the
drivers the project actually ships:

  postgres:// / postgresql:// .... postgresql+psycopg2  (sync)  | postgresql+asyncpg (async)
  mysql:// / mariadb:// .......... mysql+pymysql        (sync)  | mysql+aiomysql    (async)
  mssql:// / sqlserver:// ........ mssql+pymssql        (sync)  | no async driver exists
  oracle:// ...................... oracle+oracledb thin (sync)  | no stable async dialect
  sqlite:///path ................. sqlite               (sync)  | sqlite+aiosqlite  (async)

URLs that already pin a dialect driver (`scheme+driver://`) are respected:
the pinned driver is kept for the layer it belongs to, and the other layer
gets its family's default driver. Fully unknown schemes are passed through
end to end so users can bring any SQLAlchemy-compatible database.

No heavy imports here: drivers are lazy-loaded by SQLAlchemy itself, so a
machine without (say) pymssql still boots and only fails when someone
actually connects to an MSSQL URL.
"""
import re
from typing import Dict, Optional, Tuple

_SCHEME_RE = re.compile(r"^([a-zA-Z][a-zA-Z0-9+._-]*)://")

_FAMILY_BY_SCHEME = {
    "postgres": "postgresql",
    "postgresql": "postgresql",
    "pg": "postgresql",
    "mysql": "mysql",
    "mariadb": "mysql",
    "mssql": "mssql",
    "sqlserver": "mssql",
    "oracle": "oracle",
    "sqlite": "sqlite",
}

_SYNC_DRIVERS = {
    "postgresql": "psycopg2",
    "mysql": "pymysql",
    "mssql": "pymssql",
    "oracle": "oracledb",
}
_ASYNC_DRIVERS = {
    "postgresql": "asyncpg",
    "mysql": "aiomysql",
    "sqlite": "aiosqlite",
}
_PREFIXES_OF_ASYNC_DRIVERS = {"asyncpg", "aiomysql", "aiosqlite", "asyncpg+gssapi"}


def _parse(url: str) -> Tuple[Optional[str], Optional[str], str]:
    """Split a URL into (sqlalchemy_family, pinned_driver|None, '://rest')."""
    m = _SCHEME_RE.match(url or "")
    if not m:
        return None, None, url or ""
    scheme = m.group(1).lower()
    head, _, driver = scheme.partition("+")
    family = _FAMILY_BY_SCHEME.get(head, head)
    # '://' bytes remain part of the tail so callers can rebuild the URL
    return family, (driver or None) or None, url[m.end(1):]


# libpq connection parameters that the shipped drivers cannot accept.
# psycopg2.connect() validates keyword arguments against its own whitelist and
# raises TypeError "connect() got an unexpected keyword argument '<name>'" for
# anything else; asyncpg behaves the same. Neon's current console template
# appends `channel_binding=require` to freshly copied connection strings,
# which trips both drivers even though the URL is otherwise valid - so strip
# the offending params here instead of asking users to hand-edit pasted
# credentials.
_DRIVER_UNSUPPORTED_QUERY_PARAMS = ("channel_binding",)


def sanitize_driver_query_params(raw_uri: str) -> str:
    """Drop query parameters the shipped drivers reject (e.g. Neon's
    `channel_binding=require` from freshly copied connection strings).

    Scoped to the PostgreSQL family (the only dialect those libpq params
    apply to); TLS via `sslmode` and the password are preserved verbatim.
    Returns the URL unchanged when nothing needs stripping.
    """
    url = (raw_uri or "").strip()
    if not url or "://" not in url:
        return raw_uri or url
    family, _, _ = _parse(url)
    if family != "postgresql":
        return url
    try:
        from sqlalchemy.engine import make_url  # lazy: keep db_uri import-light
        _u = make_url(url)
    except Exception:
        # Regex fallback for URLs make_url cannot parse.
        out = url
        for _name in _DRIVER_UNSUPPORTED_QUERY_PARAMS:
            out = re.sub(r"[?&]" + re.escape(_name) + r"=[^&]*", "", out, flags=re.IGNORECASE)
        if "?" not in out and "&" in out.split("//", 1)[-1]:
            out = out.replace("&", "?", 1)
        return out
    query = dict(_u.query)
    dropped = [k for k in _DRIVER_UNSUPPORTED_QUERY_PARAMS if k in query]
    if not dropped:
        return url
    for k in dropped:
        query.pop(k, None)
    try:
        return _u.set(query=query).render_as_string(hide_password=False)
    except Exception:
        # Older SQLAlchemy without render_as_string: render manually.
        _pw = _u.password
        _auth = (_u.username + (f":{_pw}" if _pw else "") + "@") if _u.username else ""
        _hostport = (_u.host or "") + (f":{_u.port}" if _u.port else "")
        _db = ("/" + _u.database) if _u.database else ""
        _qs = "&".join(f"{k}={v}" for k, v in query.items()) if query else ""
        return f"{_u.drivername}://{_auth}{_hostport}{_db}" + (f"?{_qs}" if _qs else "")


def normalize_db_uris(raw_uri: str) -> Dict[str, Optional[str]]:
    """Normalize a user-supplied connection URL for dual-layer engine use.

    Returns:
        {
          "family":   "postgresql" | "mysql" | "mssql" | "oracle" | "sqlite" | None,
          "sync_uri": URL for SQLAlchemy.create_engine          (always set),
          "async_uri": URL for create_async_engine, or None when the family
                       has no usable async driver (mssql / oracle / exotic),
        }
    """
    url = sanitize_driver_query_params(str(raw_uri or "").strip())
    if not url:
        return {"family": None, "sync_uri": "", "async_uri": None}

    family, pinned, tail = _parse(url)
    if family is None:
        # Not a scheme URL at all (plain path, DSN name, ...) — trust the
        # caller verbatim but keep the async layer off (no async dialect).
        return {"family": None, "sync_uri": url, "async_uri": None}

    if pinned:
        if pinned in _PREFIXES_OF_ASYNC_DRIVERS:
            # User handed us an ASYNC URL; pin the sync layer to the family's
            # sync driver and keep their async URL verbatim.
            sync_uri = f"{family}+{_SYNC_DRIVERS[family]}{tail}" if family in _SYNC_DRIVERS else f"{family}{tail}"
            return {"family": family, "sync_uri": sync_uri, "async_uri": url}
        # Sync (or exotic) driver pinned: keep it verbatim so the user stays in
        # full control; async layer uses the family default when one exists.
        sync_uri = url
        async_uri = f"{family}+{_ASYNC_DRIVERS[family]}{tail}" if family in _ASYNC_DRIVERS else None
        return {"family": family, "sync_uri": sync_uri, "async_uri": async_uri}

    # Bare scheme: map both layers onto the drivers this project ships.
    sync_uri = f"{family}+{_SYNC_DRIVERS[family]}{tail}" if family in _SYNC_DRIVERS else url
    async_uri = f"{family}+{_ASYNC_DRIVERS[family]}{tail}" if family in _ASYNC_DRIVERS else None
    return {"family": family, "sync_uri": sync_uri, "async_uri": async_uri}