From eb0e2be71377184eda786d8f7b81d91c61cb8c40 Mon Sep 17 00:00:00 2001
From: Maksym Pavlenko <pavlenko.maksym@gmail.com>
Date: Mon, 18 Mar 2019 22:57:43 -0700
Subject: [PATCH] Add resolve counter

---
 cmd/resolver/function.py      | 36 +++++++++++++++++++++++++++++++++++
 cmd/resolver/function_test.py |  6 ++++++
 2 files changed, 42 insertions(+)

diff --git a/cmd/resolver/function.py b/cmd/resolver/function.py
index e89e7af..85b4115 100644
--- a/cmd/resolver/function.py
+++ b/cmd/resolver/function.py
@@ -1,6 +1,8 @@
 import os
 import youtube_dl
 import boto3
+from datetime import datetime, time
+from dateutil.relativedelta import relativedelta
 
 
 class InvalidUsage(Exception):
@@ -10,6 +12,7 @@ class InvalidUsage(Exception):
 dynamodb = boto3.resource('dynamodb')
 
 feeds_table = dynamodb.Table(os.getenv('RESOLVER_DYNAMO_FEEDS_TABLE', 'Feeds'))
+counter_table = dynamodb.Table(os.getenv('RESOLVER_DYNAMO_RESOLVE_COUNTERS_TABLE', 'ResolveCounters'))
 
 opts = {
     'quiet': True,
@@ -31,6 +34,9 @@ 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 {
@@ -78,6 +84,36 @@ def _get_metadata(feed_id):
     return dict((k.lower(), v) for k, v in item.items())
 
 
+def _update_resolve_counter(feed_id):
+    if not feed_id:
+        return
+
+    now = datetime.utcnow()
+    day = now.strftime('%Y%m%d')
+
+    expires = now + relativedelta(months=3)
+
+    response = counter_table.update_item(
+        Key={
+            'FeedID': feed_id,
+            'Day': int(day),
+        },
+        UpdateExpression='ADD #count :one SET #exp = if_not_exists(#exp, :ttl)',
+        ExpressionAttributeNames={
+            '#count': 'Count',
+            '#exp': 'Expires',
+        },
+        ExpressionAttributeValues={
+            ':one': 1,
+            ':ttl': int(expires.timestamp()),
+        },
+        ReturnValues='UPDATED_NEW',
+    )
+
+    attrs = response['Attributes']
+    return attrs['Count']
+
+
 def _resolve(url, metadata):
     if not url:
         raise InvalidUsage('Invalid URL')
diff --git a/cmd/resolver/function_test.py b/cmd/resolver/function_test.py
index 5d5d607..a10b812 100644
--- a/cmd/resolver/function_test.py
+++ b/cmd/resolver/function_test.py
@@ -34,3 +34,9 @@ class TestDynamo(unittest.TestCase):
         self.assertIsNotNone(item['format'])
         self.assertIsNotNone(item['quality'])
         self.assertIsNotNone(item['provider'])
+
+    def test_counter(self):
+        counter = ytdl._update_resolve_counter('86qZ')
+        self.assertEqual(counter, 1)
+        counter = ytdl._update_resolve_counter('86qZ')
+        self.assertEqual(counter, 2)