diff --git a/README.md b/README.md index 50463b4..de56ad0 100644 --- a/README.md +++ b/README.md @@ -240,11 +240,25 @@ worktree and wants to know exactly which DB(s) go with it. Container connection defaults (`localhost:54322`, `postgres`/`testpwd`) can be overridden with `PGDEVKIT_TESTDB_HOST`, `PGDEVKIT_TESTDB_PORT`, -`PGDEVKIT_TESTDB_USER`, `PGDEVKIT_TESTDB_PASSWORD`. Before touching the -Docker API, pgdevkit first checks (with a short timeout) whether Postgres -is already reachable at that address and skips container management if so. -Set `PGDEVKIT_SKIP_CONTAINER=1` to always assume it's already there and skip -that check too. +`PGDEVKIT_TESTDB_USER`, `PGDEVKIT_TESTDB_PASSWORD`, and the image with +`PGDEVKIT_TESTDB_IMAGE`. Before touching the Docker API, pgdevkit first +checks (with a short timeout) whether Postgres is already reachable at that +address and skips container management if so. Set `PGDEVKIT_SKIP_CONTAINER=1` +to always assume it's already there and skip that check too. + +Docker-level resource knobs are unset by default (Docker/Podman's own +defaults apply) and can be overridden with `PGDEVKIT_TESTDB_SHM_SIZE` (e.g. +`256m`, useful when Docker's 64MB default shared-memory size is too small +for parallel query workers), `PGDEVKIT_TESTDB_MEM_LIMIT` (e.g. `1g`), +`PGDEVKIT_TESTDB_CPUS` (e.g. `2` or `0.5`), and `PGDEVKIT_TESTDB_TMPFS` — +docker-CLI `--tmpfs` syntax (`/path:options`), semicolon-separated for +multiple mounts, e.g. `/var/lib/postgresql/data:size=512m`. + +All of the above (including `_HOST`/`_PORT`/`_USER`/`_PASSWORD`/`_IMAGE`) +only take effect when pgdevkit actually creates the container, not when it +finds and restarts an existing stopped one with the same name — remove the +existing container (`docker rm -f pgdevkit-postgres`) first if you've +changed any of these on a machine that already has one. Container management goes through the Docker API (the `docker` package, `docker.from_env()`, falling back to Podman's rootful/rootless socket) — it @@ -270,6 +284,17 @@ Container defaults (`localhost:14330`, `sa`/a generated complexity-valid password) can be overridden with `PGDEVKIT_TESTDB_MSSQL_HOST`, `_PORT`, `_USER`, `_PASSWORD`, `_IMAGE`, `_MEMORY_LIMIT_MB`. The container only bootstraps the `sa` login — additional logins are a known limitation. + +The same Docker-level resource knobs as the Postgres container are +available here too: `PGDEVKIT_TESTDB_MSSQL_SHM_SIZE`, +`PGDEVKIT_TESTDB_MSSQL_MEM_LIMIT` (a container-level cgroup limit, distinct +from `_MEMORY_LIMIT_MB` above which only tunes SQL Server's own internal +memory management), `PGDEVKIT_TESTDB_MSSQL_CPUS`, and +`PGDEVKIT_TESTDB_MSSQL_TMPFS`. As with the Postgres container, these only +take effect when pgdevkit creates the container, not when restarting an +existing stopped one (`docker rm -f pgdevkit-mssql` first if you've changed +any of them). + `pgdb testdb shell` execs into [`sqlcmd`](https://github.com/microsoft/go-sqlcmd) (an external prerequisite, the same category as `psql` for the Postgres path) rather than a Python diff --git a/pgdevkit/testdb/_docker.py b/pgdevkit/testdb/_docker.py index 41cfcbf..9a1bebd 100644 --- a/pgdevkit/testdb/_docker.py +++ b/pgdevkit/testdb/_docker.py @@ -39,3 +39,36 @@ def client() -> docker.DockerClient: "Could not reach a Docker-compatible API. Set DOCKER_HOST, or make sure " "Docker or Podman's API socket is running." ) + + +def _parse_tmpfs(spec: str) -> dict[str, str]: + """Parse a docker-CLI-style --tmpfs spec into the dict shape docker-py's + `containers.run(tmpfs=...)` expects. Multiple mounts are semicolon-separated + (commas are already used by docker's own comma-separated mount options), + e.g. "/path:rw,size=64m;/other/path:ro,size=32m", matching `docker run + --tmpfs /path:rw,size=64m --tmpfs /other/path:ro,size=32m`.""" + mounts: dict[str, str] = {} + for entry in spec.split(";"): + entry = entry.strip() + if not entry: + continue + path, _, options = entry.partition(":") + mounts[path] = options + return mounts + + +def resource_kwargs(*, tmpfs: str, shm_size: str, mem_limit: str, cpus: str) -> dict: + """Build the `containers.run()` kwargs for the tmpfs/shm_size/mem_limit/cpus + overrides, omitting any that aren't set so Docker's own defaults apply. + Shared by the Postgres and MSSQL container modules -- these are Docker-level + resource knobs, not anything specific to either engine.""" + kwargs: dict = {} + if tmpfs: + kwargs["tmpfs"] = _parse_tmpfs(tmpfs) + if shm_size: + kwargs["shm_size"] = shm_size + if mem_limit: + kwargs["mem_limit"] = mem_limit + if cpus: + kwargs["nano_cpus"] = int(float(cpus) * 1_000_000_000) + return kwargs diff --git a/pgdevkit/testdb/constants.py b/pgdevkit/testdb/constants.py index bdf0c03..9916041 100644 --- a/pgdevkit/testdb/constants.py +++ b/pgdevkit/testdb/constants.py @@ -5,11 +5,18 @@ from psycopg.conninfo import make_conninfo CONTAINER_NAME = "pgdevkit-postgres" -IMAGE = "pgvector/pgvector:pg18-trixie" +IMAGE = os.environ.get("PGDEVKIT_TESTDB_IMAGE", "pgvector/pgvector:pg18-trixie") HOST = os.environ.get("PGDEVKIT_TESTDB_HOST", "localhost") PORT = int(os.environ.get("PGDEVKIT_TESTDB_PORT", "54322")) USER = os.environ.get("PGDEVKIT_TESTDB_USER", "postgres") PASSWORD = os.environ.get("PGDEVKIT_TESTDB_PASSWORD", "testpwd") +# Docker-level resource overrides, passed through to `containers.run()` -- +# unset (the default) means "let Docker/Podman use their own defaults". +# TMPFS uses docker-CLI --tmpfs syntax; see _docker._parse_tmpfs(). +TMPFS = os.environ.get("PGDEVKIT_TESTDB_TMPFS", "") +SHM_SIZE = os.environ.get("PGDEVKIT_TESTDB_SHM_SIZE", "") +MEM_LIMIT = os.environ.get("PGDEVKIT_TESTDB_MEM_LIMIT", "") +CPUS = os.environ.get("PGDEVKIT_TESTDB_CPUS", "") PG_SPEED_FLAGS = ["-c", "fsync=off", "-c", "synchronous_commit=off", "-c", "full_page_writes=off"] # Production connects with session timezone=UTC; the container image's own default # (baked into its base OS, not something pgdevkit ever set) can differ, silently diff --git a/pgdevkit/testdb/container.py b/pgdevkit/testdb/container.py index f33d0ae..ac8081b 100644 --- a/pgdevkit/testdb/container.py +++ b/pgdevkit/testdb/container.py @@ -7,7 +7,7 @@ import docker.errors import psycopg -from . import constants +from . import _docker, constants from ._docker import client as _client @@ -43,6 +43,12 @@ def _create_container(client: docker.DockerClient) -> None: "PGTZ": "UTC", }, command=["postgres", *constants.PG_STARTUP_FLAGS], + **_docker.resource_kwargs( + tmpfs=constants.TMPFS, + shm_size=constants.SHM_SIZE, + mem_limit=constants.MEM_LIMIT, + cpus=constants.CPUS, + ), ) except docker.errors.APIError as e: if getattr(e, "status_code", None) == 409 or "already in use" in str(e): @@ -79,6 +85,10 @@ def ensure_container() -> None: except docker.errors.NotFound: container = None if container is not None: + # An existing container (even stopped) is just restarted as-is -- + # constants like TMPFS/SHM_SIZE/IMAGE/etc. only take effect via + # _create_container(), so changing them has no effect here until + # the stale container is removed. Documented in the README. if container.status != "running": container.start() else: diff --git a/pgdevkit/testdb/mssql/constants.py b/pgdevkit/testdb/mssql/constants.py index f321adc..4dce5e2 100644 --- a/pgdevkit/testdb/mssql/constants.py +++ b/pgdevkit/testdb/mssql/constants.py @@ -22,6 +22,15 @@ # default -- see validate_sa_password(). PASSWORD = os.environ.get("PGDEVKIT_TESTDB_MSSQL_PASSWORD", "TestPwd!2026") MEMORY_LIMIT_MB = int(os.environ.get("PGDEVKIT_TESTDB_MSSQL_MEMORY_LIMIT_MB", "2048")) +# Docker-level resource overrides, passed through to `containers.run()` -- +# unset (the default) means "let Docker/Podman use their own defaults". Distinct +# from MEMORY_LIMIT_MB above, which only controls SQL Server's own internal memory +# management, not the container's actual cgroup limit. TMPFS uses docker-CLI +# --tmpfs syntax; see _docker._parse_tmpfs(). +TMPFS = os.environ.get("PGDEVKIT_TESTDB_MSSQL_TMPFS", "") +SHM_SIZE = os.environ.get("PGDEVKIT_TESTDB_MSSQL_SHM_SIZE", "") +MEM_LIMIT = os.environ.get("PGDEVKIT_TESTDB_MSSQL_MEM_LIMIT", "") +CPUS = os.environ.get("PGDEVKIT_TESTDB_MSSQL_CPUS", "") def validate_sa_password(password: str) -> None: diff --git a/pgdevkit/testdb/mssql/container.py b/pgdevkit/testdb/mssql/container.py index 343d3d7..fb0d953 100644 --- a/pgdevkit/testdb/mssql/container.py +++ b/pgdevkit/testdb/mssql/container.py @@ -40,6 +40,12 @@ def _create_container(client: docker.DockerClient) -> None: "MSSQL_PID": "Developer", "MSSQL_MEMORY_LIMIT_MB": str(constants.MEMORY_LIMIT_MB), }, + **_docker.resource_kwargs( + tmpfs=constants.TMPFS, + shm_size=constants.SHM_SIZE, + mem_limit=constants.MEM_LIMIT, + cpus=constants.CPUS, + ), ) except docker.errors.APIError as e: if getattr(e, "status_code", None) == 409 or "already in use" in str(e): @@ -83,6 +89,9 @@ def ensure_mssql_container() -> None: except docker.errors.NotFound: container = None if container is not None: + # See the Postgres container's ensure_container() -- an existing + # (even stopped) container is just restarted as-is, so config + # constants only take effect via _create_container(). if container.status != "running": container.start() else: diff --git a/pyproject.toml b/pyproject.toml index 7217cb1..4444bc1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,7 @@ packages = ["pgdevkit"] [project] name = "pgdevkit" -version = "0.7.2" +version = "0.8.0" description = "A helper for developing with Postgres" readme = "README.md" requires-python = ">=3.14" diff --git a/tests/testdb/test_constants.py b/tests/testdb/test_constants.py index 3a927f2..a790ddb 100644 --- a/tests/testdb/test_constants.py +++ b/tests/testdb/test_constants.py @@ -10,12 +10,22 @@ def test_defaults_when_env_vars_unset(monkeypatch): monkeypatch.delenv("PGDEVKIT_TESTDB_PORT", raising=False) monkeypatch.delenv("PGDEVKIT_TESTDB_USER", raising=False) monkeypatch.delenv("PGDEVKIT_TESTDB_PASSWORD", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_IMAGE", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_TMPFS", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_SHM_SIZE", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_MEM_LIMIT", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_CPUS", raising=False) try: importlib.reload(constants) assert constants.HOST == "localhost" assert constants.PORT == 54322 assert constants.USER == "postgres" assert constants.PASSWORD == "testpwd" + assert constants.IMAGE == "pgvector/pgvector:pg18-trixie" + assert constants.TMPFS == "" + assert constants.SHM_SIZE == "" + assert constants.MEM_LIMIT == "" + assert constants.CPUS == "" finally: importlib.reload(constants) @@ -37,3 +47,21 @@ def test_env_vars_override_connection_defaults(monkeypatch): # override values right back in, leaking them into every later test. monkeypatch.undo() importlib.reload(constants) + + +def test_env_vars_override_docker_resource_defaults(monkeypatch): + monkeypatch.setenv("PGDEVKIT_TESTDB_IMAGE", "otherimage:latest") + monkeypatch.setenv("PGDEVKIT_TESTDB_TMPFS", "/var/lib/postgresql/data:size=512m") + monkeypatch.setenv("PGDEVKIT_TESTDB_SHM_SIZE", "256m") + monkeypatch.setenv("PGDEVKIT_TESTDB_MEM_LIMIT", "1g") + monkeypatch.setenv("PGDEVKIT_TESTDB_CPUS", "2") + try: + importlib.reload(constants) + assert constants.IMAGE == "otherimage:latest" + assert constants.TMPFS == "/var/lib/postgresql/data:size=512m" + assert constants.SHM_SIZE == "256m" + assert constants.MEM_LIMIT == "1g" + assert constants.CPUS == "2" + finally: + monkeypatch.undo() + importlib.reload(constants) diff --git a/tests/testdb/test_container.py b/tests/testdb/test_container.py index a338bf5..4177a5e 100644 --- a/tests/testdb/test_container.py +++ b/tests/testdb/test_container.py @@ -79,6 +79,30 @@ def test_create_container_forces_utc_timezone(): assert client.containers.run_kwargs["environment"]["PGTZ"] == "UTC" +def test_create_container_omits_resource_kwargs_by_default(): + client = _FakeClient() + + _create_container(client) + + for key in ("tmpfs", "shm_size", "mem_limit", "nano_cpus"): + assert key not in client.containers.run_kwargs + + +def test_create_container_passes_through_resource_overrides(monkeypatch): + monkeypatch.setattr(constants, "TMPFS", "/var/lib/postgresql/data:size=512m") + monkeypatch.setattr(constants, "SHM_SIZE", "256m") + monkeypatch.setattr(constants, "MEM_LIMIT", "1g") + monkeypatch.setattr(constants, "CPUS", "2") + client = _FakeClient() + + _create_container(client) + + assert client.containers.run_kwargs["tmpfs"] == {"/var/lib/postgresql/data": "size=512m"} + assert client.containers.run_kwargs["shm_size"] == "256m" + assert client.containers.run_kwargs["mem_limit"] == "1g" + assert client.containers.run_kwargs["nano_cpus"] == 2_000_000_000 + + def _admin_dsn() -> str: return constants.conninfo("postgres", connect_timeout=5) diff --git a/tests/testdb/test_container_mssql.py b/tests/testdb/test_container_mssql.py new file mode 100644 index 0000000..339a84a --- /dev/null +++ b/tests/testdb/test_container_mssql.py @@ -0,0 +1,41 @@ +from __future__ import annotations + +from pgdevkit.testdb.mssql import constants +from pgdevkit.testdb.mssql.container import _create_container + + +class _FakeContainers: + def __init__(self): + self.run_kwargs: dict | None = None + + def run(self, image, **kwargs): + self.run_kwargs = kwargs + + +class _FakeClient: + def __init__(self): + self.containers = _FakeContainers() + + +def test_create_container_omits_resource_kwargs_by_default(): + client = _FakeClient() + + _create_container(client) + + for key in ("tmpfs", "shm_size", "mem_limit", "nano_cpus"): + assert key not in client.containers.run_kwargs + + +def test_create_container_passes_through_resource_overrides(monkeypatch): + monkeypatch.setattr(constants, "TMPFS", "/var/opt/mssql:size=512m") + monkeypatch.setattr(constants, "SHM_SIZE", "256m") + monkeypatch.setattr(constants, "MEM_LIMIT", "4g") + monkeypatch.setattr(constants, "CPUS", "2") + client = _FakeClient() + + _create_container(client) + + assert client.containers.run_kwargs["tmpfs"] == {"/var/opt/mssql": "size=512m"} + assert client.containers.run_kwargs["shm_size"] == "256m" + assert client.containers.run_kwargs["mem_limit"] == "4g" + assert client.containers.run_kwargs["nano_cpus"] == 2_000_000_000 diff --git a/tests/testdb/test_docker.py b/tests/testdb/test_docker.py new file mode 100644 index 0000000..26e99ce --- /dev/null +++ b/tests/testdb/test_docker.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from pgdevkit.testdb._docker import _parse_tmpfs, resource_kwargs + + +def test_parse_tmpfs_single_mount(): + assert _parse_tmpfs("/var/lib/postgresql/data:rw,size=512m") == { + "/var/lib/postgresql/data": "rw,size=512m" + } + + +def test_parse_tmpfs_multiple_mounts_semicolon_separated(): + assert _parse_tmpfs("/data:size=512m;/tmp:rw,size=64m") == { + "/data": "size=512m", + "/tmp": "rw,size=64m", + } + + +def test_parse_tmpfs_path_without_options(): + assert _parse_tmpfs("/tmp") == {"/tmp": ""} + + +def test_parse_tmpfs_ignores_blank_entries(): + assert _parse_tmpfs("/data:size=512m;;/tmp:size=64m;") == { + "/data": "size=512m", + "/tmp": "size=64m", + } + + +def test_resource_kwargs_all_unset_returns_empty_dict(): + assert resource_kwargs(tmpfs="", shm_size="", mem_limit="", cpus="") == {} + + +def test_resource_kwargs_only_includes_set_values(): + kwargs = resource_kwargs(tmpfs="/data:size=512m", shm_size="", mem_limit="1g", cpus="") + assert kwargs == {"tmpfs": {"/data": "size=512m"}, "mem_limit": "1g"} + + +def test_resource_kwargs_cpus_converted_to_nano_cpus(): + assert resource_kwargs(tmpfs="", shm_size="", mem_limit="", cpus="1.5") == { + "nano_cpus": 1_500_000_000 + } + + +def test_resource_kwargs_shm_size_passed_through(): + assert resource_kwargs(tmpfs="", shm_size="256m", mem_limit="", cpus="") == {"shm_size": "256m"} diff --git a/tests/testdb/test_mssql_constants.py b/tests/testdb/test_mssql_constants.py index f258056..357b9f8 100644 --- a/tests/testdb/test_mssql_constants.py +++ b/tests/testdb/test_mssql_constants.py @@ -1,7 +1,10 @@ from __future__ import annotations +import importlib + import pytest +from pgdevkit.testdb.mssql import constants from pgdevkit.testdb.mssql.constants import conninfo, validate_sa_password @@ -29,3 +32,34 @@ def test_conninfo_omits_driver_and_app_keys(): assert "Driver=" not in cs assert "APP=" not in cs assert "Database=mydb" in cs + + +def test_docker_resource_defaults_unset(monkeypatch): + monkeypatch.delenv("PGDEVKIT_TESTDB_MSSQL_TMPFS", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_MSSQL_SHM_SIZE", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_MSSQL_MEM_LIMIT", raising=False) + monkeypatch.delenv("PGDEVKIT_TESTDB_MSSQL_CPUS", raising=False) + try: + importlib.reload(constants) + assert constants.TMPFS == "" + assert constants.SHM_SIZE == "" + assert constants.MEM_LIMIT == "" + assert constants.CPUS == "" + finally: + importlib.reload(constants) + + +def test_docker_resource_env_vars_override_defaults(monkeypatch): + monkeypatch.setenv("PGDEVKIT_TESTDB_MSSQL_TMPFS", "/var/opt/mssql:size=512m") + monkeypatch.setenv("PGDEVKIT_TESTDB_MSSQL_SHM_SIZE", "256m") + monkeypatch.setenv("PGDEVKIT_TESTDB_MSSQL_MEM_LIMIT", "4g") + monkeypatch.setenv("PGDEVKIT_TESTDB_MSSQL_CPUS", "2") + try: + importlib.reload(constants) + assert constants.TMPFS == "/var/opt/mssql:size=512m" + assert constants.SHM_SIZE == "256m" + assert constants.MEM_LIMIT == "4g" + assert constants.CPUS == "2" + finally: + monkeypatch.undo() + importlib.reload(constants) diff --git a/uv.lock b/uv.lock index 0fe71f7..698a146 100644 --- a/uv.lock +++ b/uv.lock @@ -313,7 +313,7 @@ wheels = [ [[package]] name = "pgdevkit" -version = "0.7.2" +version = "0.8.0" source = { editable = "." } dependencies = [ { name = "docker" },