Repository: cogent3/cogent3
Pull request #1459 (2023-05-29 03:12:44)

ENH: support db = a sqlite Connection instance
[NEW] The SqliteAnnotationDbMixin._setup_db() method can now
    handle using a connection instance. If the incoming schema
    is a subset of the classes schema, it is considered compatible.
    Any other relationship triggers a TypeError. Check is implemented
    in _compatible_schema()


Patch:
diff --git a/src/cogent3/core/annotation_db.py b/src/cogent3/core/annotation_db.py
index 74f59f229..1ab52f4ad 100644
--- a/src/cogent3/core/annotation_db.py
+++ b/src/cogent3/core/annotation_db.py
@@ -409,6 +409,26 @@ def _count_records_sql(
     return sql, vals
 
 
+def _compatible_schema(db: sqlite3.Connection, schema: dict[str, set]) -> bool:
+    """ensures the db instance is compatible with schema"""
+    for table in db.execute(
+        "SELECT name FROM sqlite_master WHERE type='table'"
+    ).fetchall():
+        table = table["name"]
+        if table not in schema:
+            return False
+
+        db_schema = {
+            (row["name"], row["type"])
+            for row in db.execute(f"pragma table_info({table!r})").fetchall()
+        }
+
+        if len(db_schema & schema[table]) != len(db_schema):
+            return False
+
+    return True
+
+
 class SqliteAnnotationDbMixin:
     # table schema for user provided annotations
     _user_schema = {
@@ -488,12 +508,25 @@ class SqliteAnnotationDbMixin:
     def table_names(self) -> tuple[str]:
         return self._table_names
 
-    def _setup_db(self, db: typing.Optional[SupportsFeatures]) -> None:
+    def _setup_db(self, db: typing.Union[SupportsFeatures, sqlite3.Connection]) -> None:
         """initialises the db, using the db passed to the constructor"""
         if isinstance(db, self.__class__):
             self._db = db.db
             return
 
+        if isinstance(db, sqlite3.Connection):
+            schema = {}
+            for table_name in self.table_names:
+                attr = getattr(self, f"_{table_name}_schema")
+                schema[table_name] = set(attr.items())
+
+            if not _compatible_schema(db, schema):
+                raise TypeError("incompatible schema")
+
+            self._db = db
+            self._init_tables()
+            return
+
         if db and not self.compatible(db):
             raise TypeError(f"cannot initialise annotation db from {type(db)}")
 


Test patch:
diff --git a/tests/test_core/test_annotation_db.py b/tests/test_core/test_annotation_db.py
index e889c691a..b63cf9a0f 100644
--- a/tests/test_core/test_annotation_db.py
+++ b/tests/test_core/test_annotation_db.py
@@ -69,6 +69,15 @@ def test_constructor_db_fail(db_name, cls, request):
         cls(db=db)
 
 
+@pytest.mark.parametrize(
+    "db_name,cls", (("gff_db", GenbankAnnotationDb), ("gb_db", GffAnnotationDb))
+)
+def test_constructor_wrong_db_schema(db_name, cls, request):
+    db = request.getfixturevalue(db_name)
+    with pytest.raises(TypeError):
+        cls(db=db.db)
+
+
 @pytest.mark.parametrize(
     "db_name,cls",
     (
@@ -78,12 +87,27 @@ def test_constructor_db_fail(db_name, cls, request):
         ("anno_db", GenbankAnnotationDb),
     ),
 )
-def test_constructor_db_works(db_name, cls, request):
+def test_constructor_db_instance_works(db_name, cls, request):
     # only compatible db's used to init
     db = request.getfixturevalue(db_name)
     cls(db=db)
 
 
+@pytest.mark.parametrize(
+    "db_name,cls",
+    (
+        ("gff_db", GffAnnotationDb),
+        ("anno_db", GffAnnotationDb),
+        ("gb_db", GenbankAnnotationDb),
+        ("anno_db", GenbankAnnotationDb),
+    ),
+)
+def test_constructor_db_connection_works(db_name, cls, request):
+    # only compatible db's used to init
+    db = request.getfixturevalue(db_name)
+    cls(db=db.db)
+
+
 def test_gff_describe(gff_db):
     result = gff_db.describe
     assert isinstance(result, _Table)