From 9aaf475c8a6d707ac19091caf60edef2808cf999 Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Tue, 7 Nov 2023 10:15:21 -0800 Subject: [PATCH] Create API webhook, endpoints for subscription payments using Stripe - Add fields to mark users as subscribed to a specific plan and subscription renewal date in DB - Add ability to unsubscribe a user using their email address - Expose webhook for stripe to callback confirming payment --- pyproject.toml | 3 +- src/database/adapters/__init__.py | 23 ++++++- ...user_subscription_renewal_date_and_more.py | 24 ++++++++ src/database/models/__init__.py | 8 +++ src/khoj/routers/api.py | 61 ++++++++++++++++++- 5 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 src/database/migrations/0013_khojuser_subscription_renewal_date_and_more.py diff --git a/pyproject.toml b/pyproject.toml index f6080ce6..10e44ac0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,8 @@ dependencies = [ "gunicorn == 21.2.0", "lxml == 4.9.3", "tzdata == 2023.3", - "rapidocr-onnxruntime == 1.3.8" + "rapidocr-onnxruntime == 1.3.8", + "stripe == 7.3.0", ] dynamic = ["version"] diff --git a/src/database/adapters/__init__.py b/src/database/adapters/__init__.py index 69a3c1f4..f4e816bf 100644 --- a/src/database/adapters/__init__.py +++ b/src/database/adapters/__init__.py @@ -1,5 +1,5 @@ from typing import Type, TypeVar, List -from datetime import date +from datetime import date, datetime, timedelta import secrets from typing import Type, TypeVar, List from datetime import date @@ -103,6 +103,27 @@ async def create_google_user(token: dict) -> KhojUser: return user +async def set_user_subscribed(email: str, type="standard") -> KhojUser: + user = await KhojUser.objects.filter(email=email).afirst() + if user: + user.subscription_type = type + start_date = user.subscription_renewal_date or datetime.now() + user.subscription_renewal_date = start_date + timedelta(days=30) + await user.asave() + return user + else: + return None + + +def is_user_subscribed(email: str, type="standard") -> bool: + user = KhojUser.objects.filter(email=email, subscription_type=type).first() + if user and user.subscription_renewal_date: + is_subscribed = user.subscription_renewal_date > date.today() + return is_subscribed + else: + return False + + async def get_user_by_token(token: dict) -> KhojUser: google_user = await GoogleUser.objects.filter(sub=token.get("sub")).select_related("user").afirst() if not google_user: diff --git a/src/database/migrations/0013_khojuser_subscription_renewal_date_and_more.py b/src/database/migrations/0013_khojuser_subscription_renewal_date_and_more.py new file mode 100644 index 00000000..d7f3df5c --- /dev/null +++ b/src/database/migrations/0013_khojuser_subscription_renewal_date_and_more.py @@ -0,0 +1,24 @@ +# Generated by Django 4.2.5 on 2023-11-07 18:19 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("database", "0012_entry_file_source"), + ] + + operations = [ + migrations.AddField( + model_name="khojuser", + name="subscription_renewal_date", + field=models.DateTimeField(default=None, null=True), + ), + migrations.AddField( + model_name="khojuser", + name="subscription_type", + field=models.CharField( + choices=[("trial", "Trial"), ("standard", "Standard")], default="trial", max_length=20 + ), + ), + ] diff --git a/src/database/models/__init__.py b/src/database/models/__init__.py index b1be9ded..ea982d06 100644 --- a/src/database/models/__init__.py +++ b/src/database/models/__init__.py @@ -14,7 +14,15 @@ class BaseModel(models.Model): class KhojUser(AbstractUser): + class SubscriptionType(models.TextChoices): + TRIAL = "trial" + STANDARD = "standard" + uuid = models.UUIDField(models.UUIDField(default=uuid.uuid4, editable=False)) + subscription_type = models.CharField( + max_length=20, choices=SubscriptionType.choices, default=SubscriptionType.TRIAL + ) + subscription_renewal_date = models.DateTimeField(null=True, default=None) def save(self, *args, **kwargs): if not self.uuid: diff --git a/src/khoj/routers/api.py b/src/khoj/routers/api.py index fabfebe1..d4176d06 100644 --- a/src/khoj/routers/api.py +++ b/src/khoj/routers/api.py @@ -1,6 +1,7 @@ # Standard Packages import concurrent.futures import math +import os import time import logging import json @@ -10,6 +11,7 @@ from typing import List, Optional, Union, Any from fastapi import APIRouter, HTTPException, Header, Request from starlette.authentication import requires from asgiref.sync import sync_to_async +import stripe # Internal Packages from khoj.configure import configure_server @@ -23,7 +25,6 @@ from khoj.utils.rawconfig import ( FullConfig, SearchConfig, SearchResponse, - TextContentConfig, GithubContentConfig, NotionContentConfig, ) @@ -723,3 +724,61 @@ async def extract_references_and_questions( compiled_references = [item.additional["compiled"] for item in result_list] return compiled_references, inferred_queries, defiltered_query + + +# Stripe integration for Khoj Cloud Subscription +stripe.api_key = os.getenv("STRIPE_API_KEY") +endpoint_secret = os.getenv("STRIPE_SIGINING_SECRET") + + +@api.post("/subscription") +async def subscribe(request: Request): + """Webhook for Stripe to send subscription events to Khoj Cloud""" + event = None + try: + payload = await request.body() + sig_header = request.headers["stripe-signature"] + event = stripe.Webhook.construct_event(payload, sig_header, endpoint_secret) + except ValueError as e: + # Invalid payload + raise e + except stripe.error.SignatureVerificationError as e: + # Invalid signature + raise e + + # Handle the event + success = True + if ( + event["type"] == "payment_intent.succeeded" + or event["type"] == "invoice.payment_succeeded" + or event["type"] == "customer.subscription.created" + ): + # 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 subscribed + user = await adapters.set_user_subscribed(customer_email) + if not user: + success = False + elif event["type"] == "customer.subscription.updated" or event["type"] == "customer.subscription.deleted": + # Retrieve the customer's details + customer_id = event["data"]["object"]["customer"] + customer = stripe.Customer.retrieve(customer_id) + + logger.info(f'Stripe subscription {event["type"]} for {customer["email"]}') + + return {"success": success} + + +@api.delete("/subscription") +@requires(["authenticated"]) +async def unsubscribe(request: Request, user_email: str): + customer = stripe.Customer.list(email=user_email).data + if not is_none_or_empty(customer): + stripe.Subscription.modify(customer[0].id, cancel_at_period_end=True) + success = True + else: + success = False + + return {"success": success}