Add logic for adding updated_at triggers, add triggers to relevant columns

Author: maxachis

alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py

@@ -0,0 +1,46 @@
+"""Add updated_at triggers
+
+Revision ID: ff4e8b2f6348
+Revises: a8f36f185694
+Create Date: 2025-10-14 18:37:07.121323
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+from src.util.alembic_helpers import create_updated_at_trigger
+
+# revision identifiers, used by Alembic.
+revision: str = 'ff4e8b2f6348'
+down_revision: Union[str, None] = 'a8f36f185694'
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    for table in [
+        "agencies",
+        "auto_record_type_suggestions",
+        "auto_relevant_suggestions",
+        "flag_url_validated",
+        "link_batch_urls",
+        "link_urls_agency",
+        "link_urls_redirect_url",
+        "link_urls_root_url",
+        "tasks",
+        "url_compressed_html",
+        "url_internet_archives_probe_metadata",
+        "url_scrape_info",
+        "url_screenshot",
+        "url_web_metadata",
+        "urls",
+        "user_record_type_suggestions",
+        "user_url_type_suggestions",
+    ]:
+        create_updated_at_trigger(table)
+
+
+def downgrade() -> None:
+    pass

src/db/client/async_.py

@@ -169,14 +169,10 @@ async def add_all(
     async def bulk_update(
         self,
         session: AsyncSession,
-        model: Base,
-        mappings: list[dict],
+        models: list[Base],
     ):
         # Note, mapping must include primary key
-        await session.execute(
-            update(model),
-            mappings
-        )
+        await sh.bulk_update(session=session, models=models)
 
     @session_manager
     async def bulk_upsert(

src/util/alembic_helpers.py

@@ -295,4 +295,35 @@ def remove_enum_value(
             f"ALTER TYPE {_q_ident(schema)}.{_q_ident(tmp_name)} "
             f"RENAME TO {_q_ident(enum_name)}"
         )
-    )
+    )
+
+
+def create_updated_at_trigger(table_name: str) -> None:
+    """
+    Adds a trigger to the given table that automatically updates the
+    'updated_at' column to the current timestamp on UPDATE.
+
+    Parameters:
+        table_name (str): Name of the table to attach the trigger to.
+    """
+
+    # Step 1: Define the trigger function (only needs to exist once)
+    op.execute("""
+    CREATE OR REPLACE FUNCTION set_updated_at()
+    RETURNS TRIGGER AS $$
+    BEGIN
+        NEW.updated_at = NOW();
+        RETURN NEW;
+    END;
+    $$ LANGUAGE plpgsql;
+    """)
+
+    # Step 2: Create the trigger for this specific table
+    trigger_name = f"{table_name}_updated_at_trigger"
+    op.execute(f"""
+    DROP TRIGGER IF EXISTS {trigger_name} ON {table_name};
+    CREATE TRIGGER {trigger_name}
+    BEFORE UPDATE ON {table_name}
+    FOR EACH ROW
+    EXECUTE FUNCTION set_updated_at();
+    """)

tests/automated/integration/db/structure/test_updated_at.py

@@ -0,0 +1,38 @@
+import asyncio
+from datetime import datetime
+
+import pytest
+
+from src.collectors.enums import URLStatus
+from src.db.models.impl.url.core.pydantic.upsert import URLUpsertModel
+from src.db.models.impl.url.core.sqlalchemy import URL
+from tests.helpers.data_creator.core import DBDataCreator
+
+
+@pytest.mark.asyncio
+async def test_updated_at(db_data_creator: DBDataCreator):
+
+    _ = await db_data_creator.create_urls(
+        count=1,
+        status=URLStatus.OK
+    )
+
+    urls: list[URL] = await db_data_creator.adb_client.get_all(URL)
+    url = urls[0]
+    assert url.updated_at is not None
+    updated_at: datetime = url.updated_at
+
+    url_upsert = URLUpsertModel(
+        id=url.id,
+        name="New Name"
+    )
+
+    await db_data_creator.adb_client.bulk_update([url_upsert])
+
+    new_urls: list[URL] = await db_data_creator.adb_client.get_all(URL)
+    new_url = new_urls[0]
+
+    new_updated_at = new_url.updated_at
+    assert new_updated_at > updated_at
+
+