Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -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

Check warning on line 10 in alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py#L10 <401>

'alembic.op' imported but unused
Raw output
./alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py:10:1: F401 'alembic.op' imported but unused
import sqlalchemy as sa

Check warning on line 11 in alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py#L11 <401>

'sqlalchemy as sa' imported but unused
Raw output
./alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py:11:1: F401 'sqlalchemy as sa' imported but unused

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:

Check warning on line 22 in alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py#L22 <103>

Missing docstring in public function
Raw output
./alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py:22:1: D103 Missing docstring in public function
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:

Check warning on line 45 in alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py#L45 <103>

Missing docstring in public function
Raw output
./alembic/versions/2025_10_14_1837-ff4e8b2f6348_add_updated_at_triggers.py:45:1: D103 Missing docstring in public function
pass
8 changes: 2 additions & 6 deletions src/db/client/async_.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
33 changes: 32 additions & 1 deletion src/util/alembic_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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();
""")
38 changes: 38 additions & 0 deletions tests/automated/integration/db/structure/test_updated_at.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
import asyncio

Check warning on line 1 in tests/automated/integration/db/structure/test_updated_at.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] tests/automated/integration/db/structure/test_updated_at.py#L1 <100>

Missing docstring in public module
Raw output
./tests/automated/integration/db/structure/test_updated_at.py:1:1: D100 Missing docstring in public module

Check warning on line 1 in tests/automated/integration/db/structure/test_updated_at.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] tests/automated/integration/db/structure/test_updated_at.py#L1 <401>

'asyncio' imported but unused
Raw output
./tests/automated/integration/db/structure/test_updated_at.py:1:1: F401 'asyncio' imported but unused
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):

Check warning on line 13 in tests/automated/integration/db/structure/test_updated_at.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] tests/automated/integration/db/structure/test_updated_at.py#L13 <103>

Missing docstring in public function
Raw output
./tests/automated/integration/db/structure/test_updated_at.py:13:1: D103 Missing docstring in public function

_ = 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


Check warning on line 38 in tests/automated/integration/db/structure/test_updated_at.py

View workflow job for this annotation

GitHub Actions / flake8

[flake8] tests/automated/integration/db/structure/test_updated_at.py#L38 <391>

blank line at end of file
Raw output
./tests/automated/integration/db/structure/test_updated_at.py:38:1: W391 blank line at end of file