[PAY-481][PAY-483][PAY-484] Add premium content write changes to DN (#…

…3752)

* Add premium content migration and new columns for track model and update track schema

* Add premium content access checker

* Update logic and write first integration tests

* Refactor

* Update premium content migration down revision

* Update track indexing integration tests

* Use keyword args for premium content access check

* Return dictionary for premium content access check

* Update premium track index

* Update function signature

* Refactor premium content access check to take a list and process batch content

* Make entity manager delete track consistent with old flow

Co-authored-by: Saliou Diallo <saliou@audius.co>

discovery-provider/alembic/versions/aab240348d73_add_premium_content_columns.py

@@ -0,0 +1,50 @@
+"""add-premium-content-columns
+
+Revision ID: aab240348d73
+Revises: 5d3f95470222
+Create Date: 2022-08-22 20:22:22.439424
+
+"""
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision = "aab240348d73"
+down_revision = "551f5fc03862"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.add_column(
+        "tracks",
+        sa.Column("is_premium", sa.Boolean(), nullable=False, server_default="false"),
+    )
+    op.add_column(
+        "tracks",
+        sa.Column("premium_conditions", postgresql.JSONB(), nullable=True),
+    )
+
+    connection = op.get_bind()
+    connection.execute(
+        """
+        begin;
+            CREATE INDEX IF NOT EXISTS track_is_premium_idx ON tracks (is_premium, is_current, is_delete);
+        commit;
+    """
+    )
+
+
+def downgrade():
+    connection = op.get_bind()
+    connection.execute(
+        """
+        begin;
+            DROP INDEX IF EXISTS track_is_premium_idx;
+        commit;
+    """
+    )
+
+    op.drop_column("tracks", "premium_conditions")
+    op.drop_column("tracks", "is_premium")

discovery-provider/integration_tests/premium_content/test_access.py

@@ -0,0 +1,80 @@
+from integration_tests.utils import populate_mock_db
+from src.premium_content.premium_content_access_checker import (
+    PremiumContentAccessChecker,
+)
+from src.utils.db_session import get_db_read_replica
+
+
+def test_track_access(app):
+    with app.app_context():
+        db = get_db_read_replica()
+
+        user_entity_1 = {"user_id": 1}
+        user_entity_2 = {"user_id": 2}
+        user_entities = [user_entity_1, user_entity_2]
+
+        non_premium_track_entity = {
+            "track_id": 1,
+            "is_premium": False,
+            "premium_conditions": None,
+        }
+        premium_track_entity = {
+            "track_id": 2,
+            "is_premium": True,
+            "premium_conditions": {"nft-collection": "some-nft-collection"},
+        }
+        track_entities = [non_premium_track_entity, premium_track_entity]
+
+        entities = {"users": user_entities, "tracks": track_entities}
+
+        populate_mock_db(db, entities)
+
+        premium_content_access_checker = PremiumContentAccessChecker()
+
+        non_exisent_track_id = 3
+
+        result = premium_content_access_checker.check_access_for_batch(
+            [
+                {
+                    "user_id": user_entity_1["user_id"],
+                    "premium_content_id": non_exisent_track_id,
+                    "premium_content_type": "track",
+                },
+                {
+                    "user_id": user_entity_2["user_id"],
+                    "premium_content_id": non_premium_track_entity["track_id"],
+                    "premium_content_type": "track",
+                },
+                {
+                    "user_id": user_entity_2["user_id"],
+                    "premium_content_id": premium_track_entity["track_id"],
+                    "premium_content_type": "track",
+                },
+            ]
+        )
+
+        track_access_result = result["track"]
+
+        # test non-existent track
+        assert user_entity_1["user_id"] not in track_access_result
+
+        # test non-premium track
+        user_2_non_premium_track_access_result = track_access_result[
+            user_entity_2["user_id"]
+        ][non_premium_track_entity["track_id"]]
+        assert (
+            not user_2_non_premium_track_access_result["is_premium"]
+            and user_2_non_premium_track_access_result["does_user_have_access"]
+        )
+
+        # test premium track with user who has access
+        user_2_premium_track_access_result = track_access_result[
+            user_entity_2["user_id"]
+        ][premium_track_entity["track_id"]]
+        assert (
+            user_2_premium_track_access_result["is_premium"]
+            and user_2_premium_track_access_result["does_user_have_access"]
+        )
+
+        # todo: test premium track with user who has no access
+        # after we implement nft infexing

discovery-provider/integration_tests/tasks/test_index_operations.py

@@ -105,6 +105,8 @@ def seed_contract_data(task, contracts, web3):
         "release_date": str(chance.date()),
         "file_type": "mp3",
         "track_segments": test_track_segments,
+        "is_premium": False,
+        "premium_conditions": None,
     }
 
     # dump metadata to file

discovery-provider/integration_tests/tasks/test_index_tracks.py

@@ -138,6 +138,8 @@ def get_delete_track_event():
             },
             "track_id": 77955,
             "stem_of": None,
+            "is_premium": False,
+            "premium_conditions": None,
         },
         multihash2: {
             "owner_id": 1,
@@ -185,6 +187,8 @@ def get_delete_track_event():
             },
             "track_id": 77955,
             "stem_of": None,
+            "is_premium": False,
+            "premium_conditions": None,
         },
     }
 )

discovery-provider/integration_tests/utils.py

@@ -171,6 +171,8 @@ def populate_mock_db(db, entities, block_offset=None):
                 created_at=track_meta.get("created_at", datetime.now()),
                 release_date=track_meta.get("release_date", None),
                 is_unlisted=track_meta.get("is_unlisted", False),
+                is_premium=track_meta.get("is_premium", False),
+                premium_conditions=track_meta.get("premium_conditions", None),
             )
             session.add(track)
         for i, playlist_meta in enumerate(playlists):

discovery-provider/src/models/tracks/track.py

@@ -64,6 +64,8 @@ class Track(Base, RepresentableMixin):
     )
     slot = Column(Integer)
     is_available = Column(Boolean, nullable=False, server_default=text("true"))
+    is_premium = Column(Boolean, nullable=False, server_default=text("false"))
+    premium_conditions = Column(JSONB())
 
     block = relationship(  # type: ignore
         "Block", primaryjoin="Track.blockhash == Block.blockhash"

discovery-provider/src/premium_content/helpers.py

@@ -0,0 +1,13 @@
+import logging
+
+# from src.utils import db_session
+
+logger = logging.getLogger(__name__)
+
+
+def does_user_have_nft_collection(user_id: int, nft_collection: str):
+    # todo: check whether user has the nft from some user_wallet_nfts table
+    # which is populated during nft indexing
+    # db = db_session.get_db_read_replica()
+    # with db.scoped_session() as session:
+    return True

discovery-provider/src/premium_content/premium_content_access_checker.py

@@ -0,0 +1,154 @@
+import json
+import logging
+from typing import Dict, List, TypedDict, cast
+
+from src.models.tracks.track import Track
+from src.premium_content.helpers import does_user_have_nft_collection
+from src.premium_content.types import PremiumContentType
+from src.utils import db_session, helpers
+
+logger = logging.getLogger(__name__)
+
+
+class PremiumContentAccessArgs(TypedDict):
+    user_id: int
+    premium_content_id: int
+    premium_content_type: PremiumContentType
+
+
+class PremiumContentAccess(TypedDict):
+    is_premium: bool
+    does_user_have_access: bool
+
+
+PremiumTrackAccessResult = Dict[int, Dict[int, PremiumContentAccess]]
+
+
+class PremiumContentAccessBatchResponse(TypedDict):
+    # track : user id -> track id -> access
+    track: PremiumTrackAccessResult
+
+
+class PremiumContentAccessChecker:
+    # Given a list of objects, each with a user id, premium content id, and premium content type,
+    # this method checks for access to the premium contents by the users.
+    #
+    # Returns a dictionary in the following format:
+    # {
+    #   <premium-content-type>: {
+    #     <user-id>: {
+    #       <track-id>: {
+    #         "is_premium": bool,
+    #         "does_user_have_access": bool
+    #       }
+    #     }
+    #   }
+    # }
+    def check_access_for_batch(
+        self, args: List[PremiumContentAccessArgs]
+    ) -> PremiumContentAccessBatchResponse:
+        # for now, we only allow tracks to be premium; premium playlists will come later
+        valid_args = list(
+            filter(lambda arg: arg["premium_content_type"] == "track", args)
+        )
+
+        if not valid_args:
+            return {"track": {}}
+
+        track_access_users = {
+            arg["premium_content_id"]: arg["user_id"] for arg in valid_args
+        }
+        premium_track_data = self._get_premium_track_data_for_batch(
+            list(track_access_users.keys())
+        )
+        track_access_result: PremiumTrackAccessResult = {}
+
+        for track_id, data in premium_track_data.items():
+            user_id = track_access_users[track_id]
+            if user_id not in track_access_result:
+                track_access_result[user_id] = {}
+
+            is_premium = data["is_premium"]
+            premium_conditions = data["premium_conditions"]
+            content_owner_id = data["content_owner_id"]
+
+            if not is_premium:
+                # premium_conditions should always be null here as it makes
+                # no sense to have a non-premium track with conditions
+                if premium_conditions:
+                    logger.warn(
+                        f"premium_content_access_checker.py | _aggregate_conditions | non-premium content with id {track_id} and type 'track' has premium conditions."
+                    )
+                track_access_result[user_id][track_id] = {
+                    "is_premium": False,
+                    "does_user_have_access": True,
+                }
+
+            # premium_conditions should always be true here because we know
+            # that is_premium is true if we get here and it makes no sense
+            # to have a premium track with no conditions
+            elif not premium_conditions:
+                logger.warn(
+                    f"premium_content_access_checker.py | _aggregate_conditions | premium content with id {track_id} and type 'track' has no premium conditions."
+                )
+                track_access_result[user_id][track_id] = {
+                    "is_premium": True,
+                    "does_user_have_access": True,
+                }
+
+            else:
+                track_access_result[user_id][track_id] = {
+                    "is_premium": True,
+                    "does_user_have_access": self._evaluate_conditions(
+                        user_id=user_id,
+                        premium_content_owner_id=cast(int, content_owner_id),
+                        premium_conditions=premium_conditions,
+                    ),
+                }
+
+        return {"track": track_access_result}
+
+    def _get_premium_track_data_for_batch(self, track_ids: List[int]):
+        db = db_session.get_db_read_replica()
+        with db.scoped_session() as session:
+            tracks = (
+                session.query(Track)
+                .filter(
+                    Track.track_id.in_(track_ids),
+                    Track.is_current == True,
+                    Track.is_delete == False,
+                )
+                .all()
+            )
+            tracks = list(map(helpers.model_to_dictionary, tracks))
+
+            return {
+                track["track_id"]: {
+                    "is_premium": track["is_premium"],
+                    "premium_conditions": track["premium_conditions"],
+                    "content_owner_id": track["owner_id"],
+                }
+                for track in tracks
+            }
+
+    # There will eventually be another step prior to this one where
+    # we aggregate multiple conditions and evaluate them altogether.
+    # For now, we only support one condition, which is the ownership
+    # of an NFT from a given collection.
+    def _evaluate_conditions(
+        self, user_id: int, premium_content_owner_id: int, premium_conditions: Dict
+    ):
+        if len(premium_conditions) != 1:
+            logging.info(
+                f"premium_content_access_checker.py | _aggregate_conditions | invalid conditions: {json.dumps(premium_conditions)}"
+            )
+            return False
+
+        condition, value = list(premium_conditions.items())[0]
+        if condition != "nft-collection":
+            return False
+
+        return does_user_have_nft_collection(user_id, value)
+
+
+premium_content_access_checker = PremiumContentAccessChecker()

discovery-provider/src/premium_content/signature.py

@@ -1,9 +1,8 @@
 from datetime import datetime
-from typing import Literal, TypedDict, cast
+from typing import TypedDict, cast
 
 from src.api_helpers import generate_signature
-
-PremiumContentType = Literal["track"]
+from src.premium_content.types import PremiumContentType
 
 
 class PremiumContentSignatureData(TypedDict):

discovery-provider/src/premium_content/types.py

@@ -0,0 +1,15 @@
+from typing import Literal
+
+# These are the different types of content that can be made premium.
+# For now, we only support tracks.
+PremiumContentType = Literal["track"]
+
+# These are the different conditions on which premium content access
+# will be gated.
+# They should match the PremiumConditions property in the track schema
+# in src/schemas/track_schema.json
+PremiumContentConditions = Literal["nft-collection"]
+
+# This is for when we support the combination of different conditions
+# for premium content access based on AND'ing / OR'ing them together.
+PremiumContentConditionGates = Literal["AND", "OR"]