Enable foreign key constraints

README.md

@@ -30,6 +30,8 @@ pip install nextline-rdb
 | -------------------- | --------------------- | --------------------------------------------------------------------------------------------- |
 | `NEXTLINE_DB__URL`   | `sqlite+aiosqlite://` | The [DB URL](https://docs.sqlalchemy.org/en/20/core/engines.html#database-urls) of SQLAlchemy |
 
+Only tested on SQLite + aiosqlite.
+
 ## License
 
 - _Nextline RDB_ is licensed under the [MIT](https://spdx.org/licenses/MIT.html) license.

src/nextline_rdb/alembic/models/rev_f9a742bb2297/model_run.py

@@ -27,6 +27,12 @@ class Run(Model):
     )
     script: Mapped[Optional['Script']] = relationship(back_populates='runs')
 
-    traces: Mapped[list["Trace"]] = relationship(back_populates="run")
-    prompts: Mapped[list["Prompt"]] = relationship(back_populates="run")
-    stdouts: Mapped[list["Stdout"]] = relationship(back_populates="run")
+    traces: Mapped[list["Trace"]] = relationship(
+        back_populates='run', cascade='all, delete-orphan'
+    )
+    prompts: Mapped[list["Prompt"]] = relationship(
+        back_populates='run', cascade='all, delete-orphan'
+    )
+    stdouts: Mapped[list["Stdout"]] = relationship(
+        back_populates='run', cascade='all, delete-orphan'
+    )

src/nextline_rdb/alembic/models/rev_f9a742bb2297/tests/test_model_run.py

@@ -1,7 +1,6 @@
 import datetime
 
 from hypothesis import given
-from hypothesis import strategies as st
 from sqlalchemy import select
 
 from nextline_rdb.db import DB
@@ -10,16 +9,32 @@
 from ..strategies import st_model_run
 
 
-@given(st.data())
-async def test_repr(data: st.DataObject):
+@given(run=st_model_run(generate_traces=False))
+async def test_repr(run: Run):
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
-            model = data.draw(st_model_run(generate_traces=False))
-            session.add(model)
+            session.add(run)
 
         async with db.session() as session:
             rows = await session.scalars(select(Run))
             for row in rows:
                 repr_ = repr(row)
                 assert Run, datetime  # type: ignore[truthy-function]
                 assert repr_ == repr(eval(repr_))
+
+
+@given(run=st_model_run(generate_traces=True))
+async def test_cascade(run: Run):
+    async with DB(use_migration=False, model_base_class=Model) as db:
+        async with db.session.begin() as session:
+            session.add(run)
+
+        async with db.session.begin() as session:
+            select_run = select(Run)
+            run = (await session.execute(select_run)).scalar_one()
+            await session.delete(run)
+
+        async with db.session() as session:
+            select_traces = select(Run)
+            traces = (await session.execute(select_traces)).scalars().all()
+            assert not traces

src/nextline_rdb/alembic/versions/2024-01-29_14:12-2fa41efc56f0.py

@@ -17,13 +17,17 @@
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    op.execute('PRAGMA foreign_keys=OFF;')
     with op.batch_alter_table('run', schema=None) as batch_op:
         batch_op.alter_column('script', new_column_name='script_old')
+    op.execute('PRAGMA foreign_keys=ON;')
     # ### end Alembic commands ###
 
 
 def downgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    op.execute('PRAGMA foreign_keys=OFF;')
     with op.batch_alter_table('run', schema=None) as batch_op:
         batch_op.alter_column('script_old', new_column_name='script')
+    op.execute('PRAGMA foreign_keys=ON;')
     # ### end Alembic commands ###

src/nextline_rdb/alembic/versions/2024-01-29_15:11-6e3cf7d9b6bf.py

@@ -17,6 +17,7 @@
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    op.execute('PRAGMA foreign_keys=OFF;')
     op.create_table('script',
     sa.Column('id', sa.Integer(), nullable=False),
     sa.Column('current', sa.Boolean(), nullable=False),
@@ -29,12 +30,13 @@ def upgrade():
     with op.batch_alter_table('run', schema=None) as batch_op:
         batch_op.add_column(sa.Column('script_id', sa.Integer(), nullable=True))
         batch_op.create_foreign_key(batch_op.f('fk_run_script_id_script'), 'script', ['script_id'], ['id'])
-
+    op.execute('PRAGMA foreign_keys=ON;')
     # ### end Alembic commands ###
 
 
 def downgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    op.execute('PRAGMA foreign_keys=OFF;')
     with op.batch_alter_table('run', schema=None) as batch_op:
         batch_op.drop_constraint(batch_op.f('fk_run_script_id_script'), type_='foreignkey')
         batch_op.drop_column('script_id')
@@ -43,4 +45,5 @@ def downgrade():
         batch_op.drop_index(batch_op.f('ix_script_id'))
 
     op.drop_table('script')
+    op.execute('PRAGMA foreign_keys=ON;')
     # ### end Alembic commands ###

src/nextline_rdb/alembic/versions/2024-01-30_13:48-f9a742bb2297.py

@@ -17,15 +17,19 @@
 
 def upgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    op.execute('PRAGMA foreign_keys=OFF;')
     with op.batch_alter_table('run', schema=None) as batch_op:
         batch_op.drop_column('script_old')
+    op.execute('PRAGMA foreign_keys=ON;')
 
     # ### end Alembic commands ###
 
 
 def downgrade():
     # ### commands auto generated by Alembic - please adjust! ###
+    op.execute('PRAGMA foreign_keys=OFF;')
     with op.batch_alter_table('run', schema=None) as batch_op:
         batch_op.add_column(sa.Column('script_old', sa.TEXT(), nullable=True))
+    op.execute('PRAGMA foreign_keys=ON;')
 
     # ### end Alembic commands ###

src/nextline_rdb/db.py

@@ -7,7 +7,7 @@
 from alembic.migration import MigrationContext
 from alembic.runtime.environment import EnvironmentContext
 from alembic.script import ScriptDirectory
-from sqlalchemy import Connection, MetaData
+from sqlalchemy import Connection, Engine, MetaData, event
 from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
 from sqlalchemy.orm import DeclarativeBase
 
@@ -148,3 +148,15 @@ def do_run_migrations(connection: Connection) -> None:
 
     async with engine.connect() as connection:
         await connection.run_sync(do_run_migrations)
+
+
+@event.listens_for(Engine, 'connect')
+def set_sqlite_pragma(dbapi_connection, connection_record):
+    '''Enable foreign key constraints in SQLite.
+
+    The code copied from the SQLAlchemy documentation:
+    https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#foreign-key-support
+    '''
+    cursor = dbapi_connection.cursor()
+    cursor.execute("PRAGMA foreign_keys=ON")
+    cursor.close()

src/nextline_rdb/models/model_run.py

@@ -27,6 +27,12 @@ class Run(Model):
     )
     script: Mapped[Optional['Script']] = relationship(back_populates='runs')
 
-    traces: Mapped[list["Trace"]] = relationship(back_populates="run")
-    prompts: Mapped[list["Prompt"]] = relationship(back_populates="run")
-    stdouts: Mapped[list["Stdout"]] = relationship(back_populates="run")
+    traces: Mapped[list["Trace"]] = relationship(
+        back_populates='run', cascade='all, delete-orphan'
+    )
+    prompts: Mapped[list["Prompt"]] = relationship(
+        back_populates='run', cascade='all, delete-orphan'
+    )
+    stdouts: Mapped[list["Stdout"]] = relationship(
+        back_populates='run', cascade='all, delete-orphan'
+    )

src/nextline_rdb/models/tests/test_model_run.py

@@ -1,7 +1,6 @@
 import datetime
 
 from hypothesis import given
-from hypothesis import strategies as st
 from sqlalchemy import select
 
 from nextline_rdb.db import DB
@@ -10,16 +9,32 @@
 from ..strategies import st_model_run
 
 
-@given(st.data())
-async def test_repr(data: st.DataObject):
+@given(run=st_model_run(generate_traces=False))
+async def test_repr(run: Run):
     async with DB(use_migration=False, model_base_class=Model) as db:
         async with db.session.begin() as session:
-            model = data.draw(st_model_run(generate_traces=False))
-            session.add(model)
+            session.add(run)
 
         async with db.session() as session:
             rows = await session.scalars(select(Run))
             for row in rows:
                 repr_ = repr(row)
                 assert Run, datetime  # type: ignore[truthy-function]
                 assert repr_ == repr(eval(repr_))
+
+
+@given(run=st_model_run(generate_traces=True))
+async def test_cascade(run: Run):
+    async with DB(use_migration=False, model_base_class=Model) as db:
+        async with db.session.begin() as session:
+            session.add(run)
+
+        async with db.session.begin() as session:
+            select_run = select(Run)
+            run = (await session.execute(select_run)).scalar_one()
+            await session.delete(run)
+
+        async with db.session() as session:
+            select_traces = select(Run)
+            traces = (await session.execute(select_traces)).scalars().all()
+            assert not traces

tests/db/test_foreign_keys.py

@@ -0,0 +1,29 @@
+import pytest
+from sqlalchemy.exc import IntegrityError
+
+from nextline_rdb.db import DB
+
+from .models import Bar, Model
+
+
+async def test_foreign_key_constraints() -> None:
+    '''Assert that foreign key constraints are enabled.
+
+    The foreign key constraints are not enabled by default in SQLite.
+
+    - https://www.sqlite.org/foreignkeys.html
+    - https://docs.sqlalchemy.org/en/20/dialects/sqlite.html#foreign-key-support
+
+    This test tries to insert a row with a foreign key that does not exist and confirms
+    that an IntegrityError is raised.
+
+    This test is copied from c5backup:
+    https://github.com/simonsobs/c5backup/blob/v0.3.0/tests/test_db.py#L48
+    '''
+    async with DB(model_base_class=Model, use_migration=False) as db:
+        async with db.session() as session:
+            bar = Bar(foo_id=1)
+            session.add(bar)
+
+            with pytest.raises(IntegrityError, match=r'(?i)foreign.*key.*constraint'):
+                await session.commit()