diff --git a/backend/alembic/versions/78fc200abac1_add_account_id_script_tables.py b/backend/alembic/versions/78fc200abac1_add_account_id_script_tables.py new file mode 100644 index 00000000..21cf8f2b --- /dev/null +++ b/backend/alembic/versions/78fc200abac1_add_account_id_script_tables.py @@ -0,0 +1,80 @@ +"""add account_id to script_builder_sessions, script_templates, script_generations + +Revision ID: 78fc200abac1 +Revises: 7f136778f5a8 +Create Date: 2026-04-09 00:00:00.000000 +""" +from typing import Sequence, Union +from alembic import op +import sqlalchemy as sa + +revision: str = '78fc200abac1' +down_revision: Union[str, None] = '7f136778f5a8' +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + for table in ('script_builder_sessions', 'script_templates', 'script_generations'): + op.add_column(table, sa.Column('account_id', sa.UUID(), nullable=True)) + op.create_foreign_key( + f'fk_{table}_account_id', table, 'accounts', + ['account_id'], ['id'], ondelete='CASCADE', + ) + + # script_builder_sessions: user_id → users.account_id + op.execute(""" + UPDATE script_builder_sessions sbs + SET account_id = u.account_id + FROM users u + WHERE sbs.user_id = u.id + AND sbs.account_id IS NULL + """) + + # script_templates: created_by → users.account_id (nullable created_by) + op.execute(""" + UPDATE script_templates st + SET account_id = u.account_id + FROM users u + WHERE st.created_by = u.id + AND st.account_id IS NULL + """) + # Fallback for script_templates with NULL created_by: team_id → team admin user + op.execute(""" + UPDATE script_templates st + SET account_id = u.account_id + FROM users u + WHERE u.team_id = st.team_id + AND u.is_team_admin = TRUE + AND u.account_id IS NOT NULL + AND st.account_id IS NULL + """) + + # script_generations: user_id → users.account_id + op.execute(""" + UPDATE script_generations sg + SET account_id = u.account_id + FROM users u + WHERE sg.user_id = u.id + AND sg.account_id IS NULL + """) + + # VERIFY + for table in ('script_builder_sessions', 'script_templates', 'script_generations'): + result = op.get_bind().execute( + sa.text(f"SELECT COUNT(*) FROM {table} WHERE account_id IS NULL") + ) + count = result.scalar() + if count > 0: + raise RuntimeError(f"ROLLBACK: {count} NULL account_id rows in {table}.") + + for table in ('script_builder_sessions', 'script_templates', 'script_generations'): + op.alter_column(table, 'account_id', nullable=False) + op.create_index(f'ix_{table}_account_id', table, ['account_id']) + + +def downgrade() -> None: + for table in ('script_builder_sessions', 'script_templates', 'script_generations'): + op.drop_index(f'ix_{table}_account_id', table_name=table) + op.drop_constraint(f'fk_{table}_account_id', table, type_='foreignkey') + op.drop_column(table, 'account_id') diff --git a/backend/app/models/script_builder_session.py b/backend/app/models/script_builder_session.py index f7075494..723a4cfb 100644 --- a/backend/app/models/script_builder_session.py +++ b/backend/app/models/script_builder_session.py @@ -29,6 +29,12 @@ class ScriptBuilderSession(Base): nullable=False, index=True, ) + account_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("accounts.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) team_id: Mapped[Optional[uuid.UUID]] = mapped_column( UUID(as_uuid=True), ForeignKey("teams.id", ondelete="SET NULL"), diff --git a/backend/app/models/script_template.py b/backend/app/models/script_template.py index 838d2f3c..3624f031 100644 --- a/backend/app/models/script_template.py +++ b/backend/app/models/script_template.py @@ -44,6 +44,12 @@ class ScriptTemplate(Base): team_id: Mapped[Optional[uuid.UUID]] = mapped_column( UUID(as_uuid=True), ForeignKey("teams.id", ondelete="CASCADE"), nullable=True, index=True ) + account_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("accounts.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) created_by: Mapped[Optional[uuid.UUID]] = mapped_column( UUID(as_uuid=True), ForeignKey("users.id", ondelete="SET NULL"), nullable=True ) @@ -97,6 +103,12 @@ class ScriptGeneration(Base): user_id: Mapped[uuid.UUID] = mapped_column( UUID(as_uuid=True), ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True ) + account_id: Mapped[uuid.UUID] = mapped_column( + UUID(as_uuid=True), + ForeignKey("accounts.id", ondelete="CASCADE"), + nullable=False, + index=True, + ) team_id: Mapped[Optional[uuid.UUID]] = mapped_column( UUID(as_uuid=True), ForeignKey("teams.id", ondelete="SET NULL"), nullable=True, index=True ) diff --git a/backend/tests/test_phase1_migrations.py b/backend/tests/test_phase1_migrations.py index 32de973c..fe89c153 100644 --- a/backend/tests/test_phase1_migrations.py +++ b/backend/tests/test_phase1_migrations.py @@ -420,3 +420,26 @@ async def test_maintenance_schedule_account_id_matches_tree(test_db: AsyncSessio ) row = result.scalar_one() assert row.account_id == account.id + + +# ── Group 7: Legacy team_id tables ─────────────────────────────────────────── + +@pytest.mark.asyncio +async def test_script_builder_session_account_id(test_db: AsyncSession): + """script_builder_sessions.account_id must match user's account_id.""" + from app.models.script_builder_session import ScriptBuilderSession + + account, user = await _make_account_and_user(test_db, "sbs1") + sbs = ScriptBuilderSession( + user_id=user.id, + account_id=account.id, + language="powershell", + ) + test_db.add(sbs) + await test_db.commit() + + result = await test_db.execute( + select(ScriptBuilderSession).where(ScriptBuilderSession.id == sbs.id) + ) + row = result.scalar_one() + assert row.account_id == account.id