Merge pull request #87 from NetherlandsForensicInstitute/feature/api-use-postgresql-database

Let FastAPI app use postgresql database

Author: akaIDIOT

citatio/citatio/api.py

@@ -1,12 +1,12 @@
-import os
 from contextlib import asynccontextmanager
 from typing import Annotated
 
+import confidence
 from asmtransformers.models.embedder import ASMEmbedder
 from fastapi import FastAPI, Request
 from fastapi.params import Body
 
-from citatio.db import SQLiteDatabase
+from citatio.db import PostgreSQLDatabase
 from citatio.models import ControlFlowGraph
 
 
@@ -15,26 +15,22 @@
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    model = os.environ.get('CITATIO_MODEL', DEFAULT_MODEL)
-    database = os.environ.get('CITATIO_SQLITE_DATABASE', ':memory:')
+    config = confidence.load_name('citatio')
+    if not config.database:
+        raise ValueError(f'missing database configuration, available: {config}')
 
-    if database == ':memory:':
-        # TODO: log warning that we're running in in-memory database
-        pass
+    app.state.model = ASMEmbedder.from_pretrained(config.model or DEFAULT_MODEL)
 
-    app.state.model = ASMEmbedder.from_pretrained(model)
-
-    with SQLiteDatabase.from_name(database) as database:
+    async with await PostgreSQLDatabase.connect(**config.database) as database:
         app.state.database = database
-
         yield
 
 
 app = FastAPI(lifespan=lifespan)
 
 
 @app.post('/api/v1/add')
-async def add_function(  # NB: function body isn't actually async, forcing it to run blocking on the event loop
+async def add_function(
     request: Request,
     name: Annotated[str, Body()],
     cfg: Annotated[ControlFlowGraph, Body()],
@@ -43,15 +39,15 @@ async def add_function(  # NB: function body isn't actually async, forcing it to
     binary_sha256: Annotated[str, Body()] = None,
 ):
     embedding = request.app.state.model.encode(str(cfg), architecture=architecture)
-    request.app.state.database.add_function(name, cfg, embedding, binary_name, binary_sha256)
+    await request.app.state.database.add_function(name, cfg, embedding, binary_name, binary_sha256)
 
 
 @app.post('/api/v1/search')
-async def search_function(  # NB: function body isn't actually async, forcing it to run blocking on the event loop
+async def search_function(
     request: Request,
     cfg: Annotated[ControlFlowGraph, Body()],
     architecture: str = 'arm64',
     top_n: Annotated[int, Body()] = 25,
 ):
     embedding = request.app.state.model.encode(str(cfg), architecture=architecture)
-    return request.app.state.database.search_function(embedding, top_n)
+    return await request.app.state.database.search_function(embedding, top_n)

citatio/citatio/db.py

@@ -119,6 +119,12 @@ async def connect(cls, **kwargs):
         await connection.execute(resources.read_text('citatio', 'schema-postgresql.sql'))
         return cls(connection)
 
+    async def __aenter__(self):
+        return 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:
@@ -147,7 +153,7 @@ async def add_function(self, name, cfg, embedding, binary_name, binary_sha256, m
     async def _search_near(self, embedding, top_n):
         return await self.connection.fetch(
             """
-            SELECT label, (2 - (embedding <-> $1)) / 2 AS similarity, binary_name, binary_sha256
+            SELECT label, (2 - (embedding <=> $1)) / 2 AS similarity, binary_name, binary_sha256
             FROM labels
                 JOIN functions ON labels.function_id = functions.id 
             ORDER BY similarity DESC

citatio/pdm.lock

@@ -19,6 +19,7 @@ requires-python = ">=3.13,<3.14"
 dependencies = [
     "asmtransformers @ file:///${PROJECT_ROOT}/../asmtransformers",
     "asyncpg>=0.31.0",
+    "confidence>=0.17.2",
     "datasets",
     "fastapi[all]>=0.115.12",
     "numpy",

citatio/pyproject.toml

@@ -3,11 +3,20 @@
 
 import numpy as np
 import pytest
+from _pytest.monkeypatch import MonkeyPatch
 
 
 TEST_ROOT = Path(__file__).parent
 
 
+@pytest.fixture(scope='session')
+def monkeypatch_session():
+    # regular monkeypatch fixture is function scope, provide a session-scoped one
+    mp = MonkeyPatch()
+    yield mp
+    mp.undo()
+
+
 @pytest.fixture(scope='session')
 def functions():
     return json.loads((TEST_ROOT / 'functions.json').read_text(encoding='utf-8'))

citatio/tests/conftest.py

@@ -1,16 +1,82 @@
 import os
+from contextlib import contextmanager
+from os import environ
 
+import asyncpg
 import pytest
 from fastapi.testclient import TestClient
+from testcontainers.postgres import PostgresContainer
 
 from citatio.api import app
 
 
+pytestmark = pytest.mark.postgresql
+
+
+@contextmanager
+def local_pgvector_container():
+    with PostgresContainer('pgvector/pgvector:pg18', driver=None) as container:
+        yield container
+
+
+def _patch_database_env(monkeypatch, env):
+    for var, value in env.items():
+        # use environment variables to communicate configuration
+        var = f'CITATIO_DATABASE_{var}'.upper()
+        if value:
+            monkeypatch.setenv(var, str(value))
+        else:
+            # env vars have no concept of 'no value' other than removing the var
+            monkeypatch.delenv(var)
+
+    # return the original env that's still suitable for use with asyncpg.connect()
+    return env
+
+
+@pytest.fixture(scope='session')
+def connect_pgvector(monkeypatch_session):
+    # translate PostgreSQL connection details into environment variables that will be picked up by confidence.load_name
+    # in the app's lifecycle
+    match environ:
+        case {'GITHUB_ACTIONS': _}:
+            # running on GHA, provide connection details to pgvector service defined in .github/workflows/citatio.yml
+            yield _patch_database_env(
+                monkeypatch_session,
+                {
+                    'host': environ.get('POSTGRES_HOST', 'postgres'),
+                    'port': environ.get('POSTGRES_PORT', 5432),
+                    'user': environ.get('POSTGRES_USER', 'postgres'),
+                    'password': environ.get('POSTGRES_PASSWORD'),
+                    'database': environ.get('POSTGRES_DATABASE', 'postgres'),
+                },
+            )
+        case _:
+            # running locally, provide connection details to local pgvector container
+            with local_pgvector_container() as container:
+                yield _patch_database_env(
+                    monkeypatch_session,
+                    {
+                        'host': container.get_container_host_ip(),
+                        'port': container.get_exposed_port(5432),
+                        'user': container.username,
+                        'password': container.password,
+                        'database': container.dbname,
+                    },
+                )
+
+
 @pytest.fixture
-def client():
+async def client(connect_pgvector):
     with TestClient(app) as client:
         yield client
 
+    # after the test (post-yield), make sure to empty the tables that we might have inserted into during the test
+    connection = await asyncpg.connect(**connect_pgvector)
+    await connection.execute("""
+        DELETE FROM labels;
+        DELETE FROM functions;
+    """)
+
 
 @pytest.mark.skipif(os.environ.get('CI') == 'true', reason="don't run this test on CI")
 def test_add_function(client, functions):

citatio/tests/test_api.py

@@ -45,7 +45,8 @@ async def connect_pgvector():
 
 @pytest.fixture(scope='session')
 async def database(connect_pgvector):
-    return await PostgreSQLDatabase.connect(**connect_pgvector)
+    async with await PostgreSQLDatabase.connect(**connect_pgvector) as database:
+        yield database
 
 
 @pytest.fixture(scope='function')