dblift 3.3.2__tar.gz → 3.3.3__tar.gz

{dblift-3.3.2 → dblift-3.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dblift
-Version: 3.3.2
+Version: 3.3.3
 Summary: A database migration tool that supports multiple databases with native Python drivers
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers

{dblift-3.3.2 → dblift-3.3.3}/db/plugins/cockroachdb/__init__.py

@@ -1,5 +1,13 @@
 """CockroachDB database provider plugin (PostgreSQL-compatible)."""
 
+# Side-effect import: register the SQLAlchemy dialect so create_engine can
+# resolve cockroachdb+psycopg even if plugin.py was not loaded yet.
+from db.plugins.cockroachdb.sqlalchemy_dialect import (  # noqa: F401
+    register_cockroach_dialect,
+)
+
+register_cockroach_dialect()
+
 __plugin_name__ = "cockroachdb"
 __plugin_version__ = "1.0.0"
 __plugin_description__ = "CockroachDB database provider"

{dblift-3.3.2 → dblift-3.3.3}/db/plugins/cockroachdb/plugin.py

@@ -11,9 +11,13 @@ from __future__ import annotations
 
 from db.plugins.cockroachdb.provider import CockroachdbProvider
 from db.plugins.cockroachdb.quirks import CockroachdbQuirks
-from db.plugins.postgresql.sqlalchemy_url import build_sqlalchemy_url
+from db.plugins.cockroachdb.sqlalchemy_dialect import register_cockroach_dialect
+from db.plugins.cockroachdb.sqlalchemy_url import build_sqlalchemy_url
 from db.provider_registry import PluginInfo
 
+# Register before any create_engine call can race plugin discovery.
+register_cockroach_dialect()
+
 PLUGIN: PluginInfo = PluginInfo(
     name="cockroachdb",
     version="1.0.0",

dblift-3.3.3/db/plugins/cockroachdb/sqlalchemy_dialect.py

@@ -0,0 +1,124 @@
+"""SQLAlchemy dialect for CockroachDB (PostgreSQL wire, CRDB version banner).
+
+SQLAlchemy's stock PostgreSQL dialect requires a ``PostgreSQL X.Y`` (or
+``EnterpriseDB``) prefix from ``select version()``. CockroachDB returns banners
+like::
+
+    CockroachDB CCL v24.3.18 (x86_64-pc-linux-gnu, built 2025/08/07 22:44:01, go1.22.8 ...)
+
+which raise ``AssertionError: Could not determine version from string`` on every
+``create_engine`` connect. This dialect reuses the psycopg PostgreSQL dialect
+and only overrides server-version parsing so the product version ``vX.Y.Z`` is
+used — never the trailing ``go1.x.y`` fragment.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import Any, Optional, Tuple
+
+from sqlalchemy.dialects.postgresql.psycopg import PGDialect_psycopg
+
+# Prefer an explicit product "vMAJOR.MINOR[.PATCH]" token (CRDB's usual form).
+_CRDB_V_TOKEN = re.compile(r"\bv(\d+)\.(\d+)(?:\.(\d+))?\b", re.IGNORECASE)
+# Fallback when the banner omits the leading ``v`` (older strings).
+_CRDB_DOTTED = re.compile(
+    r"CockroachDB(?:\s+CCL)?\s+(\d+)\.(\d+)(?:\.(\d+))?",
+    re.IGNORECASE,
+)
+
+_registered = False
+
+
+def parse_cockroach_server_version(banner: str) -> Tuple[int, ...]:
+    """Return ``(major, minor[, patch])`` from a CockroachDB ``version()`` banner.
+
+    Raises:
+        AssertionError: when no product version can be extracted (same shape as
+            SQLAlchemy's PostgreSQL dialect failure).
+    """
+    text = banner or ""
+    match = _CRDB_V_TOKEN.search(text)
+    if match is None:
+        match = _CRDB_DOTTED.search(text)
+    if match is None:
+        raise AssertionError(f"Could not determine version from string '{banner}'")
+    major = int(match.group(1))
+    minor = int(match.group(2))
+    patch = match.group(3)
+    if patch is not None:
+        return (major, minor, int(patch))
+    return (major, minor)
+
+
+class CockroachDBDialect_psycopg(PGDialect_psycopg):
+    """PostgreSQL+psycopg dialect that accepts CockroachDB version banners."""
+
+    name = "cockroachdb"
+    driver = "psycopg"
+    # Inherit PG caching semantics; SQLAlchemy warns if a subclass omits this.
+    supports_statement_cache = True
+
+    def _get_server_version_info(self, connection: Any) -> Tuple[int, ...]:
+        raw = connection.exec_driver_sql("select version()").scalar()
+        return parse_cockroach_server_version(str(raw) if raw is not None else "")
+
+
+def register_cockroach_dialect() -> None:
+    """Register ``cockroachdb`` / ``cockroachdb.psycopg`` with SQLAlchemy once."""
+    global _registered
+    if _registered:
+        return
+    from sqlalchemy.dialects import registry
+
+    registry.register(
+        "cockroachdb",
+        "db.plugins.cockroachdb.sqlalchemy_dialect",
+        "CockroachDBDialect_psycopg",
+    )
+    registry.register(
+        "cockroachdb.psycopg",
+        "db.plugins.cockroachdb.sqlalchemy_dialect",
+        "CockroachDBDialect_psycopg",
+    )
+    _registered = True
+
+
+def ensure_cockroach_drivername(drivername: str) -> str:
+    """Map a PostgreSQL-family drivername onto ``cockroachdb+psycopg``.
+
+    Only the ``psycopg`` DBAPI is registered for CockroachDB. Bare
+    ``postgresql`` / ``postgres`` / ``cockroachdb`` names and ``+psycopg``
+    variants rewrite cleanly; any other driver suffix (e.g. ``asyncpg``)
+    raises so ``create_engine`` does not fail later with an unloadable dialect.
+    """
+    if drivername in (
+        "cockroachdb",
+        "cockroachdb+psycopg",
+        "postgresql",
+        "postgres",
+        "postgresql+psycopg",
+        "postgres+psycopg",
+    ):
+        return "cockroachdb+psycopg"
+    if drivername.startswith(("postgresql+", "postgres+", "cockroachdb+")):
+        suffix = drivername.split("+", 1)[1]
+        if suffix != "psycopg":
+            raise ValueError(
+                "CockroachDB native connections require the psycopg driver "
+                f"(got {drivername!r}); use cockroachdb+psycopg or "
+                "postgresql+psycopg"
+            )
+        return "cockroachdb+psycopg"
+    raise ValueError(
+        "CockroachDB native connections require a PostgreSQL/CockroachDB "
+        f"SQLAlchemy URL (got drivername {drivername!r})"
+    )
+
+
+__all__ = [
+    "CockroachDBDialect_psycopg",
+    "ensure_cockroach_drivername",
+    "parse_cockroach_server_version",
+    "register_cockroach_dialect",
+]

dblift-3.3.3/db/plugins/cockroachdb/sqlalchemy_url.py

@@ -0,0 +1,93 @@
+"""SQLAlchemy URL construction for CockroachDB.
+
+Reuses the PostgreSQL URL shape (host/user/sslmode/search_path) but forces the
+``cockroachdb+psycopg`` drivername so SQLAlchemy loads
+:class:`CockroachDBDialect_psycopg` instead of the stock PostgreSQL dialect
+that cannot parse CockroachDB ``version()`` banners.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict
+
+from sqlalchemy.engine import URL, make_url
+
+from db.plugins.cockroachdb.sqlalchemy_dialect import (
+    ensure_cockroach_drivername,
+    register_cockroach_dialect,
+)
+
+
+def _string_mapping(values: Any) -> Dict[str, str]:
+    if not isinstance(values, dict):
+        return {}
+    return {str(key): str(value) for key, value in values.items()}
+
+
+def _query_mapping(database_config: Any, base_query: Any = None) -> Dict[str, str]:
+    query = _string_mapping(base_query)
+    query.update(_string_mapping(getattr(database_config, "extra_params", None)))
+    query.update(_string_mapping(getattr(database_config, "options", None)))
+
+    connection_timeout = getattr(database_config, "connection_timeout", None)
+    if connection_timeout:
+        query["connect_timeout"] = str(connection_timeout)
+
+    ssl_mode = getattr(database_config, "ssl_mode", None)
+    if ssl_mode:
+        query["sslmode"] = str(ssl_mode)
+
+    schema = getattr(database_config, "schema", None)
+    if schema:
+        schema_option = f"-csearch_path={schema}"
+        query["options"] = (
+            f"{query['options']} {schema_option}" if query.get("options") else schema_option
+        )
+
+    return query
+
+
+def build_sqlalchemy_url(database_config: Any) -> str:
+    """Build a CockroachDB SQLAlchemy URL (psycopg driver, CRDB dialect)."""
+    register_cockroach_dialect()
+
+    raw_url = getattr(database_config, "url", None)
+    if isinstance(raw_url, str) and raw_url:
+        # Same scheme gate as the PostgreSQL builder, plus cockroachdb:// so a
+        # pre-rewritten URL still validates. Other schemes raise immediately.
+        if not raw_url.startswith(
+            (
+                "postgresql://",
+                "postgresql+",
+                "postgres://",
+                "postgres+",
+                "cockroachdb://",
+                "cockroachdb+",
+            )
+        ):
+            raise ValueError(
+                "CockroachDB native connections require a PostgreSQL/CockroachDB SQLAlchemy URL"
+            )
+        url = make_url(raw_url)
+        url = url.set(drivername=ensure_cockroach_drivername(url.drivername))
+        username = getattr(database_config, "username", None) or url.username or None
+        password = getattr(database_config, "password", None) or url.password or None
+        if username != url.username or password != url.password:
+            url = url.set(username=username, password=password)
+        query = _query_mapping(database_config, url.query)
+        if query != dict(url.query):
+            url = url.set(query=query)
+        return url.render_as_string(hide_password=False)
+
+    return URL.create(
+        "cockroachdb+psycopg",
+        username=getattr(database_config, "username", None) or None,
+        password=getattr(database_config, "password", None) or None,
+        host=getattr(database_config, "host", None) or "localhost",
+        port=getattr(database_config, "port", None),
+        database=getattr(database_config, "database", None),
+        query=_query_mapping(database_config),
+    ).render_as_string(hide_password=False)
+
+
+__all__ = ["build_sqlalchemy_url"]

{dblift-3.3.2 → dblift-3.3.3}/dblift.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dblift
-Version: 3.3.2
+Version: 3.3.3
 Summary: A database migration tool that supports multiple databases with native Python drivers
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers

{dblift-3.3.2 → dblift-3.3.3}/dblift.egg-info/SOURCES.txt

@@ -295,6 +295,8 @@ db/plugins/cockroachdb/__init__.py
 db/plugins/cockroachdb/plugin.py
 db/plugins/cockroachdb/provider.py
 db/plugins/cockroachdb/quirks.py
+db/plugins/cockroachdb/sqlalchemy_dialect.py
+db/plugins/cockroachdb/sqlalchemy_url.py
 db/plugins/cosmosdb/__init__.py
 db/plugins/cosmosdb/config.py
 db/plugins/cosmosdb/plugin.py

{dblift-3.3.2 → dblift-3.3.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "dblift"
-version = "3.3.2"
+version = "3.3.3"
 description = "A database migration tool that supports multiple databases with native Python drivers"
 readme = "README.md"
 requires-python = ">=3.11"