feat: add support for cockroachdb (#48)

Makefile

@@ -116,9 +116,17 @@ test-spanner:
 test-mssql:
 	$(ENV_PREFIX)pytest tests -m='integration and mssql'
 
+.PHONY: test-cockroachdb-sync
+test-cockroachdb-sync:
+	$(ENV_PREFIX)pytest tests -m='integration and cockroachdb_sync'
+
+.PHONY: test-cockroachdb-async
+test-cockroachdb-async:
+	$(ENV_PREFIX)pytest tests -m='integration and cockroachdb_async'
+
 .PHONY: test-all-databases
 test-all-databases:
-	$(ENV_PREFIX)pytest tests -m='integration and integration'
+	$(ENV_PREFIX)pytest tests -m='integration'
 
 .PHONY: check-all
 check-all: lint test coverage 						## Run all linting, tests, and coverage checks

README.md

@@ -41,6 +41,7 @@ offering features such as:
   - Google Spanner via [spanner-sqlalchemy](https://github.com/googleapis/python-spanner-sqlalchemy/)
   - DuckDB via [duckdb_engine](https://github.com/Mause/duckdb_engine>)
   - Microsoft SQL Server via [pyodbc](https://github.com/mkleehammer/pyodbc>)
+  - CockroachDB via [sqlalchemy-cockroachdb (async or sync)](https://github.com/cockroachdb/sqlalchemy-cockroachdb)
 
 ## Usage
 
@@ -80,6 +81,7 @@ session_factory = sessionmaker(engine, expire_on_commit=False)
 
 with session_factory() as db_session:
   repo = UserRepository(session=db_session)
+  # 1) Create multiple users with `add_many`
   bulk_users = [
     {email: 'cody@advanced-alchemy.dev','name': 'Cody'},
     {email: 'janek@advanced-alchemy.dev','name': 'Janek'},

advanced_alchemy/types.py

@@ -40,7 +40,7 @@ def __init__(self, *args: Any, binary: bool = True, **kwargs: Any) -> None:
         self.binary = binary
 
     def load_dialect_impl(self, dialect: Dialect) -> Any:
-        if dialect.name in {"postgresql", "duckdb"}:
+        if dialect.name in {"postgresql", "duckdb", "cockroachdb"}:
             return dialect.type_descriptor(PG_UUID())
         if dialect.name == "oracle":
             return dialect.type_descriptor(ORA_RAW(16))
@@ -51,7 +51,7 @@ def load_dialect_impl(self, dialect: Dialect) -> Any:
     def process_bind_param(self, value: bytes | str | uuid.UUID | None, dialect: Dialect) -> bytes | str | None:
         if value is None:
             return value
-        if dialect.name in {"postgresql", "duckdb"}:
+        if dialect.name in {"postgresql", "duckdb", "cockroachdb"}:
             return str(value)
         value = self.to_uuid(value)
         if value is None:
@@ -170,7 +170,7 @@ def process_result_value(self, value: datetime.datetime | None, dialect: Dialect
 BigIntIdentity = BigInteger().with_variant(Integer, "sqlite")
 """A ``BigInteger`` variant that reverts to an ``Integer`` for unsupported variants."""
 
-JsonB = _JSON().with_variant(PG_JSONB, "postgresql").with_variant(ORA_JSONB, "oracle")
+JsonB = _JSON().with_variant(PG_JSONB, "postgresql", "cockroachdb").with_variant(ORA_JSONB, "oracle")
 """A JSON type that uses  native ``JSONB`` where possible and ``Binary`` or ``Blob`` as
 an alternative.
 """

pyproject.toml

@@ -65,8 +65,10 @@ dev = [
   "duckdb-engine>=0.9.2",
   "aiosqlite>=0.19.0",
   "psycopg[pool,binary]>=3.1.10",
+  "psycopg2-binary",
   "pytest-xdist>=3.3.1",
   "sqlalchemy-spanner>=1.6.2",
+  "sqlalchemy-cockroachdb>=2.0.0",
   "litestar>=2.0.1",
   "fastapi>=0.103.1",
   "msgspec>=0.18.2",
@@ -114,7 +116,7 @@ test = [
 ]
 
 [tool.pytest.ini_options]
-addopts = "-m='not integration' --dist=loadgroup"
+addopts = "-m='not integration' --dist=loadgroup -n auto"
 asyncio_mode = "auto"
 
 markers = [
@@ -129,6 +131,8 @@ markers = [
   "spanner: SQLAlchemy Google Cloud Spanner (sqlalchemy-spanner) Tests",
   "duckdb: SQLAlchemy DuckDB (duckdb-engine) Tests",
   "mssql: SQLAlchemy Microsoft SQL Server (pyodbc) Tests",
+  "cockroachdb_sync: SQLAlchemy CockroachDB (psycopg2) Tests",
+  "cockroachdb_async: SQLAlchemy CockroachDB (asyncpg) Tests",
 ]
 
 filterwarnings = [

tests/docker-compose.yml

@@ -68,3 +68,18 @@ services:
       MSSQL_PID: Developer
       ACCEPT_EULA: Accepted
       MSSQL_TCP_PORT: 1433
+  cockroachdb:
+    image: cockroachdb/cockroach:latest-v23.1
+    command: start-single-node --insecure
+    restart: "no"
+    expose:
+      - "8080"
+      - "26257"
+    ports:
+      - "26257:26257"
+      - "8880:8080"
+    healthcheck:
+      test: ["CMD", "curl", "-f", "http://localhost:8080/health?ready=1"]
+      interval: 3s
+      timeout: 3s
+      retries: 5

tests/docker_service_fixtures.py

@@ -13,6 +13,7 @@
 import asyncmy
 import asyncpg
 import oracledb
+import psycopg
 import pyodbc
 import pytest
 from google.auth.credentials import AnonymousCredentials
@@ -282,3 +283,19 @@ async def mssql_responsive(host: str) -> bool:
 @pytest.fixture()
 async def mssql_service(docker_services: DockerServiceRegistry) -> None:
     await docker_services.start("mssql", timeout=60, pause=1, check=mssql_responsive)
+
+
+async def cockroachdb_responsive(host: str) -> bool:
+    try:
+        with psycopg.connect("postgresql://root@127.0.0.1:26257/defaultdb?sslmode=disable") as conn:
+            with conn.cursor() as cursor:
+                cursor.execute("select 1 as is_available")
+                resp = cursor.fetchone()
+                return resp[0] == 1  # type: ignore
+    except Exception:
+        return False
+
+
+@pytest.fixture()
+async def cockroachdb_service(docker_services: DockerServiceRegistry) -> None:
+    await docker_services.start("cockroachdb", timeout=60, pause=1, check=cockroachdb_responsive)

tests/integration/conftest.py

@@ -174,6 +174,15 @@ def spanner_engine(docker_ip: str, spanner_service: None, monkeypatch: MonkeyPat
     )
 
 
+@pytest.fixture()
+def cockroachdb_engine(docker_ip: str, cockroachdb_service: None) -> Engine:
+    """CockroachDB instance for end-to-end testing."""
+    return create_engine(
+        url="cockroachdb://root@localhost:26257/defaultdb?sslmode=disable",
+        poolclass=NullPool,
+    )
+
+
 @pytest.fixture(
     name="engine",
     params=[
@@ -209,6 +218,14 @@ def spanner_engine(docker_ip: str, spanner_service: None, monkeypatch: MonkeyPat
                 pytest.mark.xdist_group("postgres"),
             ],
         ),
+        pytest.param(
+            "cockroachdb_engine",
+            marks=[
+                pytest.mark.cockroachdb_sync,
+                pytest.mark.integration,
+                pytest.mark.xdist_group("cockroachdb"),
+            ],
+        ),
         pytest.param(
             "mssql_engine",
             marks=[
@@ -310,6 +327,15 @@ async def psycopg_async_engine(docker_ip: str, postgres_service: None) -> AsyncE
     )
 
 
+@pytest.fixture()
+async def cockroachdb_async_engine(docker_ip: str, cockroachdb_service: None) -> AsyncEngine:
+    """Cockroach DB async engine instance for end-to-end testing."""
+    return create_async_engine(
+        url="cockroachdb+asyncpg://root@localhost:26257/defaultdb",
+        poolclass=NullPool,
+    )
+
+
 @pytest.fixture(
     name="async_engine",
     params=[
@@ -344,6 +370,14 @@ async def psycopg_async_engine(docker_ip: str, postgres_service: None) -> AsyncE
                 pytest.mark.xdist_group("postgres"),
             ],
         ),
+        pytest.param(
+            "cockroachdb_async_engine",
+            marks=[
+                pytest.mark.cockroachdb_async,
+                pytest.mark.integration,
+                pytest.mark.xdist_group("cockroachdb"),
+            ],
+        ),
     ],
 )
 def async_engine(request: FixtureRequest) -> AsyncEngine: