add support for null poolclass

Author: OriHoch

open_bus_stride_db/db.py

@@ -9,18 +9,32 @@
 SQLALCHEMY_APPLICATION_NAME = os.getenv("SQLALCHEMY_APPLICATION_NAME", "db")
 SQLALCHEMY_APPLICATION_VERSION = os.getenv("SQLALCHEMY_APPLICATION_VERSION", "-")
 
-
-engine = create_engine(
-    os.environ.get('SQLALCHEMY_URL', 'postgresql://postgres:123456@localhost'),
+_create_engine_kwargs = dict(
     future=True,
     connect_args={
         "options": "-c timezone=utc",
         "application_name": f'{SQLALCHEMY_APPLICATION_NAME} {SQLALCHEMY_APPLICATION_VERSION}'[:64],
     },
     echo=bool(os.environ.get('SQLALCHEMY_ECHO')),
-    pool_size=int(os.environ.get('SQLALCHEMY_POOL_SIZE', 10)),
-    max_overflow=int(os.environ.get('SQLALCHEMY_MAX_OVERFLOW', 20)),
-    pool_timeout=int(os.environ.get('SQLALCHEMY_POOL_TIMEOUT', 30)),
+)
+if os.getenv("SQLALCHEMY_POOLCLASS_NULLPOOL") == "yes":
+    from sqlalchemy.pool import NullPool
+    _create_engine_kwargs.update(
+        poolclass=NullPool,
+        pool_pre_ping=True,
+    )
+else:
+    _create_engine_kwargs.update(
+        dict(
+            pool_size=int(os.environ.get('SQLALCHEMY_POOL_SIZE', 10)),
+            max_overflow=int(os.environ.get('SQLALCHEMY_MAX_OVERFLOW', 20)),
+            pool_timeout=int(os.environ.get('SQLALCHEMY_POOL_TIMEOUT', 30)),
+        )
+    )
+
+engine = create_engine(
+    os.environ.get('SQLALCHEMY_URL', 'postgresql://postgres:123456@localhost'),
+    **_create_engine_kwargs,
 )
 _sessionmaker = sessionmaker(bind=engine, future=True, autoflush=False, autocommit=False)
 

tests/__init__.py

@@ -0,0 +1,47 @@
+import os
+import time
+import importlib
+import subprocess
+
+from open_bus_stride_db import db as db_module
+
+import pytest
+
+
+@pytest.fixture(scope="session")
+def get_db():
+    create_container = True
+    if os.getenv("KEEP_TEST_POSTGRES_CONTAINER") == "yes":
+        if subprocess.call([
+            "docker", "container", "inspect", "open-bus-stride-db-test-postgres"
+        ], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) == 0:
+            create_container = False
+    if create_container:
+        subprocess.call([
+            "docker", "rm", "-f", "open-bus-stride-db-test-postgres"
+        ])
+        subprocess.check_call([
+            "docker", "run", "--name", "open-bus-stride-db-test-postgres",
+            "-e", "POSTGRES_USER=postgres",
+            "-e", "POSTGRES_PASSWORD=123456",
+            "-p", "5432:5432",
+            "-d", "postgres:14"
+        ])
+        time.sleep(3)
+
+    def get_db_(poolclass_nullpool=False):
+        os.environ["SQLALCHEMY_URL"] = "postgresql://postgres:123456@localhost"
+        if poolclass_nullpool:
+            os.environ["SQLALCHEMY_POOLCLASS_NULLPOOL"] = "yes"
+        elif "SQLALCHEMY_POOLCLASS_NULLPOOL" in os.environ:
+            del os.environ["SQLALCHEMY_POOLCLASS_NULLPOOL"]
+        importlib.reload(db_module)
+        return db_module
+
+    try:
+        yield get_db_
+    finally:
+        if os.getenv("KEEP_TEST_POSTGRES_CONTAINER") != "yes":
+            subprocess.call([
+                "docker", "rm", "-f", "open-bus-stride-db-test-postgres"
+            ])

tests/conftest.py

@@ -0,0 +1 @@
+pytest==9.0.2

tests/requirements.txt

@@ -0,0 +1,13 @@
+
+def test(get_db):
+    db = get_db()
+    with db.get_session() as session:
+        assert session.execute("SELECT 1").scalar() == 1
+    assert "poolclass" not in db._create_engine_kwargs
+
+
+def test_nullpool(get_db):
+    db = get_db(poolclass_nullpool=True)
+    with db.get_session() as session:
+        assert session.execute("SELECT 1").scalar() == 1
+    assert db._create_engine_kwargs["poolclass"].__name__ == "NullPool"