feat: implement support for yugabyte (#84)

Author: cofin

.github/workflows/ci.yaml

@@ -40,7 +40,7 @@ jobs:
         uses: astral-sh/setup-uv@v7
 
       - name: Intall dependencies
-        run: uv sync --frozen
+        run: uv sync --frozen --all-extras
 
       - if: matrix.python-version == '3.12'
         name: Run tests with coverage tracking
@@ -81,7 +81,7 @@ jobs:
         uses: astral-sh/setup-uv@v7
 
       - name: Intall dependencies
-        run: uv sync --frozen
+        run: uv sync --frozen --all-extras
 
       - name: Run tests with coverage tracking
         run: uv run pytest -k elasticsearch

docs/supported-databases/index.rst

@@ -14,6 +14,7 @@ This section provides detailed information on the supported databases, including
    spanner
    bigquery
    cockroachdb
+   yugabyte
    gizmosql
    redis
    valkey

docs/supported-databases/yugabyte.rst

@@ -0,0 +1,53 @@
+Yugabyte
+========
+
+Integration with `Yugabyte DB <https://www.yugabyte.com/>`_
+
+Installation
+------------
+
+.. code-block:: bash
+
+   pip install pytest-databases[yugabyte]
+
+Usage Example
+-------------
+
+.. code-block:: python
+
+    import pytest
+    import psycopg
+    from pytest_databases.docker.yugabyte import YugabyteService
+
+    pytest_plugins = ["pytest_databases.docker.yugabyte"]
+
+    @pytest.fixture(scope="session")
+    def yugabyte_uri(yugabyte_service: YugabyteService) -> str:
+        opts = "&".join(f"{k}={v}" for k, v in yugabyte_service.driver_opts.items())
+        return f"postgresql://yugabyte:yugabyte@{yugabyte_service.host}:{yugabyte_service.port}/{yugabyte_service.database}?{opts}"
+
+    def test_yugabyte_service(yugabyte_uri: str) -> None:
+        with psycopg.connect(yugabyte_uri) as conn:
+            db_open = conn.execute("SELECT 1").fetchone()
+            assert db_open is not None and db_open[0] == 1
+
+    def test_yugabyte_connection(yugabyte_connection: psycopg.Connection) -> None:
+        yugabyte_connection.execute("CREATE TABLE if not exists simple_table as SELECT 1")
+        result = yugabyte_connection.execute("select * from simple_table").fetchone()
+        assert result is not None and result[0] == 1
+
+Available Fixtures
+------------------
+
+* ``yugabyte_image``: The Docker image to use for Yugabyte DB.
+* ``yugabyte_service``: A fixture that provides a Yugabyte DB service.
+* ``yugabyte_connection``: A fixture that provides a Yugabyte DB connection.
+* ``yugabyte_driver_opts``: A fixture that provides driver options for Yugabyte DB.
+
+Service API
+-----------
+
+.. automodule:: pytest_databases.docker.yugabyte
+   :members:
+   :undoc-members:
+   :show-inheritance:

pyproject.toml

@@ -74,34 +74,18 @@ postgres = ["psycopg>=3"]
 redis = ["redis"]
 spanner = ["google-cloud-spanner"]
 valkey = ["valkey"]
+yugabyte = ["psycopg2-yugabytedb"]
 
 [dependency-groups]
+build = ["bump-my-version"]
 dev = [
-  # tests
-  "bump-my-version",
-  "pytest-databases[azure-storage,bigquery,cockroachdb,dragonfly,elasticsearch7,elasticsearch8,gizmosql,keydb,mssql,mysql,mariadb,oracle,postgres,redis,spanner,minio,valkey]",
-  "coverage[toml]>=6.2",
-  "pytest",
-  "pytest-cov",
-  "pytest-cdist>=0.2",
-  "pytest-mock",
-  "pytest-click",
-  "pytest-xdist",
-  "slotscheck",
-  "psycopg-binary",                                                                                                                                                             # This fixes tests failing on M series CPUs.
-  # lint
-  "mypy",
-  "ruff",
-  "pyright",
-  "pre-commit",
-  "types-click",
-  "types-six",
-  "types-decorator",
-  "types-pyyaml",
-  "types-docutils",
-  "types-redis",
-  "types-pymysql",
-  # docs
+  { include-group = "extra" },
+  { include-group = "lint" },
+  { include-group = "doc" },
+  { include-group = "test" },
+  { include-group = "build" },
+]
+doc = [
   "auto-pytabs[sphinx]>=0.5.0",
   "shibuya",
   "sphinx>=7.0.0; python_version <= \"3.9\"",
@@ -118,6 +102,31 @@ dev = [
   "sphinx-autodoc-typehints",
   "sphinx-rtd-theme",
 ]
+extra = ["psycopg2-binary", "psycopg[binary,pool]", "psycopg2-yugabytedb-binary"]
+lint = [
+  "mypy",
+  "ruff",
+  "pyright",
+  "pre-commit",
+  "types-click",
+  "types-six",
+  "types-decorator",
+  "types-pyyaml",
+  "types-docutils",
+  "types-redis",
+  "types-pymysql",
+  "slotscheck",
+]
+test = [
+  "coverage[toml]>=6.2",
+  "pytest",
+  "pytest-cov",
+  "pytest-cdist>=0.2",
+  "pytest-mock",
+  "pytest-click",
+  "pytest-xdist",
+  "pytest-sugar",
+]
 
 ##################
 # External Tools #

src/pytest_databases/docker/yugabyte.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+import psycopg
+import pytest
+
+from pytest_databases._service import DockerService, ServiceContainer
+from pytest_databases.helpers import get_xdist_worker_num
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+    from pytest_databases.types import XdistIsolationLevel
+
+
+@pytest.fixture(scope="session")
+def xdist_yugabyte_isolation_level() -> XdistIsolationLevel:
+    return "database"
+
+
+@dataclass
+class YugabyteService(ServiceContainer):
+    database: str
+    driver_opts: dict[str, str]
+
+
+@pytest.fixture(scope="session")
+def yugabyte_driver_opts() -> dict[str, str]:
+    return {"sslmode": "disable"}
+
+
+@pytest.fixture(scope="session")
+def yugabyte_image() -> str:
+    return "software.yugabyte.com/yugabytedb/yugabyte:latest"
+
+
+@pytest.fixture(scope="session")
+def yugabyte_service(
+    docker_service: DockerService,
+    xdist_yugabyte_isolation_level: XdistIsolationLevel,
+    yugabyte_driver_opts: dict[str, str],
+    yugabyte_image: str,
+) -> Generator[YugabyteService, None, None]:
+    def yugabyte_responsive(_service: ServiceContainer) -> bool:
+        opts = "&".join(f"{k}={v}" for k, v in yugabyte_driver_opts.items()) if yugabyte_driver_opts else ""
+        try:
+            conn = psycopg.connect(f"postgresql://yugabyte:yugabyte@{_service.host}:{_service.port}/yugabyte?{opts}")
+        except Exception:  # noqa: BLE001
+            return False
+
+        try:
+            db_open = conn.execute("SELECT 1").fetchone()
+            return bool(db_open is not None and db_open[0] == 1)
+        finally:
+            conn.close()
+
+    container_name = "yugabyte"
+    db_name = "pytest_databases"
+    worker_num = get_xdist_worker_num()
+    if worker_num is not None:
+        suffix = f"_{worker_num}"
+        if xdist_yugabyte_isolation_level == "server":
+            container_name += suffix
+        else:
+            db_name += suffix
+
+    with docker_service.run(
+        image=yugabyte_image,
+        container_port=5433,  # YugabyteDB YSQL port (not CockroachDB's 26257)
+        check=yugabyte_responsive,
+        name=container_name,
+        command="bin/yugabyted start --background=false",
+        exec_after_start=f"sh -c 'bin/ysqlsh -h $(hostname) -U yugabyte -d yugabyte -c \"CREATE DATABASE {db_name};\"'",
+        transient=xdist_yugabyte_isolation_level == "server",
+        timeout=60,  # YugabyteDB needs longer startup time
+    ) as service:
+        yield YugabyteService(
+            host=service.host,
+            port=service.port,
+            database=db_name,
+            driver_opts=yugabyte_driver_opts,
+        )
+
+
+@pytest.fixture(scope="session")
+def yugabyte_connection(
+    yugabyte_service: YugabyteService,
+    yugabyte_driver_opts: dict[str, str],
+) -> Generator[psycopg.Connection, None, None]:
+    opts = "&".join(f"{k}={v}" for k, v in yugabyte_driver_opts.items()) if yugabyte_driver_opts else ""
+    with psycopg.connect(
+        f"postgresql://yugabyte:yugabyte@{yugabyte_service.host}:{yugabyte_service.port}/{yugabyte_service.database}?{opts}"
+    ) as conn:
+        yield conn

tests/test_elasticsearch.py

@@ -35,7 +35,7 @@ def test(elasticsearch_7_service) -> None:
 
 def test_elasticsearch_8(pytester: pytest.Pytester) -> None:
     pytester.makepyfile("""
-    from elasticsearch7 import Elasticsearch
+    from elasticsearch8 import Elasticsearch
 
     pytest_plugins = ["pytest_databases.docker.elastic_search"]
 

tests/test_yugabyte.py

@@ -0,0 +1,106 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import pytest
+
+
+def test_service_fixture(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile("""
+    import pytest
+    import psycopg
+    from pytest_databases.docker.postgres import _make_connection_string  # noqa: PLC2701
+
+    pytest_plugins = ["pytest_databases.docker.yugabyte"]
+
+    def test(yugabyte_service) -> None:
+        opts = "&".join(f"{k}={v}" for k, v in yugabyte_service.driver_opts.items())
+        with psycopg.connect(
+            f"postgresql://yugabyte:yugabyte@{yugabyte_service.host}:{yugabyte_service.port}/{yugabyte_service.database}?{opts}"
+        ) as conn:
+            db_open = conn.execute("SELECT 1").fetchone()
+            assert db_open is not None and db_open[0] == 1
+    """)
+
+    result = pytester.runpytest_subprocess("-p", "pytest_databases")
+    result.assert_outcomes(passed=1)
+
+
+def test_startup_connection_fixture(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile("""
+    import pytest
+    import psycopg
+    from pytest_databases.docker.postgres import _make_connection_string  # noqa: PLC2701
+
+
+    pytest_plugins = ["pytest_databases.docker.yugabyte"]
+
+    def test(yugabyte_connection) -> None:
+        yugabyte_connection.execute("CREATE TABLE if not exists simple_table as SELECT 1")
+        result = yugabyte_connection.execute("select * from simple_table").fetchone()
+        assert result is not None and result[0] == 1
+    """)
+
+    result = pytester.runpytest_subprocess("-p", "pytest_databases")
+    result.assert_outcomes(passed=1)
+
+
+def test_xdist_isolate_database(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile("""
+    import pytest
+    import psycopg
+    from pytest_databases.docker.postgres import _make_connection_string
+
+    pytest_plugins = ["pytest_databases.docker.yugabyte"]
+
+    def test_one(yugabyte_service) -> None:
+        opts = "&".join(f"{k}={v}" for k, v in yugabyte_service.driver_opts.items())
+        with psycopg.connect(
+            f"postgresql://yugabyte:yugabyte@{yugabyte_service.host}:{yugabyte_service.port}/{yugabyte_service.database}?{opts}"
+        ) as conn:
+            conn.execute("CREATE TABLE foo AS SELECT 1")
+
+    def test_two(yugabyte_service) -> None:
+        opts = "&".join(f"{k}={v}" for k, v in yugabyte_service.driver_opts.items())
+        with psycopg.connect(
+            f"postgresql://yugabyte:yugabyte@{yugabyte_service.host}:{yugabyte_service.port}/{yugabyte_service.database}?{opts}"
+        ) as conn:
+            conn.execute("CREATE TABLE foo AS SELECT 1")
+    """)
+
+    result = pytester.runpytest_subprocess("-n", "2")
+    result.assert_outcomes(passed=2)
+
+
+def test_xdist_isolate_server(pytester: pytest.Pytester) -> None:
+    pytester.makepyfile("""
+    import pytest
+    import psycopg
+    from pytest_databases.docker.postgres import _make_connection_string
+
+    pytest_plugins = ["pytest_databases.docker.yugabyte"]
+
+    @pytest.fixture(scope="session")
+    def xdist_yugabyte_isolation_level():
+        return "server"
+
+    def test_one(yugabyte_service) -> None:
+        opts = "&".join(f"{k}={v}" for k, v in yugabyte_service.driver_opts.items())
+        with psycopg.connect(
+            f"postgresql://yugabyte:yugabyte@{yugabyte_service.host}:{yugabyte_service.port}/{yugabyte_service.database}?{opts}",
+            autocommit=True,
+        ) as conn:
+            conn.execute("CREATE DATABASE foo")
+
+    def test_two(yugabyte_service) -> None:
+        opts = "&".join(f"{k}={v}" for k, v in yugabyte_service.driver_opts.items())
+        with psycopg.connect(
+            f"postgresql://yugabyte:yugabyte@{yugabyte_service.host}:{yugabyte_service.port}/{yugabyte_service.database}?{opts}",
+            autocommit=True,
+        ) as conn:
+            conn.execute("CREATE DATABASE foo")
+    """)
+
+    result = pytester.runpytest_subprocess("-n", "2")
+    result.assert_outcomes(passed=2)