add sqlalchemy and update writer classes

Author: jb3rndt

configs/writer/postgres.json

@@ -6,4 +6,4 @@
     "db_name": "metis_db",
     "db_port": 5432,
     "db_host": "localhost"
-}
+}

configs/writer/sqlite.json

@@ -2,4 +2,4 @@
     "writer_name": "sqlite",
     "table_name": "dqresults",
     "db_name": "dq_repository/dq_repository.db"
-}
+}

metis/database_models.py

@@ -0,0 +1,27 @@
+from datetime import datetime
+from typing import List
+
+from sqlalchemy import JSON, Double, func
+from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
+
+
+class Base(DeclarativeBase):
+    pass
+
+def register_models(results_table_name: str):
+    class DQResultModel(Base):
+        __tablename__ = results_table_name
+        __table_args__ = {"extend_existing": True}
+
+        id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+        mes_time: Mapped[datetime] = mapped_column(server_default=func.now())
+        dq_value: Mapped[float] = mapped_column(Double)
+        dq_dimension: Mapped[str]
+        dq_metric: Mapped[str]
+        column_name: Mapped[List[str] | None] = mapped_column(JSON)
+        row_index: Mapped[int | None]
+        dq_annotations: Mapped[dict | None] = mapped_column(JSON)
+        dataset: Mapped[str | None]
+        table_name: Mapped[str | None]
+
+    return DQResultModel

metis/writer/database_writer.py

@@ -0,0 +1,38 @@
+from typing import Dict, List
+
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session
+
+from metis.database_models import Base, register_models
+from metis.utils.result import DQResult
+from metis.writer.writer import DQResultWriter
+
+
+class DatabaseWriter(DQResultWriter):
+    def __init__(self, writer_config: Dict) -> None:
+        self.engine = self.create_engine(writer_config)
+
+        self.DQResultModel = register_models(writer_config.get("table_name", "dq_results"))
+        Base.metadata.create_all(self.engine)
+
+    def create_engine(self, writer_config: Dict) -> Engine:
+        raise NotImplementedError("Subclasses must implement the create_engine method.")
+
+    def write(self, results: List[DQResult]) -> None:
+        with Session(self.engine) as session:
+            db_entities = [
+                self.DQResultModel(
+                    mes_time=result.mesTime.to_pydatetime(),
+                    dq_value=result.DQvalue,
+                    dq_dimension=result.DQdimension,
+                    dq_metric=result.DQmetric,
+                    column_name=result.columnNames,
+                    row_index=result.rowIndex,
+                    dq_annotations=result.DQannotations,
+                    dataset=result.dataset,
+                    table_name=result.tableName,
+                )
+                for result in results
+            ]
+            session.add_all(db_entities)
+            session.commit()

metis/writer/postgres_writer.py

@@ -1,80 +1,17 @@
-import json
-import psycopg2
-from typing import List
+from sqlalchemy import Engine, create_engine
 
-from metis.writer.writer import DQResultWriter
-from metis.utils.result import DQResult
+from metis.writer.database_writer import DatabaseWriter
 
-class PostgresWriter(DQResultWriter):
-    def __init__(self, writer_config) -> None:
+
+class PostgresWriter(DatabaseWriter):
+    def create_engine(self, writer_config) -> Engine:
         required_keys = ("db_user", "db_pass", "db_name", "db_host", "db_port")
         if not all(k in writer_config for k in required_keys):
-            raise ValueError("Postgres writer config must include 'db_user', 'db_pass', 'db_name', 'db_host', and 'db_port' fields.")
-        
-        self.table_name = writer_config.get("table_name")
-        self.DB_USER = writer_config.get("db_user")
-        self.DB_PASS = writer_config.get("db_pass")
-        self.DB_NAME = writer_config.get("db_name")
-        self.DB_HOST = writer_config.get("db_host")
-        self.DB_PORT = writer_config.get("db_port")
-
-        conn = self.connect()
-        self.create_db_schema(conn)
-        conn.close()
+            raise ValueError(
+                "Postgres writer config must include 'db_user', 'db_pass', 'db_name', 'db_host', and 'db_port' fields."
+            )
 
-    def connect(self):
-        conn = psycopg2.connect(
-            dbname=self.DB_NAME,
-            user=self.DB_USER,
-            password=self.DB_PASS,
-            host=self.DB_HOST,
-            port=self.DB_PORT
+        return create_engine(
+            f"postgresql://{writer_config['db_user']}:{writer_config['db_pass']}@{writer_config['db_host']}:{writer_config['db_port']}/{writer_config['db_name']}",
+            echo=True,
         )
-        return conn
-
-    def create_db_schema(self, conn):
-        query = f"""
-            CREATE TABLE IF NOT EXISTS {self.table_name} (
-                id SERIAL PRIMARY KEY,
-                mes_time TIMESTAMP WITH TIME ZONE NOT NULL,
-                dq_value DOUBLE PRECISION NOT NULL,
-                dq_dimension TEXT NOT NULL,
-                dq_metric TEXT NOT NULL,
-                column_name JSONB,
-                row_index INTEGER,
-                dq_annotations JSONB,
-                dataset TEXT,
-                table_name TEXT
-            );
-        """
-        try:
-            cursor = conn.cursor()
-            cursor.execute(query=query)
-            conn.commit()
-            cursor.close()
-        except Exception as e:
-            print(f'Error when saving or connecting to DB: {e}')
-
-    def write(self, results: List[DQResult]) -> None:
-        conn = self.connect()
-        cur = conn.cursor()
-        
-        for result in results:
-            print(f"Writing result: {result.as_json()}")
-            cur.execute(f'''
-                INSERT INTO {self.table_name} (mes_time, dq_value, dq_dimension, dq_metric, column_name, row_index, dq_annotations, dataset, table_name)
-                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
-            ''', (
-                result.mesTime.to_pydatetime(),
-                result.DQvalue,
-                result.DQdimension,
-                result.DQmetric,
-                json.dumps(result.columnNames),
-                result.rowIndex,
-                json.dumps(result.DQannotations),
-                result.dataset,
-                result.tableName
-            ))
-        
-        conn.commit()
-        conn.close()

metis/writer/sqlite_writer.py

@@ -1,58 +1,11 @@
-import sqlite3
-import os
-import json
-from typing import List
+from sqlalchemy import Engine, create_engine
 
-from metis.writer.writer import DQResultWriter
-from metis.utils.result import DQResult
+from metis.writer.database_writer import DatabaseWriter
 
-class SQLiteWriter(DQResultWriter):
-    def __init__(self, writer_config) -> None:
-        if not "db_name" in writer_config or not "table_name" in writer_config:
-            raise ValueError("SQLite writer config must include 'db_name' and 'table_name' fields.")
-        self.db_name = writer_config["db_name"]
-        self.table_name = writer_config["table_name"]
 
-        if not os.path.exists(self.db_name): #TODO: Make this modular for different databases
-            conn = sqlite3.connect(self.db_name)
-            cur = conn.cursor()
-            cur.execute(f'''
-            CREATE TABLE IF NOT EXISTS {self.table_name} (
-                id INTEGER PRIMARY KEY AUTOINCREMENT,
-                mes_time TEXT NOT NULL,
-                dq_value REAL NOT NULL,
-                dq_dimension TEXT NOT NULL,
-                dq_metric TEXT NOT NULL,
-                column_name JSONB,
-                row_index INTEGER,
-                dq_annotations JSONB,
-                dataset TEXT,
-                table_name TEXT
-            )
-            ''')
-            conn.commit()
-            conn.close()
+class SQLiteWriter(DatabaseWriter):
+    def create_engine(self, writer_config) -> Engine:
+        if "db_name" not in writer_config:
+            raise ValueError("SQLite writer config must include 'db_name' field.")
 
-    def write(self, results: List[DQResult]) -> None:
-        conn = sqlite3.connect(self.db_name)
-        cur = conn.cursor()
-        
-        for result in results:
-            print(f"Writing result: {result.as_json()}")
-            cur.execute(f'''
-                INSERT INTO {self.table_name} (mes_time, dq_value, dq_dimension, dq_metric, column_name, row_index, dq_annotations, dataset, table_name)
-                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
-            ''', (
-                result.mesTime.to_pydatetime(),
-                result.DQvalue,
-                result.DQdimension,
-                result.DQmetric,
-                json.dumps(result.columnNames),
-                result.rowIndex,
-                json.dumps(result.DQannotations),
-                result.dataset,
-                result.tableName
-            ))
-        
-        conn.commit()
-        conn.close()
+        return create_engine(f"sqlite:///{writer_config['db_name']}", echo=True)

requirements.txt

@@ -1,3 +1,4 @@
 pandas
 psycopg2-binary
-sqlite3 ; sys_platform == "win32"  # sqlite3 is included with Python, but this line is for completeness
+sqlite3 ; sys_platform == "win32"  # sqlite3 is included with Python, but this line is for completeness
+sqlalchemy==2.0.44