From 2ab059eb67f496eea1abeb5fbce7b650ef77f451 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fran=C3=A7ois=20Voron?= Date: Thu, 16 May 2024 09:42:10 +0200 Subject: [PATCH] server: add a Sale model and migrate Transaction data with it --- .../versions/2024-05-15-1728_create_sale.py | 266 ++++++++++++++++++ server/polar/models/__init__.py | 2 + server/polar/models/sale.py | 60 ++++ server/polar/models/transaction.py | 37 +-- 4 files changed, 335 insertions(+), 30 deletions(-) create mode 100644 server/migrations/versions/2024-05-15-1728_create_sale.py create mode 100644 server/polar/models/sale.py diff --git a/server/migrations/versions/2024-05-15-1728_create_sale.py b/server/migrations/versions/2024-05-15-1728_create_sale.py new file mode 100644 index 0000000000..83ef82ce2b --- /dev/null +++ b/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 ### diff --git a/server/polar/models/__init__.py b/server/polar/models/__init__.py index 5cae01324c..a6ef1c50c9 100644 --- a/server/polar/models/__init__.py +++ b/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", diff --git a/server/polar/models/sale.py b/server/polar/models/sale.py new file mode 100644 index 0000000000..552db6febf --- /dev/null +++ b/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") diff --git a/server/polar/models/transaction.py b/server/polar/models/transaction.py index c4375584ff..37d1bb91da 100644 --- a/server/polar/models/transaction.py +++ b/server/polar/models/transaction.py @@ -3,12 +3,7 @@ from uuid import UUID from sqlalchemy import ForeignKey, Integer, String -from sqlalchemy.orm import ( - Mapped, - declared_attr, - mapped_column, - relationship, -) +from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship from polar.kit.db.models import RecordModel from polar.kit.extensions.sqlalchemy import PostgresUUID @@ -20,8 +15,7 @@ IssueReward, Organization, Pledge, - ProductPrice, - Subscription, + Sale, User, ) @@ -300,34 +294,17 @@ def pledge(cls) -> Mapped["Pledge | None"]: def donation(cls) -> Mapped["Donation | None"]: return relationship("Donation", lazy="raise") - subscription_id: Mapped[UUID | None] = mapped_column( + sale_id: Mapped[UUID | None] = mapped_column( PostgresUUID, - ForeignKey("subscriptions.id", ondelete="set null"), + ForeignKey("sales.id", ondelete="set null"), nullable=True, index=True, ) - """ID of the `Subscription` related to this transaction.""" + """ID of the `Sale` related to this transaction.""" @declared_attr - def subscription(cls) -> Mapped["Subscription | None"]: - return relationship("Subscription", lazy="raise") - - product_price_id: Mapped[UUID | None] = mapped_column( - PostgresUUID, - ForeignKey("product_prices.id", ondelete="set null"), - nullable=True, - index=True, - ) - """ - ID of the `ProductPrice` related to this transaction. - - Useful to keep track of the price at the time of the transaction, - which might change if the product is updated. - """ - - @declared_attr - def product_price(cls) -> Mapped["ProductPrice | None"]: - return relationship("ProductPrice", lazy="raise") + def sale(cls) -> Mapped["Sale | None"]: + return relationship("Sale", lazy="raise") issue_reward_id: Mapped[UUID | None] = mapped_column( PostgresUUID,