From dabc10bf444e77cbfb3efb220253426098029e4a Mon Sep 17 00:00:00 2001
From: Max Chis <maxachis@gmail.com>
Date: Tue, 14 Oct 2025 19:02:03 -0400
Subject: [PATCH] Add logic for adding `updated_at` triggers, add triggers to
 relevant columns

---
 ...37-ff4e8b2f6348_add_updated_at_triggers.py | 46 +++++++++++++++++++
 src/db/client/async_.py                       |  8 +---
 src/util/alembic_helpers.py                   | 33 ++++++++++++-
 .../db/structure/test_updated_at.py           | 38 +++++++++++++++
 4 files changed, 118 insertions(+), 7 deletions(-)
 create mode 100644 alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py
 create mode 100644 tests/automated/integration/db/structure/test_updated_at.py

diff --git a/alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py b/alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py
new file mode 100644
index 00000000..faf10f91
--- /dev/null
+++ b/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
diff --git a/src/db/client/async_.py b/src/db/client/async_.py
index 2a15267e..87fcb057 100644
--- a/src/db/client/async_.py
+++ b/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(
diff --git a/src/util/alembic_helpers.py b/src/util/alembic_helpers.py
index cb9d8d67..f711136d 100644
--- a/src/util/alembic_helpers.py
+++ b/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)}"
         )
-    )
\ No newline at end of file
+    )
+
+
+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();
+    """)
diff --git a/tests/automated/integration/db/structure/test_updated_at.py b/tests/automated/integration/db/structure/test_updated_at.py
new file mode 100644
index 00000000..281e6ee8
--- /dev/null
+++ b/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
+
+