diff --git a/tests/base.py b/tests/base.py index ed4892a092..69a6b0fde5 100644 --- a/tests/base.py +++ b/tests/base.py @@ -349,6 +349,7 @@ def generate_fixture_asset( description="Description Tree", asset_type_id=None, project_id=None, + episode_id=None, ): if asset_type_id is None: asset_type_id = self.asset_type.id @@ -361,9 +362,23 @@ def generate_fixture_asset( description=description, project_id=project_id, entity_type_id=asset_type_id, + source_id=episode_id, ) return self.asset + def generate_fixture_main_pack_episode( + self, + project_id, + name="MP", + ): + self.main_pack_episode = Entity.create( + name=name, + project_id=project_id, + entity_type_id=self.episode_type.id, + is_main_pack=True, + ) + return self.main_pack_episode + def generate_fixture_asset_character( self, name="Rabbit", description="Main char" ): diff --git a/tests/export/test_casting_to_csv.py b/tests/export/test_casting_to_csv.py index 4561879911..2d190cab9b 100644 --- a/tests/export/test_casting_to_csv.py +++ b/tests/export/test_casting_to_csv.py @@ -11,6 +11,8 @@ def setUp(self): self.generate_fixture_asset_types() def test_import_casting(self): + project_id = str(self.project.id) + mp_id = self.generate_fixture_main_pack_episode(project_id).id for name, asset_type_id in [ ("Lake", self.asset_type_environment.id), ("Mine", self.asset_type_environment.id), @@ -22,7 +24,9 @@ def test_import_casting(self): ("Victor", self.asset_type_character.id), ("John", self.asset_type_character.id), ]: - self.generate_fixture_asset(name, asset_type_id=asset_type_id) + self.generate_fixture_asset( + name, asset_type_id=asset_type_id, episode_id=mp_id + ) if name == "Lake": self.asset_lake_id = self.asset.id if name == "Mine": @@ -37,7 +41,6 @@ def test_import_casting(self): episode2_id = self.generate_fixture_episode("E02").id self.generate_fixture_sequence("SEQ01") self.generate_fixture_shot("SH01").id - project_id = str(self.project.id) path = "/import/csv/projects/%s/casting" % project_id self.upload_csv(path, "casting") diff --git a/tests/services/test_playlists_service.py b/tests/services/test_playlists_service.py index a19dd97d7f..317d2378b7 100644 --- a/tests/services/test_playlists_service.py +++ b/tests/services/test_playlists_service.py @@ -86,13 +86,15 @@ def test_get_playlist_for_episode(self): self.project.id, True ) self.assertEqual(len(playlists), 1) - - self.generate_fixture_playlist("Test main pack", for_entity="asset") + mp_id = self.generate_fixture_main_pack_episode(self.project.id).id + self.generate_fixture_playlist( + "Test main pack", for_entity="asset", episode_id=mp_id + ) self.generate_fixture_playlist( "Test all playlist", for_entity="asset", is_for_all=True ) playlists = playlists_service.all_playlists_for_episode( - self.project.id, "main" + self.project.id, mp_id ) self.assertEqual(len(playlists), 1) playlists = playlists_service.all_playlists_for_episode( diff --git a/zou/app/blueprints/crud/playlist.py b/zou/app/blueprints/crud/playlist.py index 7ca1f8a1a3..5ca4f5639c 100644 --- a/zou/app/blueprints/crud/playlist.py +++ b/zou/app/blueprints/crud/playlist.py @@ -17,7 +17,7 @@ def check_create_permissions(self, playlist): def update_data(self, data): data = super().update_data(data) - if "episode_id" in data and data["episode_id"] in ["all", "main"]: + if "episode_id" in data and data["episode_id"] in ["all"]: data["episode_id"] = None if "task_type_id" in data and not fields.is_valid_id( data["task_type_id"] diff --git a/zou/app/blueprints/crud/project.py b/zou/app/blueprints/crud/project.py index 42d36ecf37..6272692c45 100644 --- a/zou/app/blueprints/crud/project.py +++ b/zou/app/blueprints/crud/project.py @@ -87,6 +87,12 @@ def post_creation(self, project): project_dict["first_episode_id"] = fields.serialize_value( episode["id"] ) + episode = shots_service.create_episode( + project.id, + "MP", + created_by=persons_service.get_current_user()["id"], + is_main_pack=True, + ) user_service.clear_project_cache() projects_service.clear_project_cache("") return project_dict diff --git a/zou/app/blueprints/export/csv/casting.py b/zou/app/blueprints/export/csv/casting.py index 540d924907..79ce99e39b 100644 --- a/zou/app/blueprints/export/csv/casting.py +++ b/zou/app/blueprints/export/csv/casting.py @@ -128,9 +128,7 @@ def build_results( is_shot_casting=False, ): results = [] - if episode_id == "main": - results = self.build_main_pack_results(project_id) - elif episode_id == "all": + if episode_id == "all": results = self.build_episodes_results(project_id) elif is_shot_casting: results = self.build_shot_results( @@ -291,7 +289,6 @@ def build_asset_results( entity_link.label, ) ) - results += self.build_main_pack_results(project_id) else: query = query.add_columns( ParentAssetType.name, @@ -325,62 +322,6 @@ def build_asset_results( return results - def build_main_pack_results(self, project_id): - results = [] - ParentAsset = aliased(Entity, name="parent_asset") - ParentAssetType = aliased(EntityType, name="parent_asset_type") - Asset = aliased(Entity, name="asset") - AssetType = aliased(EntityType, name="asset_type") - shot_type = shots_service.get_shot_type() - episode_type = shots_service.get_episode_type() - - query = ( - EntityLink.query.join( - ParentAsset, EntityLink.entity_in_id == ParentAsset.id - ) - .join( - ParentAssetType, - ParentAsset.entity_type_id == ParentAssetType.id, - ) - .join(Asset, EntityLink.entity_out_id == Asset.id) - .join(AssetType, Asset.entity_type_id == AssetType.id) - .filter(ParentAsset.project_id == project_id) - .filter(ParentAsset.source_id == None) - .filter(ParentAssetType.id != shot_type["id"]) - .filter(ParentAssetType.id != episode_type["id"]) - ) - query = query.add_columns( - ParentAssetType.name, - ParentAsset.name, - AssetType.name, - Asset.name, - ).order_by( - ParentAssetType.name, - ParentAsset.name, - AssetType.name, - Asset.name, - ) - for ( - entity_link, - parent_asset_type_name, - parent_name, - asset_type_name, - asset_name, - ) in query.all(): - results.append( - ( - "MP", - "", - parent_asset_type_name, - parent_name, - asset_type_name, - asset_name, - entity_link.nb_occurences, - entity_link.label, - ) - ) - return results - def build_episodes_results(self, project_id): results = [] ParentAsset = aliased(Entity, name="parent_asset") diff --git a/zou/app/blueprints/playlists/resources.py b/zou/app/blueprints/playlists/resources.py index 8d97c2e5c0..334b7aa4d1 100644 --- a/zou/app/blueprints/playlists/resources.py +++ b/zou/app/blueprints/playlists/resources.py @@ -96,7 +96,7 @@ def get(self, project_id, episode_id): page = self.get_page() sort_by = self.get_sort_by() task_type_id = self.get_text_parameter("task_type_id") - if episode_id not in ["main", "all"]: + if episode_id not in ["all"]: shots_service.get_episode(episode_id) return playlists_service.all_playlists_for_episode( project_id, @@ -234,8 +234,6 @@ def get(self, playlist_id, build_job_id): episode_name = episode["name"] elif playlist["is_for_all"]: episode_name = "all assets" - else: - episode_name = "main pack" context_name += "_%s" % slugify.slugify( episode_name, separator="_" ) @@ -365,8 +363,6 @@ def get(self, playlist_id): episode_name = episode["name"] elif playlist["is_for_all"]: episode_name = "all assets" - else: - episode_name = "main pack" context_name += "_%s" % slugify.slugify( episode_name, separator="_" ) diff --git a/zou/app/blueprints/source/csv/assets.py b/zou/app/blueprints/source/csv/assets.py index fcc58e36c2..0a1e9d8e83 100644 --- a/zou/app/blueprints/source/csv/assets.py +++ b/zou/app/blueprints/source/csv/assets.py @@ -165,7 +165,7 @@ def import_row(self, row, project_id): episode_id = None if self.is_tv_show: - if episode_name not in [None, "MP"] + list(self.episodes.keys()): + if episode_name not in list(self.episodes.keys()): self.episodes[episode_name] = shots_service.create_episode( project_id, episode_name, created_by=self.current_user_id )["id"] diff --git a/zou/app/blueprints/source/csv/casting.py b/zou/app/blueprints/source/csv/casting.py index c34bbbc21f..cefd52e084 100644 --- a/zou/app/blueprints/source/csv/casting.py +++ b/zou/app/blueprints/source/csv/casting.py @@ -95,8 +95,6 @@ def get_episode_key(self, episode): def import_row(self, row, project_id): asset_key = slugify(f"{row['Asset Type']}{row['Asset']}") - if row.get("Episode") in ["MP", None]: - row["Episode"] = "" target_key = slugify(f"{row['Episode']}{row['Parent']}{row['Name']}") occurences = 1 if len(row["Occurences"]) > 0: diff --git a/zou/app/blueprints/source/csv/task_type_estimations.py b/zou/app/blueprints/source/csv/task_type_estimations.py index 96e0c654bf..4926121e36 100644 --- a/zou/app/blueprints/source/csv/task_type_estimations.py +++ b/zou/app/blueprints/source/csv/task_type_estimations.py @@ -60,7 +60,7 @@ def prepare_import(self, project_id, task_type_id, episode_id=None): asset_types_map[asset_type["id"]] = slugify(asset_type["name"]) criterions_assets = {"project_id": project_id} - if episode_id is not None and episode_id not in ["main", "all"]: + if episode_id is not None and episode_id not in ["all"]: criterions_assets["source_id"] = episode_id assets = assets_service.get_assets(criterions_assets) for asset in assets: @@ -72,7 +72,7 @@ def prepare_import(self, project_id, task_type_id, episode_id=None): sequences_map = {} criterions_sequences = {"project_id": project_id} - if episode_id is not None and episode_id not in ["main", "all"]: + if episode_id is not None and episode_id not in ["all"]: criterions_sequences["parent_id"] = episode_id sequences = shots_service.get_sequences(criterions_sequences) for sequence in sequences: diff --git a/zou/app/models/entity.py b/zou/app/models/entity.py index 069fb35fdf..b6acaa5483 100644 --- a/zou/app/models/entity.py +++ b/zou/app/models/entity.py @@ -102,6 +102,9 @@ class Entity(db.Model, BaseMixin, SerializerMixin): is_shared = db.Column(db.Boolean, default=False, nullable=False) + # specific to episodes + is_main_pack = db.Column(db.Boolean, default=False, nullable=False) + status = db.Column( ChoiceType(ENTITY_STATUSES), default="running", nullable=False ) diff --git a/zou/app/services/assets_service.py b/zou/app/services/assets_service.py index 32ebb13ba1..de0815549c 100644 --- a/zou/app/services/assets_service.py +++ b/zou/app/services/assets_service.py @@ -35,6 +35,7 @@ AssetNotFoundException, AssetInstanceNotFoundException, AssetTypeNotFoundException, + EpisodeNotFoundException, ) @@ -217,9 +218,7 @@ def get_assets_and_tasks(criterions={}, with_episode_ids=False): if "episode_id" in criterions: episode_id = criterions["episode_id"] - if episode_id == "main": - tasks_query = tasks_query.filter(Entity.source_id == None) - elif episode_id != "all": + if episode_id != "all": tasks_query = tasks_query.outerjoin( EntityLink, EntityLink.entity_out_id == Entity.id ) @@ -250,7 +249,6 @@ def get_assets_and_tasks(criterions={}, with_episode_ids=False): Episode.project_id == criterions["project_id"] ) if "episode_id" in criterions and criterions["episode_id"] not in [ - "main", "all", ]: episode_links_query = episode_links_query.filter( @@ -586,6 +584,36 @@ def create_asset_types(asset_type_names): return asset_types +def get_main_pack_raw(project_id): + """ + Get the main pack episode of a project. + """ + episode_type = shots_service.get_episode_type() + if episode_type is None: + episode_type = shots_service.get_episode_type() + + try: + main_pack_episode = Entity.get_by( + entity_type_id=episode_type["id"], + is_main_pack=True, + project_id=project_id, + ) + except StatementError: + raise EpisodeNotFoundException + + if main_pack_episode is None: + raise EpisodeNotFoundException + return main_pack_episode + + +@cache.memoize_function(120) +def get_main_pack(project_id): + """ + Get the main pack episode of a project as a dictionnary. + """ + return get_main_pack_raw(project_id).serialize(obj_type="Episode") + + def create_asset( project_id, asset_type_id, @@ -603,6 +631,11 @@ def create_asset( asset_type = get_asset_type_raw(asset_type_id) if source_id is not None and len(source_id) < 36: source_id = None + + if project.production_type == "tvshow" and source_id is None: + main_pack = get_main_pack_raw(project_id) + source_id = main_pack.id + asset = Entity.create( project_id=project_id, entity_type_id=asset_type_id, @@ -791,7 +824,7 @@ def get_shared_assets_used_in_project(project_id, episode_id=None): .filter(Entity.project_id != project_id) ) - if episode_id is not None and episode_id not in ["main", "all"]: + if episode_id is not None and episode_id not in ["all"]: assets = assets.filter(Sequence.parent_id == episode_id) return Entity.serialize_list(assets.all(), obj_type="Asset") diff --git a/zou/app/services/breakdown_service.py b/zou/app/services/breakdown_service.py index fb66ae92e3..4f1b9c68e0 100644 --- a/zou/app/services/breakdown_service.py +++ b/zou/app/services/breakdown_service.py @@ -148,8 +148,6 @@ def get_sequence_casting(sequence_id, project_id=None, episode_id=None): are casting for given shot. """ castings = {} - if episode_id == "main": - return {} Shot = aliased(Entity, name="shot") Sequence = aliased(Entity, name="sequence") diff --git a/zou/app/services/playlists_service.py b/zou/app/services/playlists_service.py index ac7808552d..6b5776edb1 100644 --- a/zou/app/services/playlists_service.py +++ b/zou/app/services/playlists_service.py @@ -12,7 +12,6 @@ from flask_fs.errors import FileNotFound from slugify import slugify -from sqlalchemy import or_ from zou.app import config from zou.app.stores import file_store @@ -97,15 +96,7 @@ def all_playlists_for_episode( if task_type_id is not None and len(task_type_id) > 0: query = query.filter(Playlist.task_type_id == task_type_id) - if episode_id == "main": - query = ( - query.filter(Playlist.episode_id == None) - .filter(Playlist.project_id == project_id) - .filter( - or_(Playlist.is_for_all is None, Playlist.is_for_all == False) - ) - ) - elif episode_id == "all": + if episode_id == "all": query = ( query.filter(Playlist.episode_id == None) .filter(Playlist.project_id == project_id) diff --git a/zou/app/services/shots_service.py b/zou/app/services/shots_service.py index e7390e2faa..57dfc5f3fb 100644 --- a/zou/app/services/shots_service.py +++ b/zou/app/services/shots_service.py @@ -932,6 +932,7 @@ def create_episode( description="", data={}, created_by=None, + is_main_pack=False, ): """ Create episode for given project. @@ -951,6 +952,7 @@ def create_episode( description=description, data=data, created_by=created_by, + is_main_pack=is_main_pack, ) events.emit( "episode:new", {"episode_id": episode.id}, project_id=project_id diff --git a/zou/app/services/tasks_service.py b/zou/app/services/tasks_service.py index 59c6a5adf6..5eb5e20de7 100644 --- a/zou/app/services/tasks_service.py +++ b/zou/app/services/tasks_service.py @@ -42,7 +42,6 @@ from zou.app.services.exception import ( CommentNotFoundException, - EpisodeNotFoundException, PersonNotFoundException, TaskNotFoundException, TaskStatusNotFoundException, @@ -967,11 +966,8 @@ def get_person_tasks(person_id, projects, is_done=None): if episode_id is None: episode_id = entity_source_id if episode_id is not None and episode_id != "": - try: - episode = shots_service.get_episode(episode_id) - episode_name = episode["name"] - except EpisodeNotFoundException: - episode_name = "MP" + episode = shots_service.get_episode(episode_id) + episode_name = episode["name"] task_dict = get_task(str(task.id), relations=True) if entity_type_name == "Sequence" and entity_parent_id is not None: @@ -2054,11 +2050,8 @@ def get_open_tasks( if episode_id is None: episode_id = entity_source_id if episode_id is not None and episode_id != "": - try: - episode = shots_service.get_episode(episode_id) - episode_name = episode["name"] - except EpisodeNotFoundException: - episode_name = "MP" + episode = shots_service.get_episode(episode_id) + episode_name = episode["name"] task_dict = get_task(str(task.id), relations=True) if entity_type_name == "Sequence" and entity_parent_id is not None: diff --git a/zou/migrations/versions/34ebd1f0ea86_add_entity_is_main_pack_put_existing_.py b/zou/migrations/versions/34ebd1f0ea86_add_entity_is_main_pack_put_existing_.py new file mode 100644 index 0000000000..fb3a155d25 --- /dev/null +++ b/zou/migrations/versions/34ebd1f0ea86_add_entity_is_main_pack_put_existing_.py @@ -0,0 +1,299 @@ +"""Add entity.is_main_pack + put existing main pack + +Revision ID: 34ebd1f0ea86 +Revises: 307edd8c639d +Create Date: 2025-03-11 17:24:00.856665 + +""" + +from alembic import op +import sqlalchemy as sa +from sqlalchemy.orm.session import Session +from zou.migrations.utils.base import BaseMixin +from sqlalchemy.ext.declarative import declarative_base +from sqlalchemy_utils import ChoiceType +from sqlalchemy.dialects.postgresql import JSONB +from sqlalchemy_utils import UUIDType + + +# revision identifiers, used by Alembic. +revision = "34ebd1f0ea86" +down_revision = "307edd8c639d" +branch_labels = None +depends_on = None + +base = declarative_base() + +PROJECT_STYLES = [ + ("2d", "2D Animation"), + ("2dpaper", "2D Animation (Paper)"), + ("3d", "3D Animation"), + ("2d3d", "2D/3D Animation"), + ("ar", "Augmented Reality"), + ("vfx", "VFX"), + ("stop-motion", "Stop Motion"), + ("motion-design", "Motion Design"), + ("archviz", "Archviz"), + ("commercial", "Commercial"), + ("catalog", "Catalog"), + ("immersive", "Immersive Experience"), + ("nft", "NFT Collection"), + ("video-game", "Video Game"), + ("vr", "Virtual Reality"), +] + + +class Project(base, BaseMixin): + """ + Describes a production the studio works on. + """ + + __tablename__ = "project" + + name = sa.Column(sa.String(80), nullable=False, unique=True, index=True) + code = sa.Column(sa.String(80)) + description = sa.Column(sa.Text()) + shotgun_id = sa.Column(sa.Integer) + file_tree = sa.Column(JSONB) + data = sa.Column(JSONB) + has_avatar = sa.Column(sa.Boolean(), default=False) + fps = sa.Column(sa.String(10), default=25) + ratio = sa.Column(sa.String(10), default="16:9") + resolution = sa.Column(sa.String(12), default="1920x1080") + production_type = sa.Column(sa.String(20), default="short") + production_style = sa.Column( + ChoiceType(PROJECT_STYLES), default="2d3d", nullable=False + ) + start_date = sa.Column(sa.Date()) + end_date = sa.Column(sa.Date()) + man_days = sa.Column(sa.Integer) + nb_episodes = sa.Column(sa.Integer, default=0) + episode_span = sa.Column(sa.Integer, default=0) + max_retakes = sa.Column(sa.Integer, default=0) + is_clients_isolated = sa.Column(sa.Boolean(), default=False) + is_preview_download_allowed = sa.Column(sa.Boolean(), default=False) + is_set_preview_automated = sa.Column(sa.Boolean(), default=False) + homepage = sa.Column(sa.String(80), default="assets") + is_publish_default_for_artists = sa.Column(sa.Boolean(), default=False) + hd_bitrate_compression = sa.Column(sa.Integer, default=28) + ld_bitrate_compression = sa.Column(sa.Integer, default=6) + + project_status_id = sa.Column( + UUIDType(binary=False), sa.ForeignKey("project_status.id"), index=True + ) + + default_preview_background_file_id = sa.Column( + UUIDType(binary=False), + sa.ForeignKey("preview_background_file.id"), + default=None, + index=True, + ) + + +ENTITY_STATUSES = [ + ("standby", "Stand By"), + ("running", "Running"), + ("complete", "Complete"), + ("canceled", "Canceled"), +] + + +class Entity(base, BaseMixin): + """ + Base model to represent assets, shots, sequences, episodes and scenes. + They have different meaning but they share the same behaviour toward + tasks and files. + """ + + __tablename__ = "entity" + + id = BaseMixin.id + name = sa.Column(sa.String(160), nullable=False) + code = sa.Column(sa.String(160)) # To store sanitized version of name + description = sa.Column(sa.Text()) + shotgun_id = sa.Column(sa.Integer) + canceled = sa.Column(sa.Boolean, default=False) + + nb_frames = sa.Column(sa.Integer) # Specific to shots + nb_entities_out = sa.Column(sa.Integer, default=0) + is_casting_standby = sa.Column(sa.Boolean, default=False) + + is_shared = sa.Column(sa.Boolean, default=False, nullable=False) + + # specific to episodes + is_main_pack = sa.Column(sa.Boolean, default=False, nullable=False) + + status = sa.Column( + ChoiceType(ENTITY_STATUSES), default="running", nullable=False + ) + + project_id = sa.Column( + UUIDType(binary=False), + sa.ForeignKey("project.id"), + nullable=False, + index=True, + ) + entity_type_id = sa.Column( + UUIDType(binary=False), + sa.ForeignKey("entity_type.id"), + nullable=False, + index=True, + ) + + parent_id = sa.Column( + UUIDType(binary=False), sa.ForeignKey("entity.id"), index=True + ) # sequence or episode + + source_id = sa.Column( + UUIDType(binary=False), + sa.ForeignKey("entity.id"), + index=True, + nullable=True, + ) # if the entity is generated from another one (like shots from scene). + + preview_file_id = sa.Column( + UUIDType(binary=False), + ) + data = sa.Column(JSONB) + + ready_for = sa.Column( + UUIDType(binary=False), + ) + + created_by = sa.Column( + UUIDType(binary=False), + nullable=True, + ) + + +class EntityType(base, BaseMixin): + """ + Type of entities. It can describe either an asset type, or tell if target + entity is a shot, sequence, episode or layout scene. + """ + + __tablename__ = "entity_type" + + name = sa.Column(sa.String(30), unique=True, nullable=False, index=True) + short_name = sa.Column(sa.String(20)) + description = sa.Column(sa.Text()) + archived = sa.Column(sa.Boolean(), default=False) + + +class Playlist(base, BaseMixin): + """ + Describes a playlist. The goal is to review a set of shipped materials. + """ + + __tablename__ = "playlist" + + name = sa.Column(sa.String(80), nullable=False) + shots = sa.Column(JSONB) + + project_id = sa.Column( + UUIDType(binary=False), sa.ForeignKey("project.id"), index=True + ) + episode_id = sa.Column( + UUIDType(binary=False), sa.ForeignKey("entity.id"), index=True + ) + task_type_id = sa.Column( + UUIDType(binary=False), sa.ForeignKey("task_type.id"), index=True + ) + for_client = sa.Column(sa.Boolean(), default=False, index=True) + for_entity = sa.Column(sa.String(10), default="shot", index=True) + is_for_all = sa.Column(sa.Boolean, default=False) + + +def upgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("entity", schema=None) as batch_op: + batch_op.add_column( + sa.Column( + "is_main_pack", + sa.Boolean(), + server_default=sa.text("false"), + default=False, + nullable=False, + ) + ) + with op.batch_alter_table("entity", schema=None) as batch_op: + session = Session(bind=op.get_bind()) + episode_entity_type_id = ( + session.query(EntityType.id) + .where(EntityType.name == "Episode") + .first() + )[0] + temporal_entity_type_ids = [ + i[0] + for i in session.query(EntityType.id) + .where( + EntityType.name.not_in( + ["Shot", "Scene", "Sequence", "Episode", "Edit", "Concept"] + ) + ) + .distinct() + .all() + ] + for project_id in ( + session.query(Project.id) + .where(Project.production_type == "tvshow") + .all() + ): + main_pack_episode = Entity( + name="MP", + project_id=project_id[0], + entity_type_id=episode_entity_type_id, + is_main_pack=True, + ) + session.add(main_pack_episode) + session.commit() + session.refresh(main_pack_episode) + + session.query(Entity).filter( + Entity.project_id == project_id[0], + Entity.entity_type_id.in_(temporal_entity_type_ids), + Entity.source_id == None, + ).update({Entity.source_id: main_pack_episode.id}) + session.commit() + + session.query(Playlist).filter( + Playlist.project_id == project_id[0], + Playlist.episode_id == None, + ).update({Playlist.episode_id: main_pack_episode.id}) + session.commit() + + # ### end Alembic commands ### + + +def downgrade(): + # ### commands auto generated by Alembic - please adjust! ### + with op.batch_alter_table("entity", schema=None) as batch_op: + session = Session(bind=op.get_bind()) + + for project_id in ( + session.query(Project.id) + .where(Project.production_type == "tvshow") + .all() + ): + main_pack_episode = ( + session.query(Entity) + .where( + Entity.project_id == project_id[0], + Entity.is_main_pack == True, + ) + .first() + ) + session.query(Entity).filter( + Entity.project_id == project_id[0], + Entity.source_id == main_pack_episode.id, + ).update({Entity.source_id: None}) + + session.query(Playlist).filter( + Playlist.project_id == project_id[0], + Playlist.episode_id == main_pack_episode.id, + ).update({Playlist.episode_id: None}) + session.query(Entity).where(Entity.is_main_pack == True).delete() + session.commit() + batch_op.drop_column("is_main_pack") + + # ### end Alembic commands ###