Test on both postgres and sqlite in CI

Author: jbothma

nomenklatura/db.py

@@ -1,6 +1,7 @@
 from contextlib import contextmanager
 from functools import cache
 from typing import Generator, Optional, Union
+import logging
 
 from sqlalchemy import MetaData, create_engine
 from sqlalchemy.dialects.mysql import insert as mysql_insert
@@ -13,14 +14,31 @@
 Conn = Connection
 Connish = Optional[Connection]
 
+WARNED_DB_URL = False
+
+logger = logging.getLogger(__name__)
+
 
 @cache
-def get_engine(url: str = settings.DB_URL) -> Engine:
-    if settings.TESTING:
-        url = "sqlite:///:memory:"
-    # if url.lower().startswith('sqlite'):
-    #     return create_engine(url)
-    return create_engine(url, pool_size=settings.DB_POOL_SIZE)
+def get_engine(url: Optional[str] = None) -> Engine:
+    if not url:
+        if settings.DB_URL:
+            url = settings.DB_URL
+        else:
+            url = f"sqlite:///{settings.DB_PATH.as_posix()}"
+
+            global WARNED_DB_URL
+            if not WARNED_DB_URL:
+                logger.warning(f"No DB_URL set. Using {url}")
+                WARNED_DB_URL = True
+
+    connect_args = {}
+    if url.startswith("postgres"):
+        connect_args["options"] = "-c statement_timeout=3000"
+
+    return create_engine(
+        url, pool_size=settings.DB_POOL_SIZE, connect_args=connect_args
+    )
 
 
 @cache

nomenklatura/resolver/resolver.py

@@ -96,13 +96,19 @@ def commit(self) -> None:
         self._conn.close()
         self._conn = None
 
-    def rollback(self) -> None:
-        if self._transaction is None or self._conn is None:
-            raise RuntimeError("No transaction to rollback.")
-        self._transaction.rollback()
-        self._transaction = None
-        self._conn.close()
-        self._conn = None
+    def rollback(self, force=False) -> None:
+        if self._transaction is not None:
+            self._transaction.rollback()
+            self._transaction = None
+        else:
+            if not force:
+                raise RuntimeError("No transaction to rollback.")
+        if self._conn is not None:
+            self._conn.close()
+            self._conn = None
+        else:
+            if not force:
+                raise RuntimeError("No connection to close.")
 
     def _get_connection(self) -> Connection:
         if self._transaction is None or self._conn is None:
@@ -487,6 +493,11 @@ def prune(self) -> None:
         self._invalidate()
 
     def apply_statement(self, stmt: Statement) -> Statement:
+        """
+        Canonicalise the entity ID.
+
+        Doesn't canonicalise entity ID values.
+        """
         if stmt.entity_id is not None:
             stmt.canonical_id = self.get_canonical(stmt.entity_id)
         if stmt.prop_type == registry.entity.name:

nomenklatura/settings.py

@@ -4,7 +4,7 @@
 TESTING = False
 
 DB_PATH = Path("nomenklatura.db").resolve()
-DB_URL = env_str("NOMENKLATURA_DB_URL", f"sqlite:///{DB_PATH.as_posix()}")
+DB_URL = env_str("NOMENKLATURA_DB_URL", "")
 DB_POOL_SIZE = int(env_str("NOMENKLATURA_DB_POOL_SIZE", "5"))
 
 REDIS_URL = env_str("NOMENKLATURA_REDIS_URL", "")

tests/conftest.py

@@ -1,5 +1,6 @@
 import json
 import shutil
+from sqlalchemy import MetaData
 import yaml
 import pytest
 from pathlib import Path
@@ -21,7 +22,12 @@
 
 @pytest.fixture(autouse=True)
 def wrap_test():
+    if not settings.DB_URL:
+        settings.DB_URL = "sqlite:///:memory:"
     yield
+    # Dispose of connections to let open transactions for resources not
+    # managed by the setup/teardown abort.
+    get_engine().dispose()
     get_engine.cache_clear()
     get_redis.cache_clear()
     get_metadata.cache_clear()
@@ -55,13 +61,24 @@ def donations_json(donations_path):
 @pytest.fixture(scope="function")
 def resolver():
     resolver = Resolver[CompositeEntity].make_default()
-    resolver.begin()
     yield resolver
-    resolver.rollback()
+    resolver.rollback(force=True)
+    resolver._table.drop(resolver._engine)
+
+
+@pytest.fixture(scope="function")
+def other_table_resolver():
+    engine = get_engine()
+    meta = MetaData()
+    resolver = Resolver(engine, meta, create=True, table_name="another_table")
+    yield resolver
+    resolver.rollback(force=True)
+    resolver._table.drop(engine)
 
 
 @pytest.fixture(scope="function")
 def dstore(donations_path, resolver) -> SimpleMemoryStore:
+    resolver.begin()
     return load_entity_file_store(donations_path, resolver)
 
 

tests/enrich/test_nominatim.py

@@ -88,6 +88,7 @@ def test_nominatim_match():
 
 
 def test_nominatim_match_list(resolver: Resolver[CompositeEntity]):
+    resolver.begin()
     enricher = load_enricher()
 
     full = "Kopenhagener Str. 47, Berlin"
@@ -116,6 +117,7 @@ def test_nominatim_enrich():
 
 
 def test_nominatim_enrich_list(resolver: Resolver[CompositeEntity]):
+    resolver.begin()
     enricher = load_enricher()
 
     full = "Kopenhagener Str. 47, Berlin"

tests/store/test_leveldb.py

@@ -26,6 +26,7 @@
 def test_leveldb_store_basics(
     test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 ):
+    resolver.begin()
     path = Path(tempfile.mkdtemp()) / "leveldb"
     store = LevelDBStore(test_dataset, resolver, path)
     entity = CompositeEntity.from_data(test_dataset, PERSON)
@@ -52,6 +53,7 @@ def test_leveldb_store_basics(
 def test_leveldb_graph_query(
     donations_path: Path, test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 ):
+    resolver.begin()
     path = Path(tempfile.mkdtemp()) / "xxx"
     store = LevelDBStore(test_dataset, resolver, path)
     assert len(list(store.view(test_dataset).entities())) == 0

tests/store/test_memory.py

@@ -22,6 +22,7 @@
 
 
 def test_basic_store(test_dataset: Dataset, resolver: Resolver[CompositeEntity]):
+    resolver.begin()
     store = MemoryStore(test_dataset, resolver)
     entity = CompositeEntity.from_data(test_dataset, PERSON)
     entity_ext = CompositeEntity.from_data(test_dataset, PERSON_EXT)

tests/store/test_redis.py

@@ -24,6 +24,7 @@
 
 
 def test_redis_store_basics(test_dataset: Dataset, resolver: Resolver[CompositeEntity]):
+    resolver.begin()
     redis = fakeredis.FakeStrictRedis(version=6, decode_responses=False)
     store = RedisStore(test_dataset, resolver, db=redis)
     entity = CompositeEntity.from_data(test_dataset, PERSON)
@@ -50,6 +51,7 @@ def test_redis_store_basics(test_dataset: Dataset, resolver: Resolver[CompositeE
 def test_leveldb_graph_query(
     donations_path: Path, test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 ):
+    resolver.begin()
     redis = fakeredis.FakeStrictRedis(version=6, decode_responses=False)
     store = RedisStore(test_dataset, resolver, db=redis)
     assert len(list(store.view(test_dataset).entities())) == 0

tests/store/test_resolved.py

@@ -28,6 +28,7 @@
 
 
 def test_store_basics(test_dataset: Dataset, resolver: Resolver[CompositeEntity]):
+    resolver.begin()
     redis = fakeredis.FakeStrictRedis(version=6, decode_responses=False)
     store = ResolvedStore(test_dataset, resolver, db=redis)
     entity = CompositeEntity.from_data(test_dataset, PERSON)
@@ -59,6 +60,7 @@ def test_store_basics(test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 def test_graph_query(
     donations_path: Path, test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 ):
+    resolver.begin()
     redis = fakeredis.FakeStrictRedis(version=6, decode_responses=False)
     store = ResolvedStore(test_dataset, resolver, db=redis)
     assert len(list(store.view(test_dataset).entities())) == 0
@@ -108,6 +110,7 @@ def test_graph_query(
 def test_custom_functions(
     donations_path: Path, test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 ):
+    resolver.begin()
     redis = fakeredis.FakeStrictRedis(version=6, decode_responses=False)
     prefix = "test123"
     mem_store = MemoryStore(test_dataset, resolver)

tests/store/test_stores.py

@@ -74,6 +74,7 @@ def test_store_sql(
     donations_json: List[Dict[str, Any]],
     resolver: Resolver[CompositeEntity],
 ):
+    resolver.begin()
     uri = f"sqlite:///{tmp_path / 'test.db'}"
     store = SQLStore(dataset=test_dataset, linker=resolver, uri=uri)
     assert str(store.engine.url) == uri
@@ -85,6 +86,7 @@ def test_store_memory(
     donations_json: List[Dict[str, Any]],
     resolver: Resolver[CompositeEntity],
 ):
+    resolver.begin()
     store = SimpleMemoryStore(dataset=test_dataset, linker=resolver)
     assert _run_store_test(store, test_dataset, donations_json)
 
@@ -95,6 +97,7 @@ def test_store_level(
     donations_json: List[Dict[str, Any]],
     resolver: Resolver[CompositeEntity],
 ):
+    resolver.begin()
     path = tmp_path / "level.db"
     store = LevelDBStore(dataset=test_dataset, linker=resolver, path=path)
     assert _run_store_test(store, test_dataset, donations_json)

tests/store/test_versioned.py

@@ -27,6 +27,7 @@
 
 
 def test_store_basics(test_dataset: Dataset, resolver: Resolver[CompositeEntity]):
+    resolver.begin()
     redis = fakeredis.FakeStrictRedis(version=6, decode_responses=False)
     store = VersionedRedisStore(test_dataset, resolver, db=redis)
     assert len(list(store.view(test_dataset).statements())) == 0
@@ -64,6 +65,7 @@ def test_store_basics(test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 def test_graph_query(
     donations_path: Path, test_dataset: Dataset, resolver: Resolver[CompositeEntity]
 ):
+    resolver.begin()
     redis = fakeredis.FakeStrictRedis(version=6, decode_responses=False)
     store = VersionedRedisStore(test_dataset, resolver, db=redis)
     assert len(list(store.view(test_dataset).entities())) == 0