From 4136985ced996e06f2c7737e592bd03f60867e80 Mon Sep 17 00:00:00 2001 From: Andrew Mendelsohn Date: Mon, 26 Feb 2024 13:46:08 -0800 Subject: [PATCH] [PAY-2524, PAY-2525, PAY-2526] Add collections_containing_track field to tracks (#7708) Co-authored-by: Dharit Tantiviramanond --- .../0056_playlists_containing_track.sql | 18 ++ .../entity_manager/test_playlist_tracks.py | 210 +++++++++++++++++- .../integration_tests/utils.py | 3 + packages/discovery-provider/scripts/lint.sh | 4 + .../src/models/tracks/track.py | 4 + .../tasks/entity_manager/entities/playlist.py | 39 +++- 6 files changed, 272 insertions(+), 6 deletions(-) create mode 100644 packages/discovery-provider/ddl/migrations/0056_playlists_containing_track.sql create mode 100755 packages/discovery-provider/scripts/lint.sh diff --git a/packages/discovery-provider/ddl/migrations/0056_playlists_containing_track.sql b/packages/discovery-provider/ddl/migrations/0056_playlists_containing_track.sql new file mode 100644 index 00000000000..e21ece18eb1 --- /dev/null +++ b/packages/discovery-provider/ddl/migrations/0056_playlists_containing_track.sql @@ -0,0 +1,18 @@ +BEGIN; + +ALTER TABLE tracks +ADD column IF NOT EXISTS playlists_containing_track INTEGER[] NOT NULL DEFAULT '{}'; + +UPDATE tracks SET playlists_containing_track = subquery.playlists_containing_track FROM ( + SELECT + CAST(jsonb_array_elements(playlist_contents->'track_ids')->>'track' AS INTEGER) AS track_id, + ARRAY_AGG(playlist_id) AS playlists_containing_track + FROM playlists + WHERE playlist_id IS NOT NULL + GROUP BY track_id +) AS subquery +WHERE tracks.track_id = subquery.track_id + AND subquery.track_id IS NOT NULL +AND tracks.track_id IS NOT NULL; + +COMMIT; diff --git a/packages/discovery-provider/integration_tests/tasks/entity_manager/test_playlist_tracks.py b/packages/discovery-provider/integration_tests/tasks/entity_manager/test_playlist_tracks.py index ca1e3c93b67..4450eaeae32 100644 --- a/packages/discovery-provider/integration_tests/tasks/entity_manager/test_playlist_tracks.py +++ b/packages/discovery-provider/integration_tests/tasks/entity_manager/test_playlist_tracks.py @@ -10,14 +10,17 @@ from integration_tests.utils import populate_mock_db from src.challenges.challenge_event_bus import ChallengeEventBus, setup_challenge_bus from src.models.playlists.playlist_track import PlaylistTrack +from src.models.tracks.track import Track from src.tasks.entity_manager.entity_manager import entity_manager_update from src.tasks.entity_manager.utils import PLAYLIST_ID_OFFSET from src.utils.db_session import get_db logger = logging.getLogger(__name__) -# Insert Playlist with two new tracks and check that a notification is created for the track owners +PLAYLIST_ID = PLAYLIST_ID_OFFSET + 1 + now = datetime.now() + entities = { "users": [ {"user_id": 1, "handle": "user-1", "wallet": "user1wallet"}, @@ -119,7 +122,7 @@ def get_events_side_effect(_, tx_receipt): { "args": AttributeDict( { - "_entityId": PLAYLIST_ID_OFFSET + 1, + "_entityId": PLAYLIST_ID, "_entityType": "Playlist", "_userId": 1, "_action": "Create", @@ -153,6 +156,27 @@ def test_create_playlist(app, mocker): for id in [30, 40]: assert not any([relation.track_id == id for relation in relations]) + assert session.query(Track).filter( + Track.track_id == 10 + ).first().playlists_containing_track == [PLAYLIST_ID] + assert session.query(Track).filter( + Track.track_id == 20 + ).first().playlists_containing_track == [PLAYLIST_ID] + assert ( + session.query(Track) + .filter(Track.track_id == 30) + .first() + .playlists_containing_track + == [] + ) + assert ( + session.query(Track) + .filter(Track.track_id == 40) + .first() + .playlists_containing_track + == [] + ) + # Add tracks to a playlist add_tracks_to_playlist_tx_receipts = { @@ -194,6 +218,27 @@ def test_add_tracks_to_playlist(app, mocker): for id in [10, 40]: assert not any([relation.track_id == id for relation in relations]) + assert ( + session.query(Track) + .filter(Track.track_id == 10) + .first() + .playlists_containing_track + == [] + ) + assert session.query(Track).filter( + Track.track_id == 20 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert session.query(Track).filter( + Track.track_id == 30 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert ( + session.query(Track) + .filter(Track.track_id == 40) + .first() + .playlists_containing_track + == [] + ) + # Remove a track from an playlist remove_track_from_playlist_tx_receipts = { @@ -258,6 +303,31 @@ def test_remove_track_from_playlist(app, mocker): for id in [10, 40]: assert not any([relation.track_id == id for relation in relations]) + assert ( + session.query(Track) + .filter(Track.track_id == 10) + .first() + .playlists_containing_track + == [] + ) + assert ( + session.query(Track) + .filter(Track.track_id == 20) + .first() + .playlists_containing_track + == [] + ) + assert session.query(Track).filter( + Track.track_id == 30 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert ( + session.query(Track) + .filter(Track.track_id == 40) + .first() + .playlists_containing_track + == [] + ) + # Remove a track from a playlist and then restore it restore_removed_track_to_playlist_tx_receipts = { @@ -337,6 +407,27 @@ def test_restore_removed_track_to_playlist(app, mocker): ] ) + assert ( + session.query(Track) + .filter(Track.track_id == 10) + .first() + .playlists_containing_track + == [] + ) + assert session.query(Track).filter( + Track.track_id == 20 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert session.query(Track).filter( + Track.track_id == 30 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert ( + session.query(Track) + .filter(Track.track_id == 40) + .first() + .playlists_containing_track + == [] + ) + # Create a playlist then reorder the tracks reorder_playlist_tracks_tx_receipts = { @@ -391,3 +482,118 @@ def test_reorder_playlist_tracks(app, mocker): assert any([relation.track_id == id for relation in relations]) for id in [10, 40]: assert not any([relation.track_id == id for relation in relations]) + + assert ( + session.query(Track) + .filter(Track.track_id == 10) + .first() + .playlists_containing_track + == [] + ) + assert session.query(Track).filter( + Track.track_id == 20 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert session.query(Track).filter( + Track.track_id == 30 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert ( + session.query(Track) + .filter(Track.track_id == 40) + .first() + .playlists_containing_track + == [] + ) + + +# Add track to multiple playlists +add_tracks_to_multiple_playlists_tx_receipts = { + "UpdatePlaylistTracklistUpdate": [ + { + "args": AttributeDict( + { + "_entityId": PLAYLIST_ID_OFFSET, + "_entityType": "Playlist", + "_userId": 1, + "_action": "Update", + "_metadata": f'{{"cid": "PlaylistTracklistUpdate", "data": {json.dumps(test_metadata["PlaylistTracklistUpdate"])}, "timestamp": {datetime.timestamp(now)}}}', + "_signer": "user1wallet", + } + ) + } + ], + "CreatePlaylist": [ + { + "args": AttributeDict( + { + "_entityId": PLAYLIST_ID_OFFSET + 1, + "_entityType": "Playlist", + "_userId": 1, + "_action": "Create", + "_metadata": f'{{"cid": "QmCreatePlaylist1", "data": {json.dumps(test_metadata["PlaylistToCreate"])}}}', + "_signer": "user1wallet", + } + ) + }, + ], +} + + +def test_add_track_to_multiple_playlists(app, mocker): + db, update_task, entity_manager_txs = setup_db( + app, mocker, entities, add_tracks_to_multiple_playlists_tx_receipts + ) + + with db.scoped_session() as session: + entity_manager_update( + update_task, + session, + entity_manager_txs, + block_number=0, + block_timestamp=1585336422, + block_hash=hex(0), + ) + playlist_1_relations: List[PlaylistTrack] = ( + session.query(PlaylistTrack) + .filter(PlaylistTrack.playlist_id == PLAYLIST_ID_OFFSET) + .all() + ) + assert len(playlist_1_relations) == 2 + for id in [20, 30]: + assert any([relation.track_id == id for relation in playlist_1_relations]) + for id in [10, 40]: + assert not any( + [relation.track_id == id for relation in playlist_1_relations] + ) + + playlist_2_relations: List[PlaylistTrack] = ( + session.query(PlaylistTrack) + .filter(PlaylistTrack.playlist_id == PLAYLIST_ID_OFFSET + 1) + .all() + ) + assert len(playlist_2_relations) == 2 + for id in [10, 20]: + assert any([relation.track_id == id for relation in playlist_2_relations]) + for id in [30, 40]: + assert not any( + [relation.track_id == id for relation in playlist_2_relations] + ) + + assert session.query(Track).filter( + Track.track_id == 10 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET + 1] + assert session.query(Track).filter( + Track.track_id == 20 + ).first().playlists_containing_track == [ + PLAYLIST_ID_OFFSET, + PLAYLIST_ID_OFFSET + 1, + ] + assert session.query(Track).filter( + Track.track_id == 30 + ).first().playlists_containing_track == [PLAYLIST_ID_OFFSET] + assert ( + session.query(Track) + .filter(Track.track_id == 40) + .first() + .playlists_containing_track + == [] + ) diff --git a/packages/discovery-provider/integration_tests/utils.py b/packages/discovery-provider/integration_tests/utils.py index f6962dcca5c..9b0287738f2 100644 --- a/packages/discovery-provider/integration_tests/utils.py +++ b/packages/discovery-provider/integration_tests/utils.py @@ -231,6 +231,9 @@ def populate_mock_db(db, entities, block_offset=None): is_playlist_upload=track_meta.get("is_playlist_upload", False), track_cid=track_meta.get("track_cid", None), ai_attribution_user_id=track_meta.get("ai_attribution_user_id", None), + playlists_containing_track=track_meta.get( + "playlists_containing_track", [] + ), ) session.add(track) for i, track_price_history_meta in enumerate(track_price_history): diff --git a/packages/discovery-provider/scripts/lint.sh b/packages/discovery-provider/scripts/lint.sh new file mode 100755 index 00000000000..4a7513bd01a --- /dev/null +++ b/packages/discovery-provider/scripts/lint.sh @@ -0,0 +1,4 @@ +black . +flake8 . +mypy . +isort . diff --git a/packages/discovery-provider/src/models/tracks/track.py b/packages/discovery-provider/src/models/tracks/track.py index 8314609d970..da814534abb 100644 --- a/packages/discovery-provider/src/models/tracks/track.py +++ b/packages/discovery-provider/src/models/tracks/track.py @@ -1,4 +1,5 @@ from sqlalchemy import ( + ARRAY, Boolean, Column, DateTime, @@ -86,6 +87,9 @@ class Track(Base, RepresentableMixin): is_download_gated = Column(Boolean, nullable=False, server_default=text("false")) download_conditions = Column(JSONB(True)) is_playlist_upload = Column(Boolean, nullable=False, server_default=text("false")) + playlists_containing_track = Column( + ARRAY(Integer()), server_default="ARRAY[]::INTEGER[]" + ) ai_attribution_user_id = Column(Integer, nullable=True) block1 = relationship( # type: ignore diff --git a/packages/discovery-provider/src/tasks/entity_manager/entities/playlist.py b/packages/discovery-provider/src/tasks/entity_manager/entities/playlist.py index 49c09e5cd5d..63c700fea20 100644 --- a/packages/discovery-provider/src/tasks/entity_manager/entities/playlist.py +++ b/packages/discovery-provider/src/tasks/entity_manager/entities/playlist.py @@ -6,6 +6,7 @@ from src.models.playlists.playlist import Playlist from src.models.playlists.playlist_route import PlaylistRoute from src.models.playlists.playlist_track import PlaylistTrack +from src.models.tracks.track import Track from src.tasks.entity_manager.utils import ( CHARACTER_LIMIT_DESCRIPTION, PLAYLIST_ID_OFFSET, @@ -131,10 +132,22 @@ def update_playlist_tracks(params: ManageEntityParameters, playlist_record: Play ) # delete relations that previously existed but are not in the updated list - for relation in existing_playlist_tracks: - if relation.track_id not in updated_track_ids: - relation.is_removed = True - relation.updated_at = params.block_datetime + for playlist_track in existing_playlist_tracks: + if playlist_track.track_id not in updated_track_ids: + playlist_track.is_removed = True + playlist_track.updated_at = params.block_datetime + track = ( + session.query(Track) + .filter(Track.track_id == playlist_track.track_id) + .first() + ) + if track: + track.updated_at = params.block_datetime + track.playlists_containing_track = [ + collection_id + for collection_id in (track.playlists_containing_track or []) + if collection_id != playlist["playlist_id"] + ] for track_id in updated_track_ids: # add row for each track that is not already in the table @@ -148,10 +161,28 @@ def update_playlist_tracks(params: ManageEntityParameters, playlist_record: Play ) # upsert to handle duplicates session.merge(new_playlist_track) + track = session.query(Track).filter(Track.track_id == track_id).first() + if track: + track.updated_at = params.block_datetime + track.playlists_containing_track = list( + set( + (track.playlists_containing_track or []) + + [playlist["playlist_id"]] + ) + ) elif existing_tracks[track_id].is_removed: # recover deleted relation (track was previously removed then re-added) existing_tracks[track_id].is_removed = False existing_tracks[track_id].updated_at = params.block_datetime + track = session.query(Track).filter(Track.track_id == track_id).first() + if track: + track.updated_at = params.block_datetime + track.playlists_containing_track = list( + set( + (track.playlists_containing_track or []) + + [playlist["playlist_id"]] + ) + ) params.logger.info( f"playlists.py | Updated playlist tracks for {playlist['playlist_id']}"