Add resolver quota

cmd/resolver/function.py

@@ -4,11 +4,17 @@ import boto3
 from datetime import datetime, time
 from dateutil.relativedelta import relativedelta
 
+ANONYMOUS_FEED_REQUESTS_LIMIT = 10000
+
 
 class InvalidUsage(Exception):
     pass
 
 
+class QuotaExceeded(Exception):
+    pass
+
+
 dynamodb = boto3.resource('dynamodb')
 
 feeds_table = dynamodb.Table(os.getenv('RESOLVER_DYNAMO_FEEDS_TABLE', 'Feeds'))
@@ -34,9 +40,6 @@ def handler(event, context):
     feed_id = event['feed_id']
     video_id = event['video_id']
 
-    # Update resolve requests counter
-    _update_resolve_counter(feed_id)
-
     redirect_url = download(feed_id, video_id)
 
     return {
@@ -56,6 +59,13 @@ def download(feed_id, video_id):
     # Query feed metadata info from DynamoDB
     item = _get_metadata(feed_id)
 
+    # Update resolve requests counter
+    count = _update_resolve_counter(feed_id)
+    level = int(item['featurelevel'])
+    if count > ANONYMOUS_FEED_REQUESTS_LIMIT and level == 0:
+        raise QuotaExceeded('Too many requests. Daily limit is %d. Consider upgrading account to get unlimited '
+                            'access' % ANONYMOUS_FEED_REQUESTS_LIMIT)
+
     # Build URL
     provider = item['provider']
     tpl = url_formats[provider]
@@ -70,11 +80,12 @@ def download(feed_id, video_id):
 def _get_metadata(feed_id):
     response = feeds_table.get_item(
         Key={'HashID': feed_id},
-        ProjectionExpression='#P,#F,#Q',
+        ProjectionExpression='#P,#F,#Q,#L',
         ExpressionAttributeNames={
             '#P': 'Provider',
             '#F': 'Format',
             '#Q': 'Quality',
+            '#L': 'FeatureLevel',
         },
     )
 

cmd/resolver/function_test.py

@@ -40,3 +40,11 @@ class TestDynamo(unittest.TestCase):
         self.assertEqual(counter, 1)
         counter = ytdl._update_resolve_counter('86qZ')
         self.assertEqual(counter, 2)
+
+    def test_download(self):
+        url = ytdl.download('86qZ', '7XJYLq3gviY')
+        self.assertIsNotNone(url)
+
+    def test_quota_check(self):
+        with self.assertRaises(ytdl.QuotaExceeded):
+            ytdl.download('xro548QlJ', 'j51NFs0bZ9c')