From ef5c13f9682dec0230365cf2f60a6cda957d6d67 Mon Sep 17 00:00:00 2001
From: Debanjum Singh Solanky <debanjum@gmail.com>
Date: Wed, 8 Nov 2023 12:08:36 -0800
Subject: [PATCH] Keep user subscription state. Update it when user has
 unsubscribed

---
 src/database/adapters/__init__.py               | 11 +++++++++++
 .../migrations/0014_khojuser_is_subscribed.py   | 17 +++++++++++++++++
 src/database/models/__init__.py                 |  1 +
 src/khoj/routers/api.py                         |  8 ++++++--
 4 files changed, 35 insertions(+), 2 deletions(-)
 create mode 100644 src/database/migrations/0014_khojuser_is_subscribed.py

diff --git a/src/database/adapters/__init__.py b/src/database/adapters/__init__.py
index b4186d57..3358aa79 100644
--- a/src/database/adapters/__init__.py
+++ b/src/database/adapters/__init__.py
@@ -103,10 +103,21 @@ async def create_google_user(token: dict) -> KhojUser:
     return user
 
 
+async def set_user_unsubscribed(email: str, type="standard") -> KhojUser:
+    user = await KhojUser.objects.filter(email=email, subscription_type=type).afirst()
+    if user:
+        user.is_subscribed = False
+        await user.asave()
+        return user
+    else:
+        return None
+
+
 async def set_user_subscribed(email: str, type="standard") -> KhojUser:
     user = await KhojUser.objects.filter(email=email).afirst()
     if user:
         user.subscription_type = type
+        user.is_subscribed = True
         start_date = user.subscription_renewal_date or datetime.now()
         user.subscription_renewal_date = start_date + timedelta(days=30)
         await user.asave()
diff --git a/src/database/migrations/0014_khojuser_is_subscribed.py b/src/database/migrations/0014_khojuser_is_subscribed.py
new file mode 100644
index 00000000..79b035ba
--- /dev/null
+++ b/src/database/migrations/0014_khojuser_is_subscribed.py
@@ -0,0 +1,17 @@
+# Generated by Django 4.2.5 on 2023-11-08 19:40
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("database", "0013_khojuser_subscription_renewal_date_and_more"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="khojuser",
+            name="is_subscribed",
+            field=models.BooleanField(default=False),
+        ),
+    ]
diff --git a/src/database/models/__init__.py b/src/database/models/__init__.py
index ea982d06..5a3b96c5 100644
--- a/src/database/models/__init__.py
+++ b/src/database/models/__init__.py
@@ -22,6 +22,7 @@ class KhojUser(AbstractUser):
     subscription_type = models.CharField(
         max_length=20, choices=SubscriptionType.choices, default=SubscriptionType.TRIAL
     )
+    is_subscribed = models.BooleanField(default=False)
     subscription_renewal_date = models.DateTimeField(null=True, default=None)
 
     def save(self, *args, **kwargs):
diff --git a/src/khoj/routers/api.py b/src/khoj/routers/api.py
index c4250f82..77e377e7 100644
--- a/src/khoj/routers/api.py
+++ b/src/khoj/routers/api.py
@@ -765,12 +765,16 @@ async def subscribe(request: Request):
         # Retrieve the customer's details
         customer_id = event["data"]["object"]["customer"]
         customer = stripe.Customer.retrieve(customer_id)
+        customer_email = customer["email"]
+        # Mark the customer as unsubscribed
+        user = await adapters.set_user_unsubscribed(customer_email)
+        if not user:
+            success = False
     else:
-        logger.warn(f"Unhandled Stripe event type: {event['type']}, {event['data']['object']}")
+        logger.warn(f"Unhandled Stripe event type: {event['type']}")
         return {"success": False}
 
     logger.info(f'Stripe subscription {event["type"]} for {customer["email"]}')
-
     return {"success": success}