Merge pull request #90 from NetherlandsForensicInstitute/feature/asyncify-sqlite-database-interface

Make sqlite interface async, parametrize both database and api tests to use both sqlite and postgresql in tests

Author: akaIDIOT

citatio/README.md

@@ -13,19 +13,19 @@ Configuration and runtime
 The citatio REST API takes 2 configuration options:
 
 - The model to be used for embedding (currently, only `NetherlandsForensicInstitute/ARM64BERT-embedding` is supported);
-- The database to store both assembly and embeddings in (currently, only SQLite is implemented).
+- The database to store both assembly and embeddings in, either SQLite+sqlitevec or PostgreSQL+pgvector.
 
-Both of these can be configured through environment variables:
+These can be configured through environment variables:
 
 - `CITATIO_MODEL`: a local path or huggingface model name (though again, currently on the `ARM64BERT-embedding` model is supported);
-- `CITATIO_SQLITE_DATABASE`: either `:memory:` or a local path to a SQLite database (will be created if it doesn't currently exist).
+- when using SQLite: `CITATIO_DATABASE_SQLITE`: either `:memory:` or a local path to a SQLite database (will be created if it doesn't currently exist).
+- when using PostgreSQL: either `CITATIO_DATABASE_HOST`, `..._PORT`, `..._USER`, `..._PASSWORD`, `..._DATABASE` to connect to the database in question,
+  or `CITATIO_DATABASE_DSN` with the full connection url to connect to that same database.
 
-The defaults for these values are `NetherlandsForensicInstitute/ARM64BERT-embedding` and `:memory:`, 
-resulting in a functioning but non-persistent service.
+The database configuration is required, the default model to be loaded is `NetherlandsForensicInstitute/ARM64BERT-embedding`.
 
 > [!NOTE]  
-> After observing concurrency issues with SQLite and `sqlite-vec`, the REST API is currently served fully serialized and is consequently fairly slow.
-> We're expecting to be able to solve this by using PostgreSQL and `pgvector`.
+> After observing concurrency issues with SQLite and `sqlite-vec`, the REST API is currently served fully serialized when using SQLite and is consequently fairly slow.
 
 Running the REST API service follows the default FastAPI command line setup, 
 where the application is available from the `citatio` module:
@@ -44,7 +44,7 @@ though end users are encouraged to use the [ready-made Ghidra plugin (sententia)
 Prerequisites
 -------------
 
-Python 3.12 or newer with SQLite version 3.35.0 or newer.
+Python 3.13 or newer with SQLite version 3.35.0 or newer.
 
 Requirements
 ------------

citatio/citatio/api.py

@@ -6,7 +6,7 @@
 from fastapi import FastAPI, Request
 from fastapi.params import Body
 
-from citatio.db import PostgreSQLDatabase
+from citatio.db import PostgreSQLDatabase, SQLiteDatabase
 from citatio.models import ControlFlowGraph
 
 
@@ -16,12 +16,19 @@
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     config = confidence.load_name('citatio')
-    if not config.database:
-        raise ValueError(f'missing database configuration, available: {config}')
 
-    app.state.model = ASMEmbedder.from_pretrained(config.model or DEFAULT_MODEL)
+    match config:
+        case {'database.sqlite': name}:
+            # explicit sqlite name to connect to, use SQLiteDatabase
+            database = await SQLiteDatabase.connect(name)
+        case {'database': connect}:
+            # database settings *not* mentioning sqlite, use PostgreSQLDatabase
+            database = await PostgreSQLDatabase.connect(**connect)
+        case _:
+            raise ValueError(f'missing database configuration, available: {config}')
 
-    async with await PostgreSQLDatabase.connect(**config.database) as database:
+    app.state.model = ASMEmbedder.from_pretrained(config.model or DEFAULT_MODEL)
+    async with database:
         app.state.database = database
         yield
 

citatio/citatio/db.py

@@ -8,7 +8,7 @@
 
 
 class Database:
-    def add_function(
+    async def add_function(
         self,
         name: str,
         cfg: dict[int, list[str]],
@@ -19,13 +19,13 @@ def add_function(
     ):
         pass
 
-    def search_function(self, embedding: np.array, top_n: int = 25):
+    async def search_function(self, embedding: np.array, top_n: int = 25):
         pass
 
 
 class SQLiteDatabase(Database):
     @classmethod
-    def from_name(cls, name):
+    async def connect(cls, name):
         # make sure to pass check_same_thread=False, Python 3.11+ has thread-safe sqlite
         connection = sqlite3.connect(name, check_same_thread=False)
         return cls(connection)
@@ -38,18 +38,17 @@ def __init__(self, connection):
             sqlite_vec.load(self.connection)
             self.connection.enable_load_extension(False)
 
-            # NB: read_text() shows up as deprecated in 3.11, it has since been un-deprecated (causing confusing errors)
             schema = resources.read_text('citatio', 'schema-sqlite.sql')
             self.connection.executescript(schema)
 
-    def __enter__(self):
+    async def __aenter__(self):
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
         self.connection.close()
         self.connection = None
 
-    def _insert_or_get_function(self, cursor, cfg, embedding):
+    async def _insert_or_get_function(self, cursor, cfg, embedding):
         # coerce cfg to str for database storage (keep the cfg type responsible for that (de)serialization)
         parameters = (str(cfg),)
         try:
@@ -66,10 +65,10 @@ def _insert_or_get_function(self, cursor, cfg, embedding):
 
         return function_id
 
-    def add_function(self, name, cfg, embedding, binary_name, binary_sha256, model_identifier=None):
+    async def add_function(self, name, cfg, embedding, binary_name, binary_sha256, model_identifier=None):
         with self.connection:
             cursor = self.connection.cursor()
-            function_id = self._insert_or_get_function(cursor, cfg, embedding)
+            function_id = await self._insert_or_get_function(cursor, cfg, embedding)
             cursor.execute(
                 """INSERT INTO labels (function_id, label, binary_name, binary_sha256) VALUES (?, ?, ?, ?)""",
                 (function_id, name, binary_name, binary_sha256),
@@ -78,7 +77,7 @@ def add_function(self, name, cfg, embedding, binary_name, binary_sha256, model_i
             # return the function id for convenience
             return function_id
 
-    def search_function(self, embedding, top_n=25):
+    async def search_function(self, embedding, top_n=25):
         cursor = self.connection.cursor()
         cursor.execute(
             """
@@ -125,33 +124,30 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc_val, exc_tb):
         await self.connection.close()
 
-    async def _insert_or_get_function(self, cfg, embedding):
-        cfg = str(cfg)
-        try:
-            return await self.connection.fetchval(
-                'INSERT INTO functions (cfg, embedding) VALUES ($1, $2) RETURNING id',
-                cfg,
+    async def add_function(self, name, cfg, embedding, binary_name, binary_sha256, model_identifier=None):
+        async with self.connection.transaction():
+            function_id = await self.connection.fetchval(
+                # use PostgreSQL's conflict resolution to issue an update-or-get
+                # NB: the conflict resolution update is idempotent, but needed to make sure RETURNING id works
+                """
+                INSERT INTO functions (cfg, embedding) VALUES ($1, $2)
+                ON CONFLICT (cfg) DO UPDATE SET cfg = EXCLUDED.cfg RETURNING id
+                """,
+                str(cfg),
                 embedding,
             )
-        except asyncpg.IntegrityConstraintViolationError:
-            return await self.connection.fetchval('SELECT id FROM functions WHERE cfg = $1', cfg)
-
-    async def _insert_label(self, function_id, name, binary_name, binary_sha256):
-        await self.connection.execute(
-            'INSERT INTO labels (function_id, label, binary_name, binary_sha256) VALUES ($1, $2, $3, $4)',
-            function_id,
-            name,
-            binary_name,
-            binary_sha256,
-        )
+            await self.connection.execute(
+                'INSERT INTO labels (function_id, label, binary_name, binary_sha256) VALUES ($1, $2, $3, $4)',
+                function_id,
+                name,
+                binary_name,
+                binary_sha256,
+            )
 
-    async def add_function(self, name, cfg, embedding, binary_name, binary_sha256, model_identifier=None):
-        function_id = await self._insert_or_get_function(cfg, embedding)
-        await self._insert_label(function_id, name, binary_name, binary_sha256)
         return function_id
 
-    async def _search_near(self, embedding, top_n):
-        return await self.connection.fetch(
+    async def search_function(self, embedding, top_n=25):
+        results = await self.connection.fetch(
             """
             SELECT label, (2 - (embedding <=> $1)) / 2 AS similarity, binary_name, binary_sha256
             FROM labels
@@ -163,9 +159,6 @@ async def _search_near(self, embedding, top_n):
             top_n,
         )
 
-    async def search_function(self, embedding, top_n=25):
-        results = await self._search_near(embedding, top_n)
-
         return [
             dict(
                 zip(

citatio/pdm.lock

@@ -27,7 +27,6 @@ dependencies = [
     "pydantic>=2.11.5",
     "scipy",
     "sqlite-vec>=0.1.6",
-    "testcontainers[postgres]>=4.14.2",
     "torch",
     "transformers",
 ]
@@ -44,6 +43,7 @@ test = [
     "httpx2",  # avoid deprecation warnings for fastapi / starlette tests
     "pytest>=8.3.5",
     "pytest-asyncio>=1.4.0",
+    "testcontainers[postgres]>=4.14.2",
 ]
 
 [tool.pdm.version]
@@ -74,7 +74,7 @@ test = "coverage run --branch --source citatio --module pytest --strict-markers
 asyncio_default_test_loop_scope = "session"
 asyncio_default_fixture_loop_scope = "session"
 asyncio_mode = "auto"
-markers = ["postgresql"]
+markers = ["postgresql", "sqlite"]
 
 [tool.ruff]
 format.quote-style = "single"