Merge pull request #355 from Police-Data-Accessibility-Project/mc_89_add_scraping_logic_for_non_pending_urls

mc_89_add_scraping_logic_for_non_pending_urls

Author: maxachis

alembic/env.py

@@ -1,4 +1,3 @@
-import logging
 from datetime import datetime
 from logging.config import fileConfig
 
@@ -7,7 +6,7 @@
 from sqlalchemy import pool
 
 from src.db.helpers.connect import get_postgres_connection_string
-from src.db.models.templates import Base
+from src.db.models.templates_.base import Base
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.

alembic/versions/2025_07_31_1536-99eceed6e614_add_web_status_info_table.py

@@ -0,0 +1,156 @@
+"""Add HTML Status Info table
+
+Revision ID: 99eceed6e614
+Revises: 637de6eaa3ab
+Create Date: 2025-07-31 15:36:40.966605
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from src.util.alembic_helpers import id_column, created_at_column, updated_at_column, url_id_column, switch_enum_type
+
+# revision identifiers, used by Alembic.
+revision: str = '99eceed6e614'
+down_revision: Union[str, None] = '637de6eaa3ab'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+WEB_STATUS_ENUM = sa.Enum(
+    "not_attempted",
+    "success",
+    "error",
+    "404_not_found",
+    name="web_status"
+)
+SCRAPE_STATUS_ENUM = sa.Enum(
+    "success",
+    "error",
+    name="scrape_status",
+)
+
+URL_WEB_METADATA_TABLE_NAME = 'url_web_metadata'
+URL_SCRAPE_INFO = 'url_scrape_info'
+
+
+
+
+
+def upgrade() -> None:
+    _create_url_html_info_table()
+    _add_url_probe_task_type_enum()
+    _set_up_scrape_info_table()
+    _use_existing_html_data_to_add_scrape_info()
+
+def _use_existing_html_data_to_add_scrape_info():
+    op.execute(
+        f"""
+        INSERT INTO {URL_SCRAPE_INFO} (url_id, status)
+        SELECT url_id, 'success'::scrape_status
+        FROM url_compressed_html
+        """
+    )
+    op.execute(
+        f"""
+        INSERT INTO {URL_SCRAPE_INFO} (url_id, status)
+        SELECT distinct(url_id), 'success'::scrape_status
+        FROM url_html_content
+        LEFT JOIN URL_COMPRESSED_HTML USING (url_id)
+        WHERE URL_COMPRESSED_HTML.url_id IS NULL
+        """
+    )
+
+def downgrade() -> None:
+    _drop_scrape_info_table()
+    # Drop Enums
+    WEB_STATUS_ENUM.drop(op.get_bind(), checkfirst=True)
+    _drop_url_probe_task_type_enum()
+    _tear_down_scrape_info_table()
+
+
+def _set_up_scrape_info_table():
+    op.create_table(
+        URL_SCRAPE_INFO,
+        id_column(),
+        url_id_column(),
+        sa.Column(
+            'status',
+            SCRAPE_STATUS_ENUM,
+            nullable=False,
+            comment='The status of the most recent scrape attempt.'
+        ),
+        created_at_column(),
+        updated_at_column(),
+        sa.UniqueConstraint('url_id', name='uq_url_scrape_info_url_id')
+    )
+
+
+
+
+def _tear_down_scrape_info_table():
+    op.drop_table(URL_SCRAPE_INFO)
+    # Drop enum
+    SCRAPE_STATUS_ENUM.drop(op.get_bind(), checkfirst=True)
+
+
+def _add_url_probe_task_type_enum() -> None:
+    switch_enum_type(
+        table_name='tasks',
+        column_name='task_type',
+        enum_name='task_type',
+        new_enum_values=[
+            'HTML',
+            'Relevancy',
+            'Record Type',
+            'Agency Identification',
+            'Misc Metadata',
+            'Submit Approved URLs',
+            'Duplicate Detection',
+            '404 Probe',
+            'Sync Agencies',
+            'Sync Data Sources',
+            'Push to Hugging Face',
+            'URL Probe'
+        ]
+    )
+
+def _drop_url_probe_task_type_enum() -> None:
+    switch_enum_type(
+        table_name='tasks',
+        column_name='task_type',
+        enum_name='task_type',
+        new_enum_values=[
+            'HTML',
+            'Relevancy',
+            'Record Type',
+            'Agency Identification',
+            'Misc Metadata',
+            'Submit Approved URLs',
+            'Duplicate Detection',
+            '404 Probe',
+            'Sync Agencies',
+            'Sync Data Sources',
+            'Push to Hugging Face'
+        ]
+    )
+
+def _create_url_html_info_table() -> None:
+    op.create_table(
+        URL_WEB_METADATA_TABLE_NAME,
+        id_column(),
+        url_id_column(),
+        sa.Column('accessed', sa.Boolean(), nullable=False),
+        sa.Column('status_code', sa.Integer(), nullable=True),
+        sa.Column('content_type', sa.Text(), nullable=True),
+        sa.Column('error_message', sa.Text(), nullable=True),
+        created_at_column(),
+        updated_at_column(),
+        sa.UniqueConstraint('url_id', name='uq_url_web_status_info_url_id'),
+        sa.CheckConstraint('status_code >= 100', name='ck_url_web_status_info_status_code_min'),
+        sa.CheckConstraint('status_code <= 999', name='ck_url_web_status_info_status_code_max'),
+    )
+
+def _drop_scrape_info_table() -> None:
+    op.drop_table(URL_WEB_METADATA_TABLE_NAME)

local_database/classes/DockerClient.py

@@ -1,5 +1,7 @@
 import docker
 from docker.errors import NotFound, APIError
+from docker.models.containers import Container
+from docker.models.networks import Network
 
 from local_database.DTOs import DockerfileInfo, DockerInfo
 
@@ -9,7 +11,7 @@ class DockerClient:
     def __init__(self):
         self.client = docker.from_env()
 
-    def run_command(self, command: str, container_id: str):
+    def run_command(self, command: str, container_id: str) -> None:
         exec_id = self.client.api.exec_create(
             container_id,
             cmd=command,
@@ -20,7 +22,7 @@ def run_command(self, command: str, container_id: str):
         for line in output_stream:
             print(line.decode().rstrip())
 
-    def start_network(self, network_name):
+    def start_network(self, network_name) -> Network:
         try:
             self.client.networks.create(network_name, driver="bridge")
         except APIError as e:
@@ -30,14 +32,14 @@ def start_network(self, network_name):
             print("Network already exists")
         return self.client.networks.get(network_name)
 
-    def stop_network(self, network_name):
+    def stop_network(self, network_name) -> None:
         self.client.networks.get(network_name).remove()
 
     def get_image(
             self,
             dockerfile_info: DockerfileInfo,
             force_rebuild: bool = False
-    ):
+    ) -> None:
         if dockerfile_info.dockerfile_directory:
             # Build image from Dockerfile
             self.client.images.build(
@@ -58,7 +60,7 @@ def get_image(
         except NotFound:
             self.client.images.pull(dockerfile_info.image_tag)
 
-    def get_existing_container(self, docker_info_name: str):
+    def get_existing_container(self, docker_info_name: str) -> Container | None:
         try:
             return self.client.containers.get(docker_info_name)
         except NotFound:

local_database/classes/DockerContainer.py

@@ -11,19 +11,19 @@ def __init__(self, dc: DockerClient, container: Container):
         self.dc = dc
         self.container = container
 
-    def run_command(self, command: str):
+    def run_command(self, command: str) -> None:
         self.dc.run_command(command, self.container.id)
 
-    def stop(self):
+    def stop(self) -> None:
         self.container.stop()
 
-    def log_to_file(self):
+    def log_to_file(self) -> None:
         logs = self.container.logs(stdout=True, stderr=True)
         container_name = self.container.name
         with open(f"{container_name}.log", "wb") as f:
             f.write(logs)
 
-    def wait_for_pg_to_be_ready(self):
+    def wait_for_pg_to_be_ready(self) -> None:
         for i in range(30):
             exit_code, output = self.container.exec_run("pg_isready")
             print(output)

local_database/classes/DockerManager.py

@@ -4,6 +4,8 @@
 
 import docker
 from docker.errors import APIError
+from docker.models.containers import Container
+from docker.models.networks import Network
 
 from local_database.DTOs import DockerfileInfo, DockerInfo
 from local_database.classes.DockerClient import DockerClient
@@ -20,7 +22,7 @@ def __init__(self):
         self.network = self.start_network()
 
     @staticmethod
-    def start_docker_engine():
+    def start_docker_engine() -> None:
         system = platform.system()
 
         match system:
@@ -41,7 +43,7 @@ def start_docker_engine():
                 sys.exit(1)
 
     @staticmethod
-    def is_docker_running():
+    def is_docker_running() -> bool:
         try:
             client = docker.from_env()
             client.ping()
@@ -50,16 +52,23 @@ def is_docker_running():
             print(f"Docker is not running: {e}")
             return False
 
-    def run_command(self, command: str, container_id: str):
+    def run_command(
+        self,
+        command: str,
+        container_id: str
+    ) -> None:
         self.client.run_command(command, container_id)
 
-    def start_network(self):
+    def start_network(self) -> Network:
         return self.client.start_network(self.network_name)
 
-    def stop_network(self):
+    def stop_network(self) -> None:
         self.client.stop_network(self.network_name)
 
-    def get_image(self, dockerfile_info: DockerfileInfo):
+    def get_image(
+        self,
+        dockerfile_info: DockerfileInfo
+    ) -> None:
         self.client.get_image(dockerfile_info)
 
     def run_container(
@@ -74,5 +83,5 @@ def run_container(
         )
         return DockerContainer(self.client, raw_container)
 
-    def get_containers(self):
+    def get_containers(self) -> list[Container]:
         return self.client.client.containers.list()

local_database/classes/TimestampChecker.py

@@ -1,32 +1,31 @@
-import datetime
 import os
-from typing import Optional
+from datetime import datetime, timedelta
 
 
 class TimestampChecker:
     def __init__(self):
-        self.last_run_time: Optional[datetime.datetime] = self.load_last_run_time()
+        self.last_run_time: datetime | None = self.load_last_run_time()
 
-    def load_last_run_time(self) -> Optional[datetime.datetime]:
+    def load_last_run_time(self) -> datetime | None:
         # Check if file `last_run.txt` exists
         # If it does, load the last run time
         if os.path.exists("local_state/last_run.txt"):
             with open("local_state/last_run.txt", "r") as f:
-                return datetime.datetime.strptime(
+                return datetime.strptime(
                     f.read(),
                     "%Y-%m-%d %H:%M:%S"
                 )
         return None
 
-    def last_run_within_24_hours(self):
+    def last_run_within_24_hours(self) -> bool:
         if self.last_run_time is None:
             return False
-        return datetime.datetime.now() - self.last_run_time < datetime.timedelta(days=1)
+        return datetime.now() - self.last_run_time < timedelta(days=1)
 
     def set_last_run_time(self):
         # If directory `local_state` doesn't exist, create it
         if not os.path.exists("local_state"):
             os.makedirs("local_state")
 
         with open("local_state/last_run.txt", "w") as f:
-            f.write(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
+            f.write(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))

local_database/create_database.py

@@ -15,7 +15,7 @@
 
 
 # Connect to the default 'postgres' database to create other databases
-def connect(database="postgres", autocommit=True):
+def connect(database="postgres", autocommit=True) -> psycopg2.extensions.connection:
     conn = psycopg2.connect(
         dbname=database,
         user=POSTGRES_USER,
@@ -27,7 +27,7 @@ def connect(database="postgres", autocommit=True):
         conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
     return conn
 
-def create_database(db_name):
+def create_database(db_name: str) -> None:
     conn = connect("postgres")
     with conn.cursor() as cur:
         cur.execute(sql.SQL("""
@@ -48,7 +48,7 @@ def create_database(db_name):
         except Exception as e:
             print(f"❌ Failed to create {db_name}: {e}")
 
-def main():
+def main() -> None:
     print("Creating databases...")
     create_database(LOCAL_SOURCE_COLLECTOR_DB_NAME)