Mc 122 source collector main

README.md

@@ -14,6 +14,7 @@ openai-playground | Scripts for accessing the openai API on PDAP's shared accoun
 source_collectors| Tools for extracting metadata from different sources, including CKAN data portals and Common Crawler
 collector_db | Database for storing data from source collectors
 collector_manager | A module which provides a unified interface for interacting with source collectors and relevant data
+core | A module which integrates other components, such as collector_manager and collector_db
 
 ## How to use
 

collector_db/DTOs/BatchInfo.py

@@ -0,0 +1,20 @@
+from datetime import datetime
+from typing import Optional
+
+from pydantic import BaseModel
+
+from core.enums import BatchStatus
+
+
+class BatchInfo(BaseModel):
+    strategy: str
+    status: BatchStatus
+    parameters: dict
+    count: int = 0
+    strategy_success_rate: Optional[float] = None
+    metadata_success_rate: Optional[float] = None
+    agency_match_rate: Optional[float] = None
+    record_type_match_rate: Optional[float] = None
+    record_category_match_rate: Optional[float] = None
+    compute_time: Optional[float] = None
+    date_generated: Optional[datetime] = None

collector_db/DTOs/DuplicateInfo.py

@@ -0,0 +1,8 @@
+from pydantic import BaseModel
+
+
+class DuplicateInfo(BaseModel):
+    source_url: str
+    original_url_id: int
+    duplicate_metadata: dict
+    original_metadata: dict

collector_db/DTOs/InsertURLsInfo.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+from collector_db.DTOs.DuplicateInfo import DuplicateInfo
+from collector_db.DTOs.URLMapping import URLMapping
+
+
+class InsertURLsInfo(BaseModel):
+    url_mappings: list[URLMapping]
+    duplicates: list[DuplicateInfo]

collector_db/DTOs/README.md

@@ -0,0 +1 @@
+This directory consists of data transfer objects (DTOs) for the Source Collector Database.

collector_db/DTOs/URLInfo.py

@@ -0,0 +1,13 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from collector_manager.enums import URLOutcome
+
+
+class URLInfo(BaseModel):
+    id: Optional[int] = None
+    batch_id: Optional[int] = None
+    url: str
+    url_metadata: Optional[dict] = None
+    outcome: URLOutcome = URLOutcome.PENDING

collector_db/DTOs/URLMapping.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+
+class URLMapping(BaseModel):
+    url: str
+    url_id: int

collector_db/DatabaseClient.py

@@ -1,60 +1,21 @@
 from functools import wraps
 
-from sqlalchemy import create_engine, Column, Integer, String, Float, Text, JSON, ForeignKey, CheckConstraint, TIMESTAMP, UniqueConstraint
-from sqlalchemy.orm import declarative_base, sessionmaker, relationship
-from typing import Optional, Dict, Any, List
+from sqlalchemy import create_engine
+from sqlalchemy.exc import IntegrityError
+from sqlalchemy.orm import sessionmaker
+from typing import Optional, List
 
-from collector_db.BatchInfo import BatchInfo
-from collector_db.URLInfo import URLInfo
+from collector_db.DTOs.BatchInfo import BatchInfo
+from collector_db.DTOs.DuplicateInfo import DuplicateInfo
+from collector_db.DTOs.InsertURLsInfo import InsertURLsInfo
+from collector_db.DTOs.URLMapping import URLMapping
+from collector_db.DTOs.URLInfo import URLInfo
+from collector_db.models import Base, Batch, URL
+from core.enums import BatchStatus
 
-# Base class for SQLAlchemy ORM models
-Base = declarative_base()
 
-# SQLAlchemy ORM models
-class Batch(Base):
-    __tablename__ = 'batches'
-
-    id = Column(Integer, primary_key=True)
-    strategy = Column(String, nullable=False)
-    status = Column(String, CheckConstraint("status IN ('in-process', 'complete', 'error')"), nullable=False)
-    count = Column(Integer, nullable=False)
-    date_generated = Column(TIMESTAMP, nullable=False, server_default="CURRENT_TIMESTAMP")
-    strategy_success_rate = Column(Float)
-    metadata_success_rate = Column(Float)
-    agency_match_rate = Column(Float)
-    record_type_match_rate = Column(Float)
-    record_category_match_rate = Column(Float)
-    compute_time = Column(Integer)
-    parameters = Column(JSON)
-
-    urls = relationship("URL", back_populates="batch")
-    missings = relationship("Missing", back_populates="batch")
-
-
-class URL(Base):
-    __tablename__ = 'urls'
-
-    id = Column(Integer, primary_key=True)
-    batch_id = Column(Integer, ForeignKey('batches.id'), nullable=False)
-    url = Column(Text, unique=True)
-    url_metadata = Column(JSON)
-    outcome = Column(String)
-    created_at = Column(TIMESTAMP, nullable=False, server_default="CURRENT_TIMESTAMP")
 
-    batch = relationship("Batch", back_populates="urls")
-
-
-class Missing(Base):
-    __tablename__ = 'missing'
-
-    id = Column(Integer, primary_key=True)
-    place_id = Column(Integer, nullable=False)
-    record_type = Column(String, nullable=False)
-    batch_id = Column(Integer, ForeignKey('batches.id'))
-    strategy_used = Column(Text, nullable=False)
-    date_searched = Column(TIMESTAMP, nullable=False, server_default="CURRENT_TIMESTAMP")
-
-    batch = relationship("Batch", back_populates="missings")
+# SQLAlchemy ORM models
 
 
 # Database Client
@@ -81,28 +42,86 @@
                 self.session.close()
                 self.session = None
 
+        return wrapper
+
     @session_manager
-    def insert_batch(self, batch_info: BatchInfo) -> Batch:
-        """Insert a new batch into the database."""
+    def insert_batch(self, batch_info: BatchInfo) -> int:
+        """Insert a new batch into the database and return its ID."""
         batch = Batch(
-            **batch_info.model_dump()
+            strategy=batch_info.strategy,
+            status=batch_info.status.value,
+            parameters=batch_info.parameters,
+            count=batch_info.count,
+            compute_time=batch_info.compute_time,
+            strategy_success_rate=batch_info.strategy_success_rate,
+            metadata_success_rate=batch_info.metadata_success_rate,
+            agency_match_rate=batch_info.agency_match_rate,
+            record_type_match_rate=batch_info.record_type_match_rate,
+            record_category_match_rate=batch_info.record_category_match_rate,
         )
         self.session.add(batch)
-        return batch
+        self.session.commit()
+        self.session.refresh(batch)
+        return batch.id
+
+    @session_manager
+    def update_batch_post_collection(
+        self,
+        batch_id: int,
+        url_count: int,
+        batch_status: BatchStatus,
+        compute_time: float = None,
+    ):
+        batch = self.session.query(Batch).filter_by(id=batch_id).first()
+        batch.count = url_count
+        batch.status = batch_status.value
+        batch.compute_time = compute_time
 
     @session_manager
     def get_batch_by_id(self, batch_id: int) -> Optional[BatchInfo]:
         """Retrieve a batch by ID."""
         batch = self.session.query(Batch).filter_by(id=batch_id).first()
         return BatchInfo(**batch.__dict__)
 
+    def insert_urls(self, url_infos: List[URLInfo], batch_id: int) -> InsertURLsInfo:
+        url_mappings = []
+        duplicates = []
+        for url_info in url_infos:
+            url_info.batch_id = batch_id
+            try:
+                url_id = self.insert_url(url_info)
+                url_mappings.append(URLMapping(url_id=url_id, url=url_info.url))
+            except IntegrityError:
+                orig_url_info = self.get_url_info_by_url(url_info.url)
+                duplicates.append(DuplicateInfo(
+                    source_url=url_info.url,
+                    original_url_id=orig_url_info.id,
+                    duplicate_metadata=url_info.url_metadata,
+                    original_metadata=orig_url_info.url_metadata
+                ))
+
+        return InsertURLsInfo(url_mappings=url_mappings, duplicates=duplicates)
+
+
     @session_manager
-    def insert_url(self, url_info: URLInfo):
+    def get_url_info_by_url(self, url: str) -> Optional[URLInfo]:
+        url = self.session.query(URL).filter_by(url=url).first()
+        return URLInfo(**url.__dict__)
+
+    @session_manager
+    def insert_url(self, url_info: URLInfo) -> int:
         """Insert a new URL into the database."""
         url_entry = URL(
-            **url_info.model_dump()
+            batch_id=url_info.batch_id,
+            url=url_info.url,
+            url_metadata=url_info.url_metadata,
+            outcome=url_info.outcome.value
         )
         self.session.add(url_entry)
+        self.session.commit()
+        self.session.refresh(url_entry)
+        return url_entry.id
+
 
     @session_manager
     def get_urls_by_batch(self, batch_id: int) -> List[URLInfo]:

collector_db/README.md

@@ -0,0 +1 @@
+The collector database is a database for storing collector data and associated metadata. It consists of both the database structure itself as well as the interfaces and helper functions for interacting with it.

collector_db/models.py

@@ -0,0 +1,75 @@
+"""
+SQLAlchemy ORM models
+"""
+from sqlalchemy import func, Column, Integer, String, CheckConstraint, TIMESTAMP, Float, JSON, ForeignKey, Text
+from sqlalchemy.orm import declarative_base, relationship
+
+from core.enums import BatchStatus
+from util.helper_functions import get_enum_values
+
+# Base class for SQLAlchemy ORM models
+Base = declarative_base()
+
+status_check_string = ", ".join([f"'{status}'" for status in get_enum_values(BatchStatus)])
+
+CURRENT_TIME_SERVER_DEFAULT = func.now()
+
+
+class Batch(Base):
+    __tablename__ = 'batches'
+
+    id = Column(Integer, primary_key=True)
+    strategy = Column(String, nullable=False)
+    # Gives the status of the batch
+    status = Column(String, CheckConstraint(f"status IN ({status_check_string})"), nullable=False)
+    # The number of URLs in the batch
+    # TODO: Add means to update after execution
+    count = Column(Integer, nullable=False)
+    date_generated = Column(TIMESTAMP, nullable=False, server_default=CURRENT_TIME_SERVER_DEFAULT)
+    # How often URLs ended up approved in the database
+    strategy_success_rate = Column(Float)
+    # Percentage of metadata identified by models
+    metadata_success_rate = Column(Float)
+    # Rate of matching to agencies
+    agency_match_rate = Column(Float)
+    # Rate of matching to record types
+    record_type_match_rate = Column(Float)
+    # Rate of matching to record categories
+    record_category_match_rate = Column(Float)
+    # Time taken to generate the batch
+    # TODO: Add means to update after execution
+    compute_time = Column(Float)
+    # The parameters used to generate the batch
+    parameters = Column(JSON)
+
+    urls = relationship("URL", back_populates="batch")
+    missings = relationship("Missing", back_populates="batch")
+
+
+class URL(Base):
+    __tablename__ = 'urls'
+
+    id = Column(Integer, primary_key=True)
+    # The batch this URL is associated with
+    batch_id = Column(Integer, ForeignKey('batches.id'), nullable=False)
+    url = Column(Text, unique=True)
+    # The metadata associated with the URL
+    url_metadata = Column(JSON)
+    # The outcome of the URL: submitted, human_labeling, rejected, duplicate, etc.
+    outcome = Column(String)
+    created_at = Column(TIMESTAMP, nullable=False, server_default=CURRENT_TIME_SERVER_DEFAULT)
+
+    batch = relationship("Batch", back_populates="urls")
+
+
+class Missing(Base):
+    __tablename__ = 'missing'
+
+    id = Column(Integer, primary_key=True)
+    place_id = Column(Integer, nullable=False)
+    record_type = Column(String, nullable=False)
+    batch_id = Column(Integer, ForeignKey('batches.id'))
+    strategy_used = Column(Text, nullable=False)
+    date_searched = Column(TIMESTAMP, nullable=False, server_default=CURRENT_TIME_SERVER_DEFAULT)
+
+    batch = relationship("Batch", back_populates="missings")