server: add a Sale model and migrate Transaction data with it

server/migrations/versions/2024-05-15-1728_create_sale.py

@@ -0,0 +1,266 @@
+"""Create Sale
+
+Revision ID: e553ec9cf929
+Revises: 40f397ad512a
+Create Date: 2024-05-15 17:28:54.066418
+
+"""
+
+import uu
+import uuid
+from typing import cast
+
+import sqlalchemy as sa
+import stripe
+from alembic import op
+
+from polar.integrations.stripe.service import stripe as stripe_service
+
+# Polar Custom Imports
+from polar.kit.extensions.sqlalchemy import PostgresUUID
+
+# revision identifiers, used by Alembic.
+revision = "e553ec9cf929"
+down_revision = "40f397ad512a"
+branch_labels: tuple[str] | None = None
+depends_on: tuple[str] | None = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "sales",
+        sa.Column("id", sa.UUID(), nullable=False),
+        sa.Column("created_at", sa.TIMESTAMP(timezone=True), nullable=False),
+        sa.Column("modified_at", sa.TIMESTAMP(timezone=True), nullable=True),
+        sa.Column("deleted_at", sa.TIMESTAMP(timezone=True), nullable=True),
+        sa.Column("amount", sa.Integer(), nullable=False),
+        sa.Column("tax_amount", sa.Integer(), nullable=False),
+        sa.Column("currency", sa.String(length=3), nullable=False),
+        sa.Column("stripe_invoice_id", sa.String(), nullable=True),
+        sa.Column("user_id", sa.UUID(), nullable=False),
+        sa.Column("product_id", sa.UUID(), nullable=False),
+        sa.Column("product_price_id", sa.UUID(), nullable=False),
+        sa.Column("subscription_id", sa.UUID(), nullable=True),
+        sa.ForeignKeyConstraint(
+            ["product_id"], ["products.id"], name=op.f("sales_product_id_fkey")
+        ),
+        sa.ForeignKeyConstraint(
+            ["product_price_id"],
+            ["product_prices.id"],
+            name=op.f("sales_product_price_id_fkey"),
+        ),
+        sa.ForeignKeyConstraint(
+            ["subscription_id"],
+            ["subscriptions.id"],
+            name=op.f("sales_subscription_id_fkey"),
+        ),
+        sa.ForeignKeyConstraint(
+            ["user_id"], ["users.id"], name=op.f("sales_user_id_fkey")
+        ),
+        sa.PrimaryKeyConstraint("id", name=op.f("sales_pkey")),
+    )
+    op.create_index(op.f("ix_sales_created_at"), "sales", ["created_at"], unique=False)
+    op.create_index(op.f("ix_sales_deleted_at"), "sales", ["deleted_at"], unique=False)
+    op.create_index(
+        op.f("ix_sales_modified_at"), "sales", ["modified_at"], unique=False
+    )
+    op.add_column("transactions", sa.Column("sale_id", sa.UUID(), nullable=True))
+
+    op.create_foreign_key(
+        op.f("transactions_sale_id_fkey"),
+        "transactions",
+        "sales",
+        ["sale_id"],
+        ["id"],
+        ondelete="set null",
+    )
+
+    # Create Sales from Transactions data
+    connection = op.get_bind()
+    subscription_payments = connection.execute(
+        sa.text(
+            """
+            SELECT transactions.id, transactions.created_at, transactions.amount, transactions.tax_amount, transactions.currency, subscriptions.id, subscriptions.user_id, subscriptions.stripe_subscription_id, subscriptions.product_id, transactions.product_price_id, transactions.charge_id
+            FROM transactions
+            JOIN subscriptions ON transactions.subscription_id = subscriptions.id
+            WHERE transactions.type = 'payment';
+            """
+        )
+    )
+    for (
+        transaction_id,
+        created_at,
+        amount,
+        tax_amount,
+        currency,
+        subscription_id,
+        subscription_user_id,
+        stripe_subscription_id,
+        product_id,
+        product_price_id,
+        charge_id,
+    ) in subscription_payments:
+        stripe_charge = stripe_service.get_charge(charge_id, expand=["invoice"])
+        assert (
+            stripe_charge.amount == amount + tax_amount
+        ), f"Amounts do not match for transaction {transaction_id} and charge {charge_id}"
+        assert (
+            stripe_charge.invoice is not None
+        ), f"No invoice found for charge {charge_id}"
+        stripe_invoice = cast(stripe.Invoice, stripe_charge.invoice)
+        assert (
+            stripe_invoice.subscription == stripe_subscription_id
+        ), f"Subscription IDs do not match for charge {charge_id} and subscription {stripe_subscription_id}"
+
+        sales_id = uuid.uuid4()
+
+        # Create Sale
+        connection.execute(
+            sa.text(
+                """
+                INSERT INTO sales (id, created_at, amount, tax_amount, currency, stripe_invoice_id, user_id, product_id, product_price_id, subscription_id)
+                VALUES (
+                        :id,
+                        :created_at,
+                        :amount,
+                        :tax_amount,
+                        :currency,
+                        :stripe_invoice_id,
+                        :user_id,
+                        :product_id,
+                        :product_price_id,
+                        :subscription_id
+                );
+                """
+            ),
+            {
+                "id": sales_id,
+                "created_at": created_at,
+                "amount": amount,
+                "tax_amount": tax_amount,
+                "currency": currency,
+                "stripe_invoice_id": stripe_invoice.id,
+                "user_id": subscription_user_id,
+                "product_id": product_id,
+                "product_price_id": product_price_id,
+                "subscription_id": subscription_id,
+            },
+        )
+
+        # Update the payment transaction
+        connection.execute(
+            sa.text(
+                """
+                UPDATE transactions
+                SET sale_id = :sale_id
+                WHERE id = :transaction_id;
+                """
+            ),
+            {"sale_id": sales_id, "transaction_id": transaction_id},
+        )
+
+        # Update the balance transactions
+        connection.execute(
+            sa.text(
+                """
+                UPDATE transactions
+                SET sale_id = :sale_id
+                WHERE payment_transaction_id = :transaction_id;
+                """
+            ),
+            {"sale_id": sales_id, "transaction_id": transaction_id},
+        )
+
+        # Update the reversal balance transactions
+        connection.execute(
+            sa.text(
+                """
+                UPDATE transactions
+                SET sale_id = :sale_id
+                WHERE balance_reversal_transaction_id IN (SELECT id FROM transactions WHERE payment_transaction_id = :transaction_id);
+                """
+            ),
+            {"sale_id": sales_id, "transaction_id": transaction_id},
+        )
+
+    op.drop_index("ix_transactions_product_price_id", table_name="transactions")
+    op.drop_index("ix_transactions_subscription_id", table_name="transactions")
+    op.create_index(
+        op.f("ix_transactions_sale_id"), "transactions", ["sale_id"], unique=False
+    )
+    op.drop_constraint(
+        "transactions_subscription_tier_price_id_fkey",
+        "transactions",
+        type_="foreignkey",
+    )
+    op.drop_constraint(
+        "transactions_subscription_id_fkey", "transactions", type_="foreignkey"
+    )
+    op.drop_column("transactions", "subscription_id")
+    op.drop_column("transactions", "product_price_id")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "transactions",
+        sa.Column("product_price_id", sa.UUID(), autoincrement=False, nullable=True),
+    )
+    op.add_column(
+        "transactions",
+        sa.Column("subscription_id", sa.UUID(), autoincrement=False, nullable=True),
+    )
+    op.create_foreign_key(
+        "transactions_subscription_id_fkey",
+        "transactions",
+        "subscriptions",
+        ["subscription_id"],
+        ["id"],
+        ondelete="SET NULL",
+    )
+    op.create_foreign_key(
+        "transactions_subscription_tier_price_id_fkey",
+        "transactions",
+        "product_prices",
+        ["product_price_id"],
+        ["id"],
+        ondelete="SET NULL",
+    )
+
+    # Set subscription_id and product_price_id from Sale
+    op.execute(
+        """
+        UPDATE transactions
+        SET
+            subscription_id = sales.subscription_id,
+            product_price_id = sales.product_price_id
+        FROM sales
+        WHERE transactions.sale_id = sales.id;
+        """
+    )
+
+    op.drop_constraint(
+        op.f("transactions_sale_id_fkey"), "transactions", type_="foreignkey"
+    )
+
+    op.drop_index(op.f("ix_transactions_sale_id"), table_name="transactions")
+    op.create_index(
+        "ix_transactions_subscription_id",
+        "transactions",
+        ["subscription_id"],
+        unique=False,
+    )
+    op.create_index(
+        "ix_transactions_product_price_id",
+        "transactions",
+        ["product_price_id"],
+        unique=False,
+    )
+    op.drop_column("transactions", "sale_id")
+    op.drop_index(op.f("ix_sales_modified_at"), table_name="sales")
+    op.drop_index(op.f("ix_sales_deleted_at"), table_name="sales")
+    op.drop_index(op.f("ix_sales_created_at"), table_name="sales")
+    op.drop_table("sales")
+    # ### end Alembic commands ###

server/polar/models/__init__.py

@@ -28,6 +28,7 @@
 from .product_price import ProductPrice
 from .pull_request import PullRequest
 from .repository import Repository
+from .sale import Sale
 from .subscription import Subscription
 from .traffic import Traffic
 from .transaction import Transaction
@@ -67,6 +68,7 @@
     "PledgeTransaction",
     "PullRequest",
     "Repository",
+    "Sale",
     "Subscription",
     "Product",
     "ProductBenefit",

server/polar/models/sale.py

@@ -0,0 +1,60 @@
+from typing import TYPE_CHECKING
+from uuid import UUID
+
+from sqlalchemy import ForeignKey, Integer, String
+from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship
+
+from polar.kit.db.models import RecordModel
+from polar.kit.extensions.sqlalchemy import PostgresUUID
+
+if TYPE_CHECKING:
+    from polar.models import Product, ProductPrice, Subscription, User
+
+
+class Sale(RecordModel):
+    __tablename__ = "sales"
+
+    amount: Mapped[int] = mapped_column(Integer, nullable=False)
+    tax_amount: Mapped[int] = mapped_column(Integer, nullable=False)
+    currency: Mapped[str] = mapped_column(String(3), nullable=False)
+    stripe_invoice_id: Mapped[str | None] = mapped_column(String, nullable=True)
+
+    user_id: Mapped[UUID] = mapped_column(
+        PostgresUUID,
+        ForeignKey("users.id"),
+        nullable=False,
+    )
+
+    @declared_attr
+    def user(cls) -> Mapped["User"]:
+        return relationship("User", lazy="raise")
+
+    product_id: Mapped[UUID] = mapped_column(
+        PostgresUUID,
+        ForeignKey("products.id"),
+        nullable=False,
+    )
+
+    @declared_attr
+    def product(cls) -> Mapped["Product"]:
+        return relationship("Product", lazy="raise")
+
+    product_price_id: Mapped[UUID] = mapped_column(
+        PostgresUUID,
+        ForeignKey("product_prices.id"),
+        nullable=False,
+    )
+
+    @declared_attr
+    def product_price(cls) -> Mapped["ProductPrice"]:
+        return relationship("ProductPrice", lazy="raise")
+
+    subscription_id: Mapped[UUID | None] = mapped_column(
+        PostgresUUID,
+        ForeignKey("subscriptions.id"),
+        nullable=True,
+    )
+
+    @declared_attr
+    def subscription(cls) -> Mapped["Subscription"]:
+        return relationship("Subscription", lazy="raise")