[fix] Resolve issues with false tag reassignment (#3344)

Author: mihran113

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Changelog
 
+## Unreleased:
+
+### Fixes: 
+- Fix issues with tag false reassignment (mihran113)
+
 ## 3.29.1 May 8, 2025:
 
 ### Enhancements:

aim/sdk/data_version.py

@@ -1 +1 @@
-DATA_VERSION = (1, 3)
+DATA_VERSION = (1, 4)

aim/storage/migrations/versions/661514b12ee1_.py

@@ -0,0 +1,69 @@
+"""empty message
+
+Revision ID: 661514b12ee1
+Revises: 46b89d830ad8
+Create Date: 2025-06-05 19:52:31.221392
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from alembic.context import get_context
+
+
+# revision identifiers, used by Alembic.
+revision = '661514b12ee1'
+down_revision = '46b89d830ad8'
+branch_labels = None
+depends_on = None
+
+
+
+def upgrade():
+    # Get the SQLite connection context
+    context = get_context()
+    naming_convention = {
+        "fk":
+            "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+    }
+    # Use batch operations for SQLite
+    with op.batch_alter_table('run_tag', naming_convention=naming_convention) as batch_op:
+        # First drop the existing foreign key
+        batch_op.drop_constraint('fk_run_tag_run_id_run', type_='foreignkey')
+        batch_op.drop_constraint('fk_run_tag_tag_id_tag', type_='foreignkey')
+
+        # Then create a new one with CASCADE
+        batch_op.create_foreign_key('fk_run_tag_run_id_run', 'run', ['run_id'], ['id'], ondelete='CASCADE')
+        batch_op.create_foreign_key('fk_run_tag_tag_id_tag', 'tag', ['tag_id'], ['id'], ondelete='CASCADE')
+
+
+    with op.batch_alter_table('note', naming_convention=naming_convention) as batch_op:
+        # First drop the existing foreign key
+        batch_op.drop_constraint('fk_note_run_id_run', type_='foreignkey')
+
+        # Then create a new one with CASCADE
+        batch_op.create_foreign_key('fk_note_run_id_run', 'run', ['run_id'], ['id'], ondelete='CASCADE')
+
+
+def downgrade():
+    # Use batch operations for SQLite
+    naming_convention = {
+        "fk":
+            "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
+    }
+    # Use batch operations for SQLite
+    with op.batch_alter_table('run_tag', naming_convention=naming_convention) as batch_op:
+        # Drop the CASCADE foreign key
+        batch_op.drop_constraint('fk_run_tag_run_id_run', type_='foreignkey')
+        batch_op.drop_constraint('fk_run_tag_tag_id_tag', type_='foreignkey')
+
+        # Then create a new one with CASCADE
+        batch_op.create_foreign_key('fk_run_tag_run_id_run', 'run', ['run_id'], ['id'],)
+        batch_op.create_foreign_key('fk_run_tag_tag_id_tag', 'tag', ['tag_id'], ['id'],)
+
+    with op.batch_alter_table('note', naming_convention=naming_convention) as batch_op:
+        # First drop the existing foreign key
+        batch_op.drop_constraint('fk_note_run_id_run', type_='foreignkey')
+
+        # Then create a new one with CASCADE
+        batch_op.create_foreign_key('fk_note_run_id_run', 'run', ['run_id'], ['id'],)
+

aim/storage/structured/db.py

@@ -9,7 +9,7 @@
 )
 from aim.storage.types import SafeNone
 from aim.web.configs import AIM_LOG_LEVEL_KEY
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, event
 from sqlalchemy.orm import scoped_session, sessionmaker
 
 
@@ -66,6 +66,7 @@ def __init__(self, path: str, readonly: bool = False):
             pool_size=10,
             max_overflow=20,
         )
+        event.listen(self.engine, 'connect', lambda c, _: c.execute('pragma foreign_keys=on'))
         self.session_cls = scoped_session(sessionmaker(autoflush=False, bind=self.engine))
         self._upgraded = None
 

aim/storage/structured/sql_engine/entities.py

@@ -87,11 +87,9 @@ def from_hash(cls, runhash: str, created_at, session) -> 'ModelMappedRun':
 
     @classmethod
     def delete_run(cls, runhash: str, session) -> bool:
-        try:
-            rows_affected = session.query(RunModel).filter(RunModel.hash == runhash).delete()
-            session_commit_or_flush(session)
-        except Exception:
-            return False
+        rows_affected = session.query(RunModel).filter(RunModel.hash == runhash).delete()
+        session_commit_or_flush(session)
+
         return rows_affected > 0
 
     @classmethod

aim/storage/structured/sql_engine/models.py

@@ -28,7 +28,7 @@ def default_to_run_hash(context):
 run_tags = Table(
     'run_tag',
     Base.metadata,
-    Column('run_id', Integer, ForeignKey('run.id'), primary_key=True, nullable=False),
+    Column('run_id', Integer, ForeignKey('run.id', ondelete='CASCADE'), primary_key=True, nullable=False),
     Column('tag_id', Integer, ForeignKey('tag.id'), primary_key=True, nullable=False),
 )
 
@@ -51,7 +51,9 @@ class Run(Base):
     experiment_id = Column(ForeignKey('experiment.id'), nullable=True)
 
     experiment = relationship('Experiment', backref=backref('runs', uselist=True, order_by='Run.created_at.desc()'))
-    tags = relationship('Tag', secondary=run_tags, backref=backref('runs', uselist=True))
+    tags = relationship(
+        'Tag', secondary=run_tags, backref=backref('runs', uselist=True), cascade='all, delete', passive_deletes=True
+    )
     notes = relationship('Note', back_populates='run')
 
     def __init__(self, run_hash, created_at=None):
@@ -106,7 +108,7 @@ class Note(Base):
 
     id = Column(Integer, autoincrement=True, primary_key=True)
     content = Column(Text, nullable=False, default='')
-    run_id = Column(Integer, ForeignKey('run.id'))
+    run_id = Column(Integer, ForeignKey('run.id', ondelete='CASCADE'),)
     experiment_id = Column(Integer, ForeignKey('experiment.id'))
 
     created_at = Column(DateTime, default=datetime.datetime.utcnow)