Spaces:
Running on Zero
Running on Zero
| """ | |
| 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} | |