feat: save valid datasets to database

Author: tillywoodfield

README.md

@@ -7,6 +7,7 @@ A Python application to validate and store published OC4IDS datasets.
 ### Prerequisites
 
 - Python 3.12
+- Postgres
 
 ### Install Python requirements
 
@@ -16,6 +17,12 @@ source .venv/bin/activate
 pip install -r requirements_dev.txt
 ```
 
+### Set database enrivonment variable
+
+```
+export DATABASE_URL="postgresql://oc4ids_datastore@localhost/oc4ids_datastore"
+```
+
 ### Run app
 
 ```

oc4ids_datastore_pipeline/database.py

@@ -0,0 +1,48 @@
+import datetime
+import logging
+import os
+
+from sqlalchemy import (
+    DateTime,
+    Engine,
+    String,
+    create_engine,
+)
+from sqlalchemy.orm import DeclarativeBase, Mapped, Session, mapped_column
+
+logger = logging.getLogger(__name__)
+
+_engine = None
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class Dataset(Base):
+    __tablename__ = "dataset"
+
+    dataset_id: Mapped[str] = mapped_column(String, primary_key=True)
+    source_url: Mapped[str] = mapped_column(String)
+    publisher_name: Mapped[str] = mapped_column(String)
+    json_url: Mapped[str] = mapped_column(String)
+    updated_at: Mapped[datetime.datetime] = mapped_column(DateTime(timezone=True))
+
+
+def get_engine() -> Engine:
+    global _engine
+    if _engine is None:
+        _engine = create_engine(os.environ["DATABASE_URL"])
+    return _engine
+
+
+def create_tables() -> None:
+    logger.info("Creating database tables")
+    Base.metadata.create_all(get_engine())
+    logger.info("Created tables successfully")
+
+
+def save_dataset(dataset: Dataset) -> None:
+    with Session(get_engine()) as session:
+        session.merge(dataset)
+        session.commit()

oc4ids_datastore_pipeline/pipeline.py

@@ -1,3 +1,4 @@
+import datetime
 import json
 import logging
 import os
@@ -6,6 +7,8 @@
 import requests
 from libcoveoc4ids.api import oc4ids_json_output
 
+from oc4ids_datastore_pipeline.database import Dataset, create_tables, save_dataset
+
 logger = logging.getLogger(__name__)
 
 
@@ -62,12 +65,34 @@ def write_json_to_file(file_name: str, json_data: Any) -> None:
         raise Exception("Error while writing to JSON file", e)
 
 
+def save_dataset_metadata(
+    dataset_name: str, source_url: str, publisher_name: str, file_name: str
+) -> None:
+    logger.info(f"Saving metadata for dataset {dataset_name}")
+    dataset = Dataset(
+        dataset_id=dataset_name,
+        source_url=source_url,
+        publisher_name=publisher_name,
+        json_url=file_name,
+        updated_at=datetime.datetime.now(datetime.UTC),
+    )
+    save_dataset(dataset)
+
+
 def process_dataset(dataset_name: str, dataset_url: str) -> None:
     logger.info(f"Processing dataset {dataset_name}")
     try:
         json_data = download_json(dataset_url)
         validate_json(dataset_name, json_data)
-        write_json_to_file(f"data/{dataset_name}.json", json_data)
+        file_name = f"data/{dataset_name}.json"
+        write_json_to_file(file_name, json_data)
+        publisher_name = json_data.get("publisher", {}).get("name", "")
+        save_dataset_metadata(
+            dataset_name=dataset_name,
+            source_url=dataset_url,
+            publisher_name=publisher_name,
+            file_name=file_name,
+        )
         logger.info(f"Processed dataset {dataset_name}")
     except Exception as e:
         logger.warning(f"Failed to process dataset {dataset_name} with error {e}")
@@ -80,4 +105,5 @@ def process_datasets() -> None:
 
 
 def run() -> None:
+    create_tables()
     process_datasets()

pyproject.toml

@@ -9,7 +9,9 @@ version = "0.1.0"
 readme = "README.md"
 dependencies = [
   "libcoveoc4ids",
+  "psycopg2",
   "requests",
+  "sqlalchemy",
 ]
 
 [project.optional-dependencies]

requirements_dev.txt

@@ -103,6 +103,8 @@ platformdirs==4.3.6
     #   requests-cache
 pluggy==1.5.0
     # via pytest
+psycopg2==2.9.10
+    # via oc4ids-datastore-pipeline (pyproject.toml)
 pycodestyle==2.12.1
     # via flake8
 pycparser==2.22
@@ -146,6 +148,8 @@ six==1.17.0
     # via
     #   rfc3339-validator
     #   url-normalize
+sqlalchemy==2.0.37
+    # via oc4ids-datastore-pipeline (pyproject.toml)
 transaction==5.0
     # via zodb
 types-requests==2.32.0.20241016
@@ -154,6 +158,7 @@ typing-extensions==4.12.2
     # via
     #   mypy
     #   referencing
+    #   sqlalchemy
 url-normalize==1.4.3
     # via requests-cache
 urllib3==2.3.0