postgresql-testing 0.0.2__tar.gz → 0.0.4__tar.gz

{postgresql_testing-0.0.2 → postgresql_testing-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: postgresql-testing
-Version: 0.0.2
+Version: 0.0.4
 Summary: PostgreSQL testing helpers
 Author: Oli Russell
 Author-email: Oli Russell <mrxoliver@gmail.com>

{postgresql_testing-0.0.2 → postgresql_testing-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "postgresql-testing"
-version = "0.0.2"
+version = "0.0.4"
 description = "PostgreSQL testing helpers"
 readme = "README.md"
 authors = [

{postgresql_testing-0.0.2 → postgresql_testing-0.0.4}/src/postgresql_testing/__init__.py

@@ -14,6 +14,9 @@ SUPERUSER = "postgres"
 ROOT_DATABASE = "postgres"
 
 
+class PostgresqlTestingError(RuntimeError): ...
+
+
 @dataclass(kw_only=True)
 class ClusterConfig:
     host: str
@@ -73,7 +76,7 @@ def initdb(
 ) -> None:
     if directory.exists():
         if on_existing == "raise":
-            raise RuntimeError(f"Directory {directory} already exists")
+            raise PostgresqlTestingError(f"Directory {directory} already exists")
         if on_existing == "replace":
             shutil.rmtree(directory)
         if on_existing == "use":
@@ -116,8 +119,16 @@ def serve(c: ClusterConfig) -> Iterator[None]:
             _try_connect(c.superuser().dsn)
             yield
         finally:
+            # Stop any ongoing queries
+            with psycopg.connect(c.superuser().dsn) as conn, conn.cursor() as cur:
+                cur.execute("SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE pid <> pg_backend_pid()")
+            # Shutdown
             process.terminate()
-            process.wait()
+            try:
+                process.wait(timeout=1)
+            except subprocess.TimeoutExpired:
+                process.kill()
+                process.wait(timeout=5)
 
 
 def ensure_user(c: DatabaseConfig) -> None:
@@ -132,7 +143,7 @@ def create_database(
     c: DatabaseConfig,
     *,
     template: str | None = None,
-    on_existing: Literal["raise", "replace"] = "replace",
+    on_existing: Literal["raise", "use", "replace"] = "replace",
 ) -> None:
     template_str = "" if template is None else f'WITH TEMPLATE "{template}"'
     sql_create_database = f'CREATE DATABASE "{c.database}" {template_str} OWNER "{c.user}" STRATEGY=FILE_COPY'
@@ -143,7 +154,9 @@ def create_database(
             cur.execute(sql_create_database)
         except psycopg.errors.DuplicateDatabase:
             if on_existing == "raise":
-                raise RuntimeError(f"Database {c.database} already exists")
+                raise PostgresqlTestingError(f"Database {c.database} already exists")
+            if on_existing == "use":
+                pass
             if on_existing == "replace":
                 cur.execute(f'DROP DATABASE "{c.database}"')
                 cur.execute(sql_create_database)
@@ -164,7 +177,7 @@ def create_template(
         template_exists = bool(cur.execute(f"SELECT 1 FROM pg_database WHERE datname = '{template}'").fetchall())
         if template_exists:
             if on_existing == "raise":
-                raise RuntimeError(f"Template database {template} already exists")
+                raise PostgresqlTestingError(f"Template database {template} already exists")
             if on_existing == "replace":
                 cur.execute(f"UPDATE pg_database SET datistemplate = false WHERE datname='{template}'")
                 cur.execute(f'DROP DATABASE "{template}"')
@@ -182,7 +195,7 @@ def dump_archive(
 ) -> None:
     if archive.exists():
         if on_existing == "raise":
-            raise RuntimeError(f"Archive {archive} already exists")
+            raise PostgresqlTestingError(f"Archive {archive} already exists")
         if on_existing == "replace":
             archive.unlink()
 
@@ -224,4 +237,4 @@ def _try_connect(dsn: str) -> None:
             return
         except psycopg.OperationalError:
             pass
-    raise RuntimeError(f"Could not connect to {dsn}")
+    raise PostgresqlTestingError(f"Could not connect to {dsn}")