Merge branch 'master' into glossary_new

Procfile

@@ -17,6 +17,6 @@ admin: poetry run python manage.py runserver 127.0.0.1:8000
 
 dashboard: poetry run streamlit run Home.py --server.port 8501 --server.headless true
 
-celery: poetry run celery -A celeryapp worker
+celery: poetry run celery -A celeryapp worker -P threads -c 16 -l DEBUG
 
 ui: cd ../gooey-ui/; PORT=3000 npm run dev

app_users/admin.py

@@ -1,7 +1,8 @@
 from django.contrib import admin
 
 from app_users import models
-from bots.admin_links import open_in_new_tab
+from bots.admin_links import open_in_new_tab, list_related_html_url
+from bots.models import SavedRun
 
 
 # Register your models here.
@@ -15,6 +16,7 @@ class AppUserAdmin(admin.ModelAdmin):
         "email",
         "phone_number",
         "balance",
+        "is_paying",
         "is_anonymous",
         "is_disabled",
         "created_at",
@@ -30,12 +32,14 @@ class AppUserAdmin(admin.ModelAdmin):
     list_filter = [
         "is_anonymous",
         "is_disabled",
+        "is_paying",
         "created_at",
         "upgraded_from_anonymous_at",
     ]
     readonly_fields = [
         "created_at",
         "upgraded_from_anonymous_at",
+        "view_transactions",
         "open_in_firebase",
         "open_in_stripe",
     ]
@@ -62,3 +66,65 @@ def open_in_stripe(self, user: models.AppUser):
         )
 
     open_in_stripe.short_description = "Open in Stripe"
+
+    @admin.display(description="View transactions")
+    def view_transactions(self, user: models.AppUser):
+        return list_related_html_url(user.transactions, show_add=False)
+
+
+class SavedRunInline(admin.StackedInline):
+    model = SavedRun
+    extra = 0
+    fields = readonly_fields = [
+        "open_in_gooey",
+        "price",
+        "created_at",
+        "updated_at",
+        "run_time",
+    ]
+    show_change_link = True
+
+    def has_change_permission(self, request, obj=None):
+        return False
+
+    def has_add_permission(self, request, obj):
+        return False
+
+    def has_delete_permission(self, request, obj=None):
+        return False
+
+
+class IsStripeFilter(admin.SimpleListFilter):
+    title = "Is Stripe Invoice"
+    parameter_name = "is_stripe_invoice"
+
+    def lookups(self, request, model_admin):
+        return (
+            ("1", "Yes"),
+            ("0", "No"),
+        )
+
+    def queryset(self, request, queryset):
+        value = self.value()
+        if value is None:
+            return queryset
+        if int(value):
+            return queryset.filter(invoice_id__startswith="in_")
+        else:
+            return queryset.exclude(invoice_id__startswith="in_")
+
+
+@admin.register(models.AppUserTransaction)
+class AppUserTransactionAdmin(admin.ModelAdmin):
+    autocomplete_fields = ["user"]
+    list_display = [
+        "invoice_id",
+        "user",
+        "amount",
+        "created_at",
+        "end_balance",
+    ]
+    readonly_fields = ["created_at"]
+    list_filter = ["created_at", IsStripeFilter]
+    inlines = [SavedRunInline]
+    ordering = ["-created_at"]

app_users/apps.py

@@ -4,6 +4,7 @@
 class AppUsersConfig(AppConfig):
     default_auto_field = "django.db.models.BigAutoField"
     name = "app_users"
+    verbose_name = "App Users"
 
     def ready(self):
         from . import signals

app_users/migrations/0007_appuser_is_paying.py

@@ -0,0 +1,17 @@
+# Generated by Django 4.2.5 on 2023-10-12 13:54
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("app_users", "0006_appuser_disable_safety_checker"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="appuser",
+            name="is_paying",
+            field=models.BooleanField(default=False),
+        ),
+    ]

app_users/migrations/0008_appusertransaction.py

@@ -0,0 +1,48 @@
+# Generated by Django 4.2.5 on 2023-10-13 10:55
+
+from django.db import migrations, models
+import django.db.models.deletion
+import django.utils.timezone
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("app_users", "0007_appuser_is_paying"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="AppUserTransaction",
+            fields=[
+                (
+                    "id",
+                    models.BigAutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("invoice_id", models.CharField(max_length=255)),
+                ("amount", models.IntegerField()),
+                ("end_balance", models.IntegerField()),
+                (
+                    "created_at",
+                    models.DateTimeField(
+                        blank=True, default=django.utils.timezone.now, editable=False
+                    ),
+                ),
+                (
+                    "user",
+                    models.ForeignKey(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        related_name="transactions",
+                        to="app_users.appuser",
+                    ),
+                ),
+            ],
+            options={
+                "unique_together": {("user", "invoice_id")},
+            },
+        ),
+    ]

app_users/migrations/0009_alter_appusertransaction_options_and_more.py

@@ -0,0 +1,25 @@
+# Generated by Django 4.2.5 on 2023-10-18 14:33
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("app_users", "0008_appusertransaction"),
+    ]
+
+    operations = [
+        migrations.AlterModelOptions(
+            name="appusertransaction",
+            options={"verbose_name": "Transaction"},
+        ),
+        migrations.AlterUniqueTogether(
+            name="appusertransaction",
+            unique_together=set(),
+        ),
+        migrations.AlterField(
+            model_name="appusertransaction",
+            name="invoice_id",
+            field=models.CharField(max_length=255, unique=True),
+        ),
+    ]

app_users/models.py

@@ -1,5 +1,3 @@
-import datetime
-
 import requests
 import stripe
 from django.db import models, IntegrityError, transaction
@@ -89,6 +87,7 @@ class AppUser(models.Model):
     photo_url = CustomURLField(default="", blank=True)
 
     stripe_customer_id = models.CharField(max_length=255, default="", blank=True)
+    is_paying = models.BooleanField(default=False)
 
     created_at = models.DateTimeField(editable=False, blank=True, default=timezone.now)
     upgraded_from_anonymous_at = models.DateTimeField(null=True, blank=True)
@@ -105,43 +104,40 @@ def first_name(self):
             return ""
         return self.display_name.split(" ")[0]
 
-    def add_balance(self, amount: int, invoice_id: str, **invoice_items):
-        from google.cloud import firestore
-        from google.cloud.firestore_v1.transaction import Transaction
-
-        @firestore.transactional
-        @db_middleware
-        def _update_user_balance_in_txn(txn: Transaction):
-            user_doc_ref = db.get_user_doc_ref(self.uid)
-
-            invoice_ref: firestore.DocumentReference
-            invoice_ref = user_doc_ref.collection("invoices").document(invoice_id)
-            # if an invoice entry exists
-            if invoice_ref.get(transaction=txn).exists:
-                # avoid updating twice for same invoice
-                return
-
-            obj = self.add_balance_direct(amount)
-
-            # create invoice entry
-            txn.create(
-                invoice_ref,
-                {
-                    "amount": amount,
-                    "end_balance": obj.balance,
-                    "timestamp": datetime.datetime.utcnow(),
-                    **invoice_items,
-                },
-            )
+    @db_middleware
+    @transaction.atomic
+    def add_balance(self, amount: int, invoice_id: str) -> "AppUserTransaction":
+        """
+        Used to add/deduct credits when they are bought or consumed.
 
-        _update_user_balance_in_txn(db.get_client().transaction())
+        When credits are bought with stripe -- invoice_id is the stripe
+        invoice ID.
+        When credits are deducted due to a run -- invoice_id is of the
+        form "gooey_in_{uuid}"
+        """
 
-    @transaction.atomic
-    def add_balance_direct(self, amount):
-        obj: AppUser = self.__class__.objects.select_for_update().get(pk=self.pk)
-        obj.balance += amount
-        obj.save(update_fields=["balance"])
-        return obj
+        # if an invoice entry exists
+        try:
+            # avoid updating twice for same invoice
+            return AppUserTransaction.objects.get(invoice_id=invoice_id)
+        except AppUserTransaction.DoesNotExist:
+            pass
+
+        # select_for_update() is very important here
+        # transaction.atomic alone is not enough!
+        # It won't lock this row for reads, and multiple threads can update the same row leading incorrect balance
+        #
+        # Also we're not using .update() here because it won't give back the updated end balance
+        user: AppUser = AppUser.objects.select_for_update().get(pk=self.pk)
+        user.balance += amount
+        user.save(update_fields=["balance"])
+
+        return AppUserTransaction.objects.create(
+            user=self,
+            invoice_id=invoice_id,
+            amount=amount,
+            end_balance=user.balance,
+        )
 
     def copy_from_firebase_user(self, user: auth.UserRecord) -> "AppUser":
         # copy data from firebase user
@@ -150,6 +146,7 @@ def copy_from_firebase_user(self, user: auth.UserRecord) -> "AppUser":
         self.display_name = user.display_name or ""
         self.email = user.email
         self.phone_number = user.phone_number
+        provider_list = user.provider_data
         self.created_at = timezone.datetime.fromtimestamp(
             user.user_metadata.creation_timestamp / 1000
         )
@@ -175,14 +172,15 @@ def copy_from_firebase_user(self, user: auth.UserRecord) -> "AppUser":
         self.is_anonymous = is_anonymous_now
 
         # get existing balance or set free credits
+        default_balance = settings.LOGIN_USER_FREE_CREDITS
+        if self.is_anonymous:
+            default_balance = settings.ANON_USER_FREE_CREDITS
+        elif provider_list[-1].provider_id == "password":
+            default_balance = settings.EMAIL_USER_FREE_CREDITS
         self.balance = db.get_doc_field(
             doc_ref=db.get_user_doc_ref(user.uid),
             field=db.USER_BALANCE_FIELD,
-            default=(
-                settings.ANON_USER_FREE_CREDITS
-                if self.is_anonymous
-                else settings.LOGIN_USER_FREE_CREDITS
-            ),
+            default=default_balance,
         )
 
         return self
@@ -215,3 +213,19 @@ def search_stripe_customer(self) -> stripe.Customer | None:
             self.stripe_customer_id = customer.id
             self.save()
             return customer
+
+
+class AppUserTransaction(models.Model):
+    user = models.ForeignKey(
+        "AppUser", on_delete=models.CASCADE, related_name="transactions"
+    )
+    invoice_id = models.CharField(max_length=255, unique=True)
+    amount = models.IntegerField()
+    end_balance = models.IntegerField()
+    created_at = models.DateTimeField(editable=False, blank=True, default=timezone.now)
+
+    class Meta:
+        verbose_name = "Transaction"
+
+    def __str__(self):
+        return f"{self.invoice_id} ({self.amount})"