elefast 0.1.0a1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,59 @@
1
+ Metadata-Version: 2.4
2
+ Name: elefast
3
+ Version: 0.1.0a1
4
+ Summary: Database testing toolkit integrating Postgres, SQLAlchemy, Pytest and Docker
5
+ License-Expression: MIT
6
+ Classifier: Development Status :: 3 - Alpha
7
+ Classifier: License :: OSI Approved :: MIT License
8
+ Classifier: Framework :: AsyncIO
9
+ Classifier: Framework :: Pytest
10
+ Classifier: Programming Language :: Python
11
+ Classifier: Programming Language :: Python :: 3.14
12
+ Classifier: Topic :: Education :: Testing
13
+ Classifier: Topic :: Software Development :: Testing
14
+ Classifier: Topic :: Software Development :: Documentation
15
+ Classifier: Topic :: Database
16
+ Requires-Dist: sqlalchemy>=2.0.45
17
+ Requires-Dist: docker>=7.1.0 ; extra == 'docker'
18
+ Requires-Dist: filelock>=3.20.3 ; extra == 'docker'
19
+ Requires-Python: >=3.14
20
+ Provides-Extra: docker
21
+ Description-Content-Type: text/markdown
22
+
23
+ > [!WARNING]
24
+ > Alpha-level software.
25
+
26
+ # Elefast 🐘⚡
27
+
28
+ Using an actual database for testing is nice, but setting everything up can be a pain.
29
+ Generating the schema, isolating the state, and supporting parallel execution while keeping everything reasonably quick.
30
+ Elefast helps you by providing utilities for all of these tasks
31
+
32
+ To learn more, visit [the documentation](https://niclasvaneyk.github.io/elefast/getting-started) or run through the quickstart guide below.
33
+
34
+ ## Quickstart
35
+
36
+ Install the package from PyPi through your package manager
37
+
38
+ ```shell
39
+ uv add 'elefast[docker]'
40
+ ```
41
+
42
+ add the recommended set of fixtures to your `conftest.py`
43
+
44
+ ```shell
45
+ mkdir tests/ && uv run elefast init >> tests/conftest.py
46
+ ```
47
+
48
+ create a test using your new set of fixtures
49
+
50
+ ```python
51
+ # tests/test_something_with_the_database.py
52
+ from sqlalchemy import Connection, text
53
+
54
+ def test_database_math(db_connection: Connection):
55
+ result = db_connection.execute(text("SELECT 1 + 1")).scalar_one()
56
+ assert result == 2
57
+ ```
58
+
59
+ When you now run `pytest`, we'll automatically start a Docker container (if one is not already running), and pass a database connection into the test
@@ -0,0 +1,37 @@
1
+ > [!WARNING]
2
+ > Alpha-level software.
3
+
4
+ # Elefast 🐘⚡
5
+
6
+ Using an actual database for testing is nice, but setting everything up can be a pain.
7
+ Generating the schema, isolating the state, and supporting parallel execution while keeping everything reasonably quick.
8
+ Elefast helps you by providing utilities for all of these tasks
9
+
10
+ To learn more, visit [the documentation](https://niclasvaneyk.github.io/elefast/getting-started) or run through the quickstart guide below.
11
+
12
+ ## Quickstart
13
+
14
+ Install the package from PyPi through your package manager
15
+
16
+ ```shell
17
+ uv add 'elefast[docker]'
18
+ ```
19
+
20
+ add the recommended set of fixtures to your `conftest.py`
21
+
22
+ ```shell
23
+ mkdir tests/ && uv run elefast init >> tests/conftest.py
24
+ ```
25
+
26
+ create a test using your new set of fixtures
27
+
28
+ ```python
29
+ # tests/test_something_with_the_database.py
30
+ from sqlalchemy import Connection, text
31
+
32
+ def test_database_math(db_connection: Connection):
33
+ result = db_connection.execute(text("SELECT 1 + 1")).scalar_one()
34
+ assert result == 2
35
+ ```
36
+
37
+ When you now run `pytest`, we'll automatically start a Docker container (if one is not already running), and pass a database connection into the test
@@ -0,0 +1,47 @@
1
+ [project]
2
+ name = "elefast"
3
+ version = "0.1.0a1"
4
+ description = "Database testing toolkit integrating Postgres, SQLAlchemy, Pytest and Docker"
5
+ readme = "README.md"
6
+ authors = []
7
+ requires-python = ">=3.14"
8
+ dependencies = [
9
+ "sqlalchemy>=2.0.45",
10
+ ]
11
+ license = "MIT"
12
+ classifiers = [
13
+ "Development Status :: 3 - Alpha",
14
+ "License :: OSI Approved :: MIT License",
15
+ "Framework :: AsyncIO",
16
+ "Framework :: Pytest",
17
+ "Programming Language :: Python",
18
+ "Programming Language :: Python :: 3.14",
19
+ "Topic :: Education :: Testing",
20
+ "Topic :: Software Development :: Testing",
21
+ "Topic :: Software Development :: Documentation",
22
+ "Topic :: Database",
23
+ ]
24
+
25
+ [project.optional-dependencies]
26
+ docker = [
27
+ "docker>=7.1.0",
28
+ "filelock>=3.20.3",
29
+ ]
30
+
31
+ [project.scripts]
32
+ elefast = "elefast.cli:main"
33
+
34
+ [build-system]
35
+ requires = ["uv_build>=0.9.18,<0.10.0"]
36
+ build-backend = "uv_build"
37
+
38
+ [dependency-groups]
39
+ dev = [
40
+ "zensical>=0.0.15",
41
+ ]
42
+
43
+ [tool.uv.workspace]
44
+ members = [
45
+ "examples/fastapi-async",
46
+ "examples/simple-sync",
47
+ ]
@@ -0,0 +1,17 @@
1
+ from elefast.asyncio import (
2
+ AsyncDatabase,
3
+ AsyncDatabaseServer,
4
+ CanBeTurnedIntoAsyncEngine,
5
+ )
6
+ from elefast.sync import DatabaseServer, Database, CanBeTurnedIntoEngine
7
+
8
+ __all__ = [
9
+ "AsyncDatabase",
10
+ "AsyncDatabaseServer",
11
+ "BasicAsyncDatabaseServer",
12
+ "BasicDatabaseServer",
13
+ "Database",
14
+ "DatabaseServer",
15
+ "CanBeTurnedIntoEngine",
16
+ "CanBeTurnedIntoAsyncEngine",
17
+ ]
@@ -0,0 +1,127 @@
1
+ from __future__ import annotations
2
+
3
+ import time
4
+ from asyncio import sleep
5
+ from collections.abc import Callable
6
+ from contextlib import AbstractAsyncContextManager
7
+ from uuid import uuid4
8
+
9
+ from sqlalchemy import URL, MetaData, NullPool, text
10
+ from sqlalchemy.ext.asyncio import (
11
+ AsyncEngine,
12
+ AsyncSession,
13
+ async_sessionmaker,
14
+ create_async_engine,
15
+ )
16
+
17
+ from elefast.errors import DatabaseNotReadyError
18
+
19
+ type CanBeTurnedIntoAsyncEngine = AsyncEngine | URL | str
20
+
21
+
22
+ class AsyncDatabase(AbstractAsyncContextManager):
23
+ def __init__(
24
+ self,
25
+ engine: AsyncEngine,
26
+ server: AsyncDatabaseServer,
27
+ sessionmaker_factory: Callable[
28
+ [AsyncEngine], Callable[[], AsyncSession]
29
+ ] = async_sessionmaker,
30
+ ) -> None:
31
+ self.engine = engine
32
+ self.server = server
33
+ self.sessionmaker = sessionmaker_factory(self.engine)
34
+ assert self.engine.url.database
35
+ self.name = self.engine.url.database
36
+
37
+ async def __aexit__(self, exc_type, exc, tb):
38
+ await self.drop()
39
+
40
+ @property
41
+ def url(self) -> URL:
42
+ return self.engine.url
43
+
44
+ async def drop(self) -> None:
45
+ await self.engine.dispose()
46
+ await self.server.drop_database(self.name)
47
+
48
+ def session(self) -> AsyncSession:
49
+ return self.sessionmaker()
50
+
51
+
52
+ class AsyncDatabaseServer:
53
+ def __init__(
54
+ self, engine: CanBeTurnedIntoAsyncEngine, metadata: MetaData | None = None
55
+ ) -> None:
56
+ self._metadata = metadata
57
+ self._engine = _build_engine(engine)
58
+ self._template_db_name: str | None = None
59
+
60
+ @property
61
+ def url(self) -> URL:
62
+ return self._engine.url
63
+
64
+ async def ensure_is_ready(
65
+ self, timeout: int | float = 30, interval: int | float = 0.5
66
+ ) -> None:
67
+ deadline = time.monotonic() + timeout
68
+ attempts = 0
69
+
70
+ while True:
71
+ try:
72
+ async with self._engine.connect() as conn:
73
+ await conn.execute(text("SELECT 1"))
74
+ return
75
+ except Exception as error:
76
+ attempts += 1
77
+ if time.monotonic() >= deadline:
78
+ raise DatabaseNotReadyError(
79
+ f"Reached the configured timeout of {timeout} seconds after {attempts} attempts connecting to the database."
80
+ ) from error
81
+ await sleep(interval)
82
+
83
+ async def create_database(self) -> AsyncDatabase:
84
+ template_db = self._template_db_name
85
+ if template_db is None:
86
+ engine = await _prepare_async_database(self._engine)
87
+ if self._metadata:
88
+ async with engine.begin() as connection:
89
+ await connection.run_sync(self._metadata.drop_all)
90
+ await connection.run_sync(self._metadata.create_all)
91
+ await engine.dispose()
92
+ template_db = engine.url.database
93
+ assert isinstance(template_db, str)
94
+ self._template_db_name = template_db
95
+
96
+ engine = await _prepare_async_database(self._engine, template=template_db)
97
+ return AsyncDatabase(engine=engine, server=self)
98
+
99
+ async def drop_database(self, name: str) -> None:
100
+ async with self._engine.begin() as connection:
101
+ statement = f'DROP DATABASE "{name}"'
102
+ await connection.execute(text(statement))
103
+
104
+
105
+ async def _prepare_async_database(
106
+ engine: AsyncEngine, encoding: str = "utf8", template: str | None = None
107
+ ) -> AsyncEngine:
108
+ database = f"pytest-elephantastic-{uuid4()}"
109
+ async with engine.begin() as connection:
110
+ statement = (
111
+ f"CREATE DATABASE \"{database}\" ENCODING '{encoding}' TEMPLATE template0"
112
+ if template is None
113
+ else f'CREATE DATABASE "{database}" WITH TEMPLATE "{template}"'
114
+ )
115
+ await connection.execute(text(statement))
116
+ await connection.commit()
117
+ return create_async_engine(engine.url.set(database=database))
118
+
119
+
120
+ def _build_engine(input: CanBeTurnedIntoAsyncEngine) -> AsyncEngine:
121
+ if isinstance(input, AsyncEngine):
122
+ return input
123
+ if isinstance(input, URL | str):
124
+ return create_async_engine(
125
+ input, isolation_level="autocommit", poolclass=NullPool
126
+ )
127
+ raise TypeError()
@@ -0,0 +1,11 @@
1
+ from elefast.cli.parser import build_parser
2
+
3
+
4
+ def main():
5
+ cli = build_parser()
6
+ args = cli.parse_args()
7
+ args.func(args)
8
+
9
+
10
+ if __name__ == "__main__":
11
+ main()
@@ -0,0 +1,20 @@
1
+ from elefast.cli.types import ParentParser
2
+
3
+
4
+ def install_docker_command_if_available(commands: ParentParser) -> None:
5
+ try:
6
+ import elefast.docker # noqa: F401
7
+ except Exception:
8
+ return
9
+
10
+ parser = commands.add_parser(
11
+ "docker",
12
+ help="Utilities for managing elefasts docker container",
13
+ description="A helpful description.", # TODO
14
+ )
15
+ parser.add_argument("--reuse", action="store_true")
16
+ parser.set_defaults(func=docker_command)
17
+
18
+
19
+ def docker_command(args) -> None:
20
+ print(args)
@@ -0,0 +1,134 @@
1
+ from argparse import Namespace
2
+ from sys import stderr
3
+
4
+ from elefast.cli.types import ParentParser
5
+
6
+ # NOTE: psycopg can do both, so we cannot infer anything.
7
+ STRICTLY_ASYNC_DRIVERS = ["asyncpg", "aiopg"]
8
+ STRICTLY_SYNC_DRIVERS = ["psycopg2", "pg8000"]
9
+
10
+
11
+ def install_init_command(commands: ParentParser) -> None:
12
+ init_command_parser = commands.add_parser(
13
+ "init",
14
+ help="Helps you the recommended default fixtures",
15
+ description="Helps you get started with the recommended default fixtures by printing them to the screen. The output of this command is intended to be redirected into a file (e.g. `elefast init > conftest.py && mv conftest.py tests/`) If you omit the flags, you'll be prompted for values.",
16
+ )
17
+ init_command_parser.set_defaults(func=init_command)
18
+ init_command_parser.add_argument(
19
+ "--driver", help="The name of the driver you intend to use."
20
+ )
21
+ init_command_parser.add_argument(
22
+ "--async",
23
+ action="append_const",
24
+ dest="async_preference",
25
+ const=True,
26
+ help="Use async methods and classes where necessary. Inferred from the driver if not passed explicitly.",
27
+ )
28
+ init_command_parser.add_argument(
29
+ "--sync",
30
+ action="append_const",
31
+ dest="async_preference",
32
+ const=False,
33
+ help="Use sync methods and classes. Inferred from the driver if not passed explicitly.",
34
+ )
35
+ init_command_parser.add_argument(
36
+ "--no-interaction",
37
+ action="store_true",
38
+ help="Use defaults instead of asking interactive questions for options",
39
+ )
40
+
41
+
42
+ def init_command(args: Namespace):
43
+ _init_command(
44
+ driver=args.driver,
45
+ allow_interaction=not args.no_interaction,
46
+ async_preference=args.async_preference or [],
47
+ )
48
+
49
+
50
+ def _init_command(
51
+ driver: str | None,
52
+ allow_interaction: bool,
53
+ async_preference: list[bool],
54
+ ):
55
+ driver = _figure_out_driver(driver)
56
+ use_async = _figure_out_if_we_should_use_async(
57
+ driver, async_preference, allow_interaction
58
+ )
59
+
60
+ class_prefix = "Async" if use_async else ""
61
+ maybe_async = "async " if use_async else ""
62
+ maybe_await = "await " if use_async else ""
63
+ template = f'''
64
+ import os
65
+
66
+ import pytest
67
+ from elefast import {class_prefix}Database, {class_prefix}DatabaseServer, docker
68
+
69
+
70
+ @pytest.fixture(scope="session")
71
+ def db_server() -> {class_prefix}DatabaseServer:
72
+ explicit_url = os.getenv("TESTING_DB_URL")
73
+ db_url = explicit_url if explicit_url else docker.postgres("{driver}")
74
+ # If you have a shared Base-class, import it above and use
75
+ # `metadata=YourBaseClass.metadata` below.
76
+ return {class_prefix}DatabaseServer(db_url, metadata=None)
77
+
78
+
79
+ @pytest.fixture
80
+ {maybe_async}def db(db_server: {class_prefix}DatabaseServer):
81
+ {maybe_async}with {maybe_await}db_server.create_database() as database:
82
+ yield database
83
+
84
+
85
+ @pytest.fixture
86
+ {maybe_async}def db_connection(db: {class_prefix}Database):
87
+ {maybe_async}with db.engine.begin() as connection:
88
+ yield connection
89
+
90
+
91
+ @pytest.fixture
92
+ {maybe_async}def db_session(db: {class_prefix}Database):
93
+ {maybe_async}with db.session() as session:
94
+ yield session
95
+ '''
96
+ print(template.strip() + "\n")
97
+
98
+
99
+ def _figure_out_driver(driver: str | None) -> str:
100
+ if driver:
101
+ return driver
102
+
103
+ for name in STRICTLY_SYNC_DRIVERS + STRICTLY_ASYNC_DRIVERS + ["psycopg"]:
104
+ from importlib.util import find_spec
105
+
106
+ if find_spec(name):
107
+ print(
108
+ f"Using '{name}' as the driver, since no explit --driver argument was passed, and {name} is installed.",
109
+ file=stderr,
110
+ )
111
+ return name
112
+
113
+ print(
114
+ "No --driver was passed and no popular one was found in installed packages. Falling back to psycopg2...",
115
+ file=stderr,
116
+ )
117
+ return "psycopg2"
118
+
119
+
120
+ def _figure_out_if_we_should_use_async(
121
+ driver: str, async_preference: list[bool], allow_interaction: bool
122
+ ) -> bool:
123
+ if async_preference == []:
124
+ if driver in STRICTLY_ASYNC_DRIVERS:
125
+ return True
126
+ elif driver in STRICTLY_SYNC_DRIVERS:
127
+ return False
128
+ elif not allow_interaction:
129
+ return False
130
+ else:
131
+ print("Are you intending to use asyncio? [y/N]: ", file=stderr)
132
+ return "y" in input()
133
+
134
+ return any(async_preference)
@@ -0,0 +1,17 @@
1
+ from elefast.cli.docker import install_docker_command_if_available
2
+ from elefast.cli.init import install_init_command
3
+ from argparse import ArgumentParser
4
+
5
+ DESCRIPTION = """
6
+ A helpful description
7
+ """.strip()
8
+
9
+
10
+ def build_parser() -> ArgumentParser:
11
+ cli = ArgumentParser(prog="elefast", description=DESCRIPTION)
12
+ commands = cli.add_subparsers(required=True, title="subcommands")
13
+
14
+ install_init_command(commands)
15
+ install_docker_command_if_available(commands)
16
+
17
+ return cli
@@ -0,0 +1,3 @@
1
+ from argparse import _SubParsersAction, ArgumentParser
2
+
3
+ type ParentParser = _SubParsersAction[ArgumentParser]
@@ -0,0 +1,8 @@
1
+ """
2
+ This module is part of the 'docker' extra and needs to be explicitly installed
3
+ before it is safe to use!
4
+ """
5
+
6
+ from elefast.docker.integration import postgres
7
+
8
+ __all__ = ["postgres"]
@@ -0,0 +1,45 @@
1
+ from dataclasses import dataclass, field
2
+
3
+
4
+ @dataclass(frozen=True, slots=True, kw_only=True)
5
+ class Optimizations:
6
+ """
7
+ Configuration overrides that make Postgres more suitable for fast testing runs.
8
+ """
9
+
10
+ tmpfs_size_mb: int | None = 512
11
+ fsync_off: bool = True
12
+ synchronous_commit_off: bool = True
13
+ full_page_writes_off: bool = True
14
+ wal_level_minimal: bool = True
15
+ disable_wal_senders: bool = True
16
+ disable_archiving: bool = True
17
+ autovacuum_off: bool = True
18
+ jit_off: bool = True
19
+ no_locale: bool = True
20
+ shared_buffers_mb: int | None = 128
21
+ work_mem_mb: int | None = None
22
+ maintenance_work_mem_mb: int | None = None
23
+
24
+
25
+ @dataclass(frozen=True, slots=True, kw_only=True)
26
+ class Container:
27
+ name: str = "elefast"
28
+ image: str = "postgres"
29
+ version: str = "latest"
30
+ ports: dict[str, str] = field(default_factory=lambda: {"5432": "5432"})
31
+
32
+
33
+ @dataclass(frozen=True, slots=True, kw_only=True)
34
+ class Credentials:
35
+ user: str = "postgres"
36
+ password: str = "elefast"
37
+ host: str = "127.0.0.1"
38
+ port: int = 5432
39
+
40
+
41
+ @dataclass(frozen=True, slots=True, kw_only=True)
42
+ class Configuration:
43
+ container: Container = Container()
44
+ credentials: Credentials = Credentials()
45
+ optimizations: Optimizations = Optimizations()
@@ -0,0 +1,29 @@
1
+ from pathlib import Path
2
+ from tempfile import gettempdir
3
+ from filelock import FileLock
4
+ from docker import DockerClient
5
+ from elefast.docker.orchestration import ensure_db_server_started
6
+ from elefast.docker.configuration import Configuration
7
+ from sqlalchemy import URL
8
+
9
+
10
+ def postgres(
11
+ driver: str,
12
+ docker: DockerClient | None = None,
13
+ config: Configuration | None = None,
14
+ keep_container_around: bool = True,
15
+ ) -> URL:
16
+ docker = docker if docker else DockerClient.from_env()
17
+ config = config if config else Configuration()
18
+
19
+ with FileLock(Path(gettempdir()) / "elefast-docker.lock"):
20
+ ensure_db_server_started(
21
+ docker=docker, config=config, keep_container_around=keep_container_around
22
+ )
23
+ return URL.create(
24
+ drivername=f"postgresql+{driver}",
25
+ username=config.credentials.user,
26
+ password=config.credentials.password,
27
+ host=config.credentials.host,
28
+ port=config.credentials.port,
29
+ )
@@ -0,0 +1,83 @@
1
+ from typing import cast
2
+
3
+ from docker import DockerClient
4
+ from docker.models.containers import Container
5
+
6
+ from elefast.docker.configuration import Configuration
7
+
8
+
9
+ def get_docker() -> DockerClient:
10
+ return DockerClient.from_env()
11
+
12
+
13
+ def ensure_db_server_started(
14
+ docker: DockerClient | None = None,
15
+ config: Configuration | None = None,
16
+ keep_container_around: bool = False,
17
+ ) -> Container:
18
+ if docker is None:
19
+ docker = get_docker()
20
+ if config is None:
21
+ config = Configuration()
22
+
23
+ if container := get_db_server_container(docker, config.container.name):
24
+ return container
25
+ return start_db_server_container(docker, config, keep_container_around)
26
+
27
+
28
+ def get_db_server_container(docker: DockerClient, name: str) -> Container | None:
29
+ containers = cast(list[Container], docker.containers.list())
30
+ for container in containers:
31
+ if container.name == name:
32
+ return container
33
+
34
+
35
+ def start_db_server_container(
36
+ docker: DockerClient, config: Configuration, keep_container_around: bool
37
+ ) -> Container:
38
+ optimizations = config.optimizations
39
+ command: list[str] = []
40
+ env: dict[str, str] = {
41
+ "POSTGRES_USER": config.credentials.user,
42
+ "POSTGRES_PASSWORD": config.credentials.password,
43
+ }
44
+
45
+ if optimizations.fsync_off:
46
+ command += ["-c", "fsync=off"]
47
+ if optimizations.synchronous_commit_off:
48
+ command += ["-c", "synchronous_commit=off"]
49
+ if optimizations.full_page_writes_off:
50
+ command += ["-c", "full_page_writes=off"]
51
+ if optimizations.wal_level_minimal:
52
+ command += ["-c", "wal_level=minimal"]
53
+ if optimizations.disable_wal_senders:
54
+ command += ["-c", "max_wal_senders=0"]
55
+ if optimizations.disable_archiving:
56
+ command += ["-c", "archive_mode=off"]
57
+ if optimizations.autovacuum_off:
58
+ command += ["-c", "autovacuum=off"]
59
+ if optimizations.jit_off:
60
+ command += ["-c", "jit=off"]
61
+ if optimizations.shared_buffers_mb is not None:
62
+ command += ["-c", f"shared_buffers={optimizations.shared_buffers_mb}MB"]
63
+ if optimizations.work_mem_mb is not None:
64
+ command += ["-c", f"work_mem={optimizations.work_mem_mb}MB"]
65
+ if optimizations.maintenance_work_mem_mb is not None:
66
+ command += [
67
+ "-c",
68
+ f"maintenance_work_mem={optimizations.maintenance_work_mem_mb}MB",
69
+ ]
70
+ if optimizations.no_locale:
71
+ env["POSTGRES_INITDB_ARGS"] = "--no-locale"
72
+
73
+ return docker.containers.run(
74
+ image=f"{config.container.image}:{config.container.version}",
75
+ name=config.container.name,
76
+ ports=config.container.ports,
77
+ environment=env,
78
+ tmpfs={"/var/lib/postgresql": f"rw,size={optimizations.tmpfs_size_mb}m"}
79
+ if optimizations.tmpfs_size_mb is not None
80
+ else {},
81
+ remove=not keep_container_around,
82
+ detach=True,
83
+ )
@@ -0,0 +1,6 @@
1
+ class ElefastError(Exception):
2
+ """The base class for errors of elefat"""
3
+
4
+
5
+ class DatabaseNotReadyError(ElefastError):
6
+ """We tried to connect to the DB, but even after lots of attempts we could not run a simple query."""
@@ -0,0 +1,119 @@
1
+ from __future__ import annotations
2
+ from uuid import uuid4
3
+ import time
4
+ from elefast.errors import DatabaseNotReadyError
5
+
6
+ from collections.abc import Callable
7
+ from contextlib import AbstractContextManager
8
+
9
+ from sqlalchemy import URL, Engine, MetaData, text, create_engine, NullPool
10
+ from sqlalchemy.orm import Session, sessionmaker
11
+
12
+
13
+ type CanBeTurnedIntoEngine = Engine | URL | str
14
+
15
+
16
+ class Database(AbstractContextManager):
17
+ def __init__(
18
+ self,
19
+ engine: Engine,
20
+ server: DatabaseServer,
21
+ sessionmaker_factory: Callable[[Engine], Callable[[], Session]] = sessionmaker,
22
+ ) -> None:
23
+ self.engine = engine
24
+ self.server = server
25
+ self.sessionmaker = sessionmaker_factory(self.engine)
26
+ assert self.engine.url.database
27
+ self.name = self.engine.url.database
28
+
29
+ def __exit__(self, exc_type, exc, tb):
30
+ self.drop()
31
+
32
+ @property
33
+ def url(self) -> URL:
34
+ return self.engine.url
35
+
36
+ def drop(self) -> None:
37
+ self.engine.dispose()
38
+ self.server.drop_database(self.name)
39
+
40
+ def session(self) -> Session:
41
+ return self.sessionmaker()
42
+
43
+
44
+ class DatabaseServer:
45
+ def __init__(
46
+ self,
47
+ engine: CanBeTurnedIntoEngine,
48
+ metadata: MetaData | None = None,
49
+ ) -> None:
50
+ self._metadata = metadata
51
+ self._engine = _build_engine(engine)
52
+ self._template_db_name: str | None = None
53
+
54
+ @property
55
+ def url(self) -> URL:
56
+ return self._engine.url
57
+
58
+ def ensure_is_ready(
59
+ self, timeout: int | float = 30, interval: int | float = 0.5
60
+ ) -> None:
61
+ deadline = time.monotonic() + timeout
62
+ attempts = 0
63
+
64
+ while True:
65
+ try:
66
+ with self._engine.connect() as conn:
67
+ conn.execute(text("SELECT 1"))
68
+ return
69
+ except Exception as error:
70
+ attempts += 1
71
+ if time.monotonic() >= deadline:
72
+ raise DatabaseNotReadyError(
73
+ f"Reached the configured timeout of {timeout} seconds after {attempts} attempts connecting to the database."
74
+ ) from error
75
+ time.sleep(interval)
76
+
77
+ def create_database(self) -> Database:
78
+ template_db = self._template_db_name
79
+ if template_db is None:
80
+ engine = _prepare_database(self._engine)
81
+ if self._metadata:
82
+ with engine.begin() as connection:
83
+ self._metadata.drop_all(bind=connection)
84
+ self._metadata.create_all(bind=connection)
85
+ engine.dispose()
86
+ template_db = engine.url.database
87
+ assert isinstance(template_db, str)
88
+ self._template_db_name = template_db
89
+
90
+ engine = _prepare_database(self._engine, template=template_db)
91
+ return Database(engine=engine, server=self)
92
+
93
+ def drop_database(self, name: str) -> None:
94
+ with self._engine.begin() as connection:
95
+ statement = f'DROP DATABASE "{name}"'
96
+ connection.execute(text(statement))
97
+
98
+
99
+ def _build_engine(input: CanBeTurnedIntoEngine) -> Engine:
100
+ if isinstance(input, Engine):
101
+ return input
102
+ if isinstance(input, URL | str):
103
+ return create_engine(input, isolation_level="autocommit", poolclass=NullPool)
104
+ raise TypeError()
105
+
106
+
107
+ def _prepare_database(
108
+ engine: Engine, encoding: str = "utf8", template: str | None = None
109
+ ) -> Engine:
110
+ database = f"pytest-elephantastic-{uuid4()}"
111
+ with engine.begin() as connection:
112
+ statement = (
113
+ f"CREATE DATABASE \"{database}\" ENCODING '{encoding}' TEMPLATE template0"
114
+ if template is None
115
+ else f'CREATE DATABASE "{database}" WITH TEMPLATE "{template}"'
116
+ )
117
+ connection.execute(text(statement))
118
+ connection.commit()
119
+ return create_engine(engine.url.set(database=database))