from __future__ import annotations
import dataclasses
import os
from contextlib import contextmanager
from typing import TYPE_CHECKING
import psycopg
import pytest
from pytest_databases.helpers import get_xdist_worker_num
from pytest_databases.types import ServiceContainer, XdistIsolationLevel
if TYPE_CHECKING:
from collections.abc import Generator
from pytest_databases._service import DockerService
def _make_connection_string(host: str, port: int, user: str, password: str, database: str) -> str:
return f"dbname={database} user={user} host={host} port={port} password={password}"
[docs]
@pytest.fixture(scope="session")
def xdist_postgres_isolation_level() -> XdistIsolationLevel:
return "database"
[docs]
@dataclasses.dataclass
class PostgresService(ServiceContainer):
database: str
password: str
user: str
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_host() -> str:
return os.environ.get("POSTGRES_HOST", "127.0.0.1")
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_password() -> str:
return "super-secret"
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_user() -> str:
return "postgres"
@contextmanager
def _provide_postgres_service(
docker_service: DockerService,
image: str,
name: str,
host: str,
user: str,
password: str,
xdist_postgres_isolate: XdistIsolationLevel,
) -> Generator[PostgresService, None, None]:
def check(_service: ServiceContainer) -> bool:
try:
with psycopg.connect(
_make_connection_string(
host=_service.host,
port=_service.port,
user=user,
password=password,
database="postgres",
)
) as conn:
db_open = conn.execute("SELECT 1").fetchone()
return bool(db_open is not None and db_open[0] == 1)
except Exception: # noqa: BLE001
return False
worker_num = get_xdist_worker_num()
db_name = "pytest_databases"
if worker_num is not None:
suffix = f"_{worker_num}"
if xdist_postgres_isolate == "server":
name += suffix
else:
db_name += suffix
with docker_service.run(
image=image,
check=check,
container_host=host,
container_port=5432,
name=name,
env={
"POSTGRES_PASSWORD": password,
},
exec_after_start=f"psql -U postgres -d postgres -c 'CREATE DATABASE {db_name};'",
transient=xdist_postgres_isolate == "server",
) as service:
yield PostgresService(
database=db_name,
host=service.host,
port=service.port,
user=user,
password=password,
)
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_11_service(
docker_service: DockerService,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image="postgres:11",
name="postgres-11",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_12_service(
docker_service: DockerService,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image="postgres:12",
name="postgres-12",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_13_service(
docker_service: DockerService,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image="postgres:13",
name="postgres-13",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_14_service(
docker_service: DockerService,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image="postgres:14",
name="postgres-14",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_15_service(
docker_service: DockerService,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image="postgres:15",
name="postgres-15",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_16_service(
docker_service: DockerService,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image="postgres:16",
name="postgres-16",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_17_service(
docker_service: DockerService,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image="postgres:17",
name="postgres-17",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_11_connection(
postgres_11_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_11_service.host,
port=postgres_11_service.port,
user=postgres_11_service.user,
password=postgres_11_service.password,
database=postgres_11_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_12_connection(
postgres_12_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_12_service.host,
port=postgres_12_service.port,
user=postgres_12_service.user,
password=postgres_12_service.password,
database=postgres_12_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_13_connection(
postgres_13_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_13_service.host,
port=postgres_13_service.port,
user=postgres_13_service.user,
password=postgres_13_service.password,
database=postgres_13_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_14_connection(
postgres_14_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_14_service.host,
port=postgres_14_service.port,
user=postgres_14_service.user,
password=postgres_14_service.password,
database=postgres_14_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_15_connection(
postgres_15_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_15_service.host,
port=postgres_15_service.port,
user=postgres_15_service.user,
password=postgres_15_service.password,
database=postgres_15_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_16_connection(
postgres_16_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_16_service.host,
port=postgres_16_service.port,
user=postgres_16_service.user,
password=postgres_16_service.password,
database=postgres_16_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_17_connection(
postgres_17_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_17_service.host,
port=postgres_17_service.port,
user=postgres_17_service.user,
password=postgres_17_service.password,
database=postgres_17_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_image() -> str:
return "postgres:17"
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_service(
docker_service: DockerService,
postgres_image: str,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image=postgres_image,
name="postgres",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def postgres_connection(
postgres_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=postgres_service.host,
port=postgres_service.port,
user=postgres_service.user,
password=postgres_service.password,
database=postgres_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def pgvector_image() -> str:
return "pgvector/pgvector:pg15"
[docs]
@pytest.fixture(autouse=False, scope="session")
def pgvector_service(
docker_service: DockerService,
pgvector_image: str,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image=pgvector_image,
name="pgvector",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def pgvector_connection(
pgvector_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=pgvector_service.host,
port=pgvector_service.port,
user=pgvector_service.user,
password=pgvector_service.password,
database=pgvector_service.database,
),
) as conn:
yield conn
[docs]
@pytest.fixture(autouse=False, scope="session")
def alloydb_omni_image() -> str:
return "google/alloydbomni:16"
[docs]
@pytest.fixture(autouse=False, scope="session")
def alloydb_omni_service(
docker_service: DockerService,
pgvector_image: str,
xdist_postgres_isolation_level: XdistIsolationLevel,
postgres_host: str,
postgres_user: str,
postgres_password: str,
) -> Generator[PostgresService, None, None]:
with _provide_postgres_service(
docker_service,
image=pgvector_image,
name="alloydb-omni",
xdist_postgres_isolate=xdist_postgres_isolation_level,
host=postgres_host,
user=postgres_user,
password=postgres_password,
) as service:
yield service
[docs]
@pytest.fixture(autouse=False, scope="session")
def alloydb_omni_connection(
pgvector_service: PostgresService,
) -> Generator[psycopg.Connection, None, None]:
with psycopg.connect(
_make_connection_string(
host=pgvector_service.host,
port=pgvector_service.port,
user=pgvector_service.user,
password=pgvector_service.password,
database=pgvector_service.database,
),
) as conn:
yield conn
