"""Add department and creator ownership to operator-managed resources.""" from alembic import context, op import sqlalchemy as sa from sqlalchemy import inspect revision = "0026_admin_ownership" down_revision = "0025_admin_rbac" branch_labels = None depends_on = None DEFAULT_DEPT_ID = "00000000-0000-0000-0000-000000000001" OWNED_TABLES = ( "MediaAsset", "HeroSlide", "DestinationHero", "VehicleOption", "VehicleServiceConfig", "HomeExperience", "HomeTeamBuilding", "HomeWildArchive", "HomeWanfaRecommendation", "WanfaCategory", "WanfaRoute", "DetailRecord", "ConciergeAdvisor", "DemandHero", "DemandFeatureCard", "DemandForm", "Lead", ) def is_offline_mode() -> bool: try: return context.is_offline_mode() except (NameError, AttributeError): return False def _schema_snapshot() -> tuple[set[str], dict[str, dict[str, dict]], dict[str, set[str]], dict[str, list[dict]]]: if is_offline_mode(): return set(), {}, {}, {} inspector = inspect(op.get_bind()) table_names = set(inspector.get_table_names()) columns = { table_name: {column["name"]: column for column in inspector.get_columns(table_name)} for table_name in table_names } indexes = { table_name: { index["name"] for index in inspector.get_indexes(table_name) if index.get("name") } for table_name in table_names } foreign_keys = { table_name: inspector.get_foreign_keys(table_name) for table_name in table_names } return table_names, columns, indexes, foreign_keys def _has_foreign_key( foreign_keys: list[dict], constrained_column: str, referred_table: str, ) -> bool: return any( foreign_key.get("constrained_columns") == [constrained_column] and foreign_key.get("referred_table") == referred_table and foreign_key.get("referred_columns") == ["id"] for foreign_key in foreign_keys ) def upgrade() -> None: table_names, columns, indexes, foreign_keys = _schema_snapshot() for table_name in OWNED_TABLES: if table_name not in table_names: continue table_columns = columns[table_name] added_dept_column = "deptId" not in table_columns added_creator_column = "createdById" not in table_columns if added_dept_column: op.add_column( table_name, sa.Column( "deptId", sa.String(), nullable=True, server_default=sa.text(f"'{DEFAULT_DEPT_ID}'"), ), ) table_columns["deptId"] = {"name": "deptId", "nullable": True} if added_creator_column: op.add_column(table_name, sa.Column("createdById", sa.String(), nullable=True)) table_columns["createdById"] = {"name": "createdById", "nullable": True} table_indexes = indexes[table_name] if f"ix_{table_name}_deptId" not in table_indexes: op.create_index(f"ix_{table_name}_deptId", table_name, ["deptId"]) table_indexes.add(f"ix_{table_name}_deptId") if f"ix_{table_name}_createdById" not in table_indexes: op.create_index(f"ix_{table_name}_createdById", table_name, ["createdById"]) table_indexes.add(f"ix_{table_name}_createdById") table_foreign_keys = foreign_keys[table_name] if not _has_foreign_key(table_foreign_keys, "deptId", "AdminDepartment"): op.create_foreign_key( f"fk_{table_name}_deptId", table_name, "AdminDepartment", ["deptId"], ["id"], ondelete="RESTRICT", ) table_foreign_keys.append( { "constrained_columns": ["deptId"], "referred_table": "AdminDepartment", "referred_columns": ["id"], } ) if not _has_foreign_key(table_foreign_keys, "createdById", "AdminUser"): op.create_foreign_key( f"fk_{table_name}_createdById", table_name, "AdminUser", ["createdById"], ["id"], ondelete="SET NULL", ) table_foreign_keys.append( { "constrained_columns": ["createdById"], "referred_table": "AdminUser", "referred_columns": ["id"], } ) quoted_table = f'"{table_name}"' needs_dept_backfill = added_dept_column or table_columns["deptId"].get("nullable") is True if needs_dept_backfill: op.execute( sa.text(f'UPDATE {quoted_table} SET "deptId" = :dept_id WHERE "deptId" IS NULL').bindparams( dept_id=DEFAULT_DEPT_ID ) ) op.alter_column( table_name, "deptId", existing_type=sa.String(), nullable=False, server_default=sa.text(f"'{DEFAULT_DEPT_ID}'"), ) if added_creator_column: op.execute( sa.text( f'UPDATE {quoted_table} SET "createdById" = ' '(SELECT "id" FROM "AdminUser" ORDER BY "createdAt" ASC LIMIT 1) ' 'WHERE "createdById" IS NULL' ) ) def downgrade() -> None: for table_name in reversed(OWNED_TABLES): op.drop_constraint(f"fk_{table_name}_createdById", table_name, type_="foreignkey") op.drop_constraint(f"fk_{table_name}_deptId", table_name, type_="foreignkey") op.drop_index(f"ix_{table_name}_createdById", table_name=table_name) op.drop_index(f"ix_{table_name}_deptId", table_name=table_name) op.drop_column(table_name, "createdById") op.drop_column(table_name, "deptId")