[PAY-2524, PAY-2525, PAY-2526] Add collections_containing_track field…

… to tracks (#7708)

Co-authored-by: Dharit Tantiviramanond <dharit.tan@gmail.com>

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;

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
+            == []
+        )

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):

packages/discovery-provider/scripts/lint.sh

@@ -0,0 +1,4 @@
+black .
+flake8 .
+mypy .
+isort .

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

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']}"