Merge pull request #85 from NetherlandsForensicInstitute/feature/add-postgresql-database-interface

Add PostgreSQL database interface

Author: akaIDIOT

.github/workflows/citatio.yml

@@ -27,20 +27,34 @@ jobs:
   test:
     runs-on: ubuntu-latest
     needs: check
-    strategy:
-      matrix:
-        python-version: ['3.13']
     defaults:
       run:
         working-directory: citatio/
+    services:
+      pgvector:
+        image: pgvector/pgvector:pg18
+        ports:
+          - 5432:5432
+        env:
+          POSTGRES_PASSWORD: tee4ieCh1jeerahkie3u
+        options: >-
+          --health-cmd pg_isready
+          --health-interval 10s
+          --health-timeout 5s
+          --health-retries 5
     steps:
       - uses: actions/checkout@v6
         with:
           fetch-depth: 20
       - uses: pdm-project/setup-pdm@v4
         with:
-          python-version: ${{ matrix.python-version }}
+          python-version: '3.13'
           cache: true
           cache-dependency-path: 'citatio/pdm.lock'
       - run: pdm install --group test
       - run: pdm run test
+      - run: pdm run integration-test
+        env:
+          POSTGRES_HOST: localhost
+          POSTGRES_PORT: 5432
+          POSTGRES_PASSWORD: tee4ieCh1jeerahkie3u

citatio/citatio/db.py

@@ -1,8 +1,10 @@
 import sqlite3
 from importlib import resources
 
+import asyncpg
 import numpy as np
 import sqlite_vec
+from pgvector.asyncpg import register_vector
 
 
 class Database:
@@ -106,11 +108,65 @@ def search_function(self, embedding, top_n=25):
 
 
 class PostgreSQLDatabase:
-    def __init__(self):
-        raise TypeError
+    def __init__(self, connection):
+        self.connection = connection
 
-    def add_function(self, name, cfg, embedding, binary_name, binary_sha256, model_identifier=None):
-        raise NotImplementedError
+    @classmethod
+    async def connect(cls, **kwargs):
+        connection = await asyncpg.connect(**kwargs)
+        await connection.execute('CREATE EXTENSION IF NOT EXISTS vector')
+        await register_vector(connection)
+        await connection.execute(resources.read_text('citatio', 'schema-postgresql.sql'))
+        return cls(connection)
 
-    def search_function(self, embedding, top_n=25):
-        raise NotImplementedError
+    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,
+                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,
+        )
+
+    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(
+            """
+            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
+            LIMIT $2
+            """,
+            embedding,
+            top_n,
+        )
+
+    async def search_function(self, embedding, top_n=25):
+        results = await self._search_near(embedding, top_n)
+
+        return [
+            dict(
+                zip(
+                    ['function', 'similarity', 'binary_name', 'binary_sha256'],
+                    result,
+                    strict=True,
+                )
+            )
+            for result in results
+        ]

citatio/citatio/schema-postgresql.sql

@@ -0,0 +1,19 @@
+CREATE TABLE IF NOT EXISTS functions (
+    id SERIAL PRIMARY KEY,
+    -- UNIQUE constraint allows using an error trigger as a 'cfg is already known' signal
+    cfg TEXT NOT NULL UNIQUE,
+    embedding VECTOR(768)
+);
+
+CREATE INDEX IF NOT EXISTS embeddings_cosine ON functions USING hnsw (embedding vector_cosine_ops);
+
+CREATE TABLE IF NOT EXISTS labels (
+    id SERIAL PRIMARY KEY,
+    function_id INTEGER NOT NULL,
+    label TEXT NOT NULL,
+    -- (binary_name, binary_sha256) to identify the source of this label
+    binary_name TEXT NOT NULL,
+    binary_sha256 TEXT NOT NULL,
+
+    FOREIGN KEY (function_id) REFERENCES functions(id)
+);