blakeblackshear.frigate/migrations/030_create_user_review_status.py
Josh Hawkins d87268acfe
Small tweaks (#17168)
* Clean up repeated code in auth

* Ensure review status is migrated for anonymous users
2025-03-15 07:11:45 -06:00

88 lines
3.4 KiB
Python

"""Peewee migrations -- 030_create_user_review_status.py.
This migration creates the UserReviewStatus table to track per-user review states,
migrates existing has_been_reviewed data from ReviewSegment to all users in the user table,
and drops the has_been_reviewed column. Rollback drops UserReviewStatus and restores the column.
Some examples (model - class or model_name)::
> Model = migrator.orm['model_name'] # Return model in current state by name
> migrator.sql(sql) # Run custom SQL
> migrator.python(func, *args, **kwargs) # Run python code
> migrator.create_model(Model) # Create a model (could be used as decorator)
> migrator.remove_model(model, cascade=True) # Remove a model
> migrator.add_fields(model, **fields) # Add fields to a model
> migrator.change_fields(model, **fields) # Change fields
> migrator.remove_fields(model, *field_names, cascade=True)
> migrator.rename_field(model, old_field_name, new_field_name)
> migrator.rename_table(model, new_table_name)
> migrator.add_index(model, *col_names, unique=False)
> migrator.drop_index(model, *col_names)
> migrator.add_not_null(model, *field_names)
> migrator.drop_not_null(model, *field_names)
> migrator.add_default(model, field_name, default)
"""
import peewee as pw
from frigate.models import User, UserReviewStatus
SQL = pw.SQL
def migrate(migrator, database, fake=False, **kwargs):
User._meta.database = database
UserReviewStatus._meta.database = database
migrator.sql(
"""
CREATE TABLE IF NOT EXISTS "userreviewstatus" (
"id" INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT,
"user_id" VARCHAR(30) NOT NULL,
"review_segment_id" VARCHAR(30) NOT NULL,
"has_been_reviewed" INTEGER NOT NULL DEFAULT 0,
FOREIGN KEY ("review_segment_id") REFERENCES "reviewsegment" ("id") ON DELETE CASCADE
)
"""
)
# Add unique index on (user_id, review_segment_id)
migrator.sql(
'CREATE UNIQUE INDEX IF NOT EXISTS "userreviewstatus_user_segment" ON "userreviewstatus" ("user_id", "review_segment_id")'
)
# Migrate existing has_been_reviewed data to UserReviewStatus for all users
def migrate_data():
all_users = list(User.select())
if not all_users:
return
cursor = database.execute_sql(
'SELECT "id" FROM "reviewsegment" WHERE "has_been_reviewed" = 1'
)
reviewed_segment_ids = [row[0] for row in cursor.fetchall()]
# also migrate for anonymous (unauthenticated users)
usernames = [user.username for user in all_users] + ["anonymous"]
for segment_id in reviewed_segment_ids:
for username in usernames:
UserReviewStatus.create(
user_id=username,
review_segment=segment_id,
has_been_reviewed=True,
)
if not fake: # Only run data migration if not faking
migrator.python(migrate_data)
migrator.sql('ALTER TABLE "reviewsegment" DROP COLUMN "has_been_reviewed"')
def rollback(migrator, database, fake=False, **kwargs):
migrator.sql('DROP TABLE IF EXISTS "userreviewstatus"')
# Restore has_been_reviewed column to reviewsegment (no data restoration)
migrator.sql(
'ALTER TABLE "reviewsegment" ADD COLUMN "has_been_reviewed" INTEGER NOT NULL DEFAULT 0'
)