feat(customer): Generate a unique short ID for invoice order number

Set the base26 encoding of the new short id as an infix for the invoice order number.

Author: Yopi

server/migrations/functions.py

@@ -0,0 +1,31 @@
+# PostgreSQL functions that are applied, that needs to be both migrated and recreated
+# in our tests.
+
+SEQUENCE_CREATE_CUSTOMER_SHORT_ID = "CREATE SEQUENCE customer_short_id_seq START 1;"
+
+# ID generation algorithm based on https://instagram-engineering.com/sharding-ids-at-instagram-1cf5a71e5a5c
+FUNCTION_GENERATE_CUSTOMER_SHORT_ID = """
+CREATE OR REPLACE FUNCTION generate_customer_short_id(creation_timestamp TIMESTAMP WITH TIME ZONE DEFAULT clock_timestamp())
+RETURNS bigint AS $$
+DECLARE
+    our_epoch bigint := 1672531200000; -- 2023-01-01 in milliseconds
+    seq_id bigint;
+    now_millis bigint;
+    result bigint;
+BEGIN
+    -- Get sequence number modulo 1024 (10 bits)
+    SELECT nextval('customer_short_id_seq') % 1024 INTO seq_id;
+
+    -- Use provided timestamp (defaults to clock_timestamp())
+    SELECT FLOOR(EXTRACT(EPOCH FROM creation_timestamp) * 1000) INTO now_millis;
+
+    -- 42 bits timestamp (milliseconds) | 10 bits sequence = 52 bits total
+    -- Capacity: 1,024 IDs per millisecond (over 1 million per second)
+    -- Combine: (timestamp - epoch) << 10 | sequence
+    result := (now_millis - our_epoch) << 10;
+    result := result | seq_id;
+
+    RETURN result;
+END;
+$$ LANGUAGE plpgsql;
+"""

server/migrations/versions/2025-10-24-1139_add_short_id_to_customers.py

@@ -0,0 +1,73 @@
+"""add short_id to customers
+
+Revision ID: aee75623c4bb
+Revises: 8c9ef8060e95
+Create Date: 2025-10-24 11:39:54.946342
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# Polar Custom Imports
+from migrations.functions import (
+    FUNCTION_GENERATE_CUSTOMER_SHORT_ID,
+    SEQUENCE_CREATE_CUSTOMER_SHORT_ID,
+)
+
+# revision identifiers, used by Alembic.
+revision = "aee75623c4bb"
+down_revision = "8c9ef8060e95"
+branch_labels: tuple[str] | None = None
+depends_on: tuple[str] | None = None
+
+
+def upgrade() -> None:
+    op.execute(SEQUENCE_CREATE_CUSTOMER_SHORT_ID)
+
+    # Create Instagram-style ID generator
+    # Combines timestamp (milliseconds since epoch) + sequence number
+    op.execute(FUNCTION_GENERATE_CUSTOMER_SHORT_ID)
+
+    op.add_column(
+        "customers",
+        sa.Column("short_id", sa.BigInteger(), nullable=True),
+    )
+
+    op.execute(
+        """
+        UPDATE customers
+        SET short_id = generate_customer_short_id(created_at);
+        """
+    )
+
+    op.alter_column("customers", "short_id", nullable=False)
+
+    op.execute(
+        """
+        ALTER TABLE customers
+        ALTER COLUMN short_id SET DEFAULT generate_customer_short_id();
+        """
+    )
+
+    op.create_unique_constraint(
+        op.f("customers_organization_id_short_id_key"),
+        "customers",
+        ["organization_id", "short_id"],
+    )
+    op.create_index(
+        op.f("ix_customers_short_id"), "customers", ["short_id"], unique=False
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_customers_short_id"), table_name="customers")
+    op.drop_constraint(
+        op.f("customers_organization_id_short_id_key"), "customers", type_="unique"
+    )
+    op.drop_column("customers", "short_id")
+    op.execute("DROP FUNCTION IF EXISTS generate_customer_short_id();")
+    op.execute("DROP SEQUENCE IF EXISTS customer_short_id_seq;")
+    # ### end Alembic commands ###

server/polar/models/customer.py

@@ -1,11 +1,13 @@
 import dataclasses
+import string
 import time
 from collections.abc import Sequence
 from datetime import datetime
 from enum import StrEnum
 from typing import TYPE_CHECKING, Any
 from uuid import UUID
 
+import sqlalchemy as sa
 from sqlalchemy import (
     TIMESTAMP,
     Boolean,
@@ -37,6 +39,21 @@
     from .subscription import Subscription
 
 
+def short_id_to_base26(short_id: int) -> str:
+    """Convert a numeric short_id to an 8-character base-26 string (A-Z)."""
+    chars = string.ascii_uppercase
+    result = ""
+    num = short_id
+
+    # Convert to base-26
+    while num > 0:
+        result = chars[num % 26] + result
+        num = num // 26
+
+    # Pad with 'A' to ensure 8 characters
+    return result.rjust(8, "A")
+
+
 class CustomerOAuthPlatform(StrEnum):
     github = "github"
     discord = "discord"
@@ -92,9 +109,16 @@ class Customer(MetadataMixin, RecordModel):
             postgresql_nulls_not_distinct=True,
         ),
         UniqueConstraint("organization_id", "external_id"),
+        UniqueConstraint("organization_id", "short_id"),
     )
 
     external_id: Mapped[str | None] = mapped_column(String, nullable=True, default=None)
+    short_id: Mapped[int] = mapped_column(
+        sa.BigInteger,
+        nullable=False,
+        index=True,
+        server_default=sa.text("generate_customer_short_id()"),
+    )
     email: Mapped[str] = mapped_column(String(320), nullable=False)
     email_verified: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
     stripe_customer_id: Mapped[str | None] = mapped_column(
@@ -217,6 +241,11 @@ def remove_oauth_account(
     def oauth_accounts(self) -> dict[str, Any]:
         return self._oauth_accounts
 
+    @property
+    def short_id_str(self) -> str:
+        """Get the base-26 string representation of the short_id."""
+        return short_id_to_base26(self.short_id)
+
     @property
     def legacy_user_id(self) -> UUID:
         return self._legacy_user_id or self.id

server/polar/organization/service.py

@@ -370,10 +370,9 @@ async def get_next_invoice_number(
     ) -> str:
         match organization.invoice_numbering:
             case InvoiceNumbering.customer:
-                customer_suffix = str(customer.id).split("-")[0].upper()
-                invoice_number = f"{organization.customer_invoice_prefix}-{customer_suffix}-{customer.invoice_next_number:04d}"
+                invoice_number = f"{organization.customer_invoice_prefix}-{customer.short_id_str}-{customer.invoice_next_number:04d}"
                 customer_repository = CustomerRepository.from_session(session)
-                await customer_repository.update(
+                customer = await customer_repository.update(
                     customer,
                     update_dict={
                         "invoice_next_number": customer.invoice_next_number + 1

server/tests/customer/test_short_id.py

@@ -0,0 +1,55 @@
+import pytest
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from polar.models.customer import short_id_to_base26
+
+
+@pytest.mark.asyncio
+async def test_generate_customer_short_id_function(session: AsyncSession) -> None:
+    """Test that the PostgreSQL generate_customer_short_id() function works correctly."""
+
+    result1 = await session.execute(text("SELECT generate_customer_short_id()"))
+    short_id1 = result1.scalar()
+    assert short_id1 is not None
+
+    result2 = await session.execute(text("SELECT generate_customer_short_id()"))
+    short_id2 = result2.scalar()
+    assert short_id2 is not None
+
+    result3 = await session.execute(text("SELECT generate_customer_short_id()"))
+    short_id3 = result3.scalar()
+    assert short_id3 is not None
+
+    assert short_id1 != short_id2
+    assert short_id2 != short_id3
+    assert short_id1 != short_id3
+
+    assert short_id1 > 0
+    assert short_id2 > 0
+    assert short_id3 > 0
+
+    assert short_id2 > short_id1
+    assert short_id3 > short_id2
+
+
+@pytest.mark.asyncio
+async def test_customer_short_id_to_base26(session: AsyncSession) -> None:
+    """Test that short_id converts to base-26 correctly."""
+
+    # Test some known conversions
+    assert short_id_to_base26(0) == "AAAAAAAA"
+    assert short_id_to_base26(1) == "AAAAAAAB"
+    assert short_id_to_base26(25) == "AAAAAAAZ"
+    assert short_id_to_base26(26) == "AAAAAABA"
+
+    # Test with a realistic generated ID
+    result = await session.execute(text("SELECT generate_customer_short_id()"))
+    short_id = result.scalar()
+    assert short_id is not None
+
+    base26 = short_id_to_base26(short_id)
+
+    assert len(base26) == 8
+    assert base26.isupper()
+    assert all(c in "ABCDEFGHIJKLMNOPQRSTUVWXYZ" for c in base26)

server/tests/fixtures/database.py

@@ -8,6 +8,10 @@
 from sqlalchemy.sql import text
 from sqlalchemy_utils import create_database, database_exists, drop_database
 
+from migrations.functions import (
+    FUNCTION_GENERATE_CUSTOMER_SHORT_ID,
+    SEQUENCE_CREATE_CUSTOMER_SHORT_ID,
+)
 from polar.config import settings
 from polar.kit.db.postgres import create_async_engine
 from polar.models import Model
@@ -45,6 +49,8 @@ async def initialize_test_database(worker_id: str) -> AsyncIterator[None]:
     async with engine.begin() as conn:
         await conn.execute(text("CREATE EXTENSION IF NOT EXISTS citext"))
         await conn.execute(text('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"'))
+        await conn.execute(text(SEQUENCE_CREATE_CUSTOMER_SHORT_ID))
+        await conn.execute(text(FUNCTION_GENERATE_CUSTOMER_SHORT_ID))
         await conn.run_sync(Model.metadata.create_all)
     await engine.dispose()
 

server/tests/order/test_service.py

@@ -3893,7 +3893,6 @@ async def test_different_customers_different_invoice_numbers(
         }
         await save_fixture(organization)
 
-        # Create two different customers
         customer_1 = await create_customer(
             save_fixture,
             organization=organization,
@@ -3909,7 +3908,6 @@ async def test_different_customers_different_invoice_numbers(
             stripe_customer_id="STRIPE_CUSTOMER_2",
         )
 
-        # Create checkout and order for customer 1
         checkout_1 = await create_checkout(
             save_fixture,
             products=[product_one_time],
@@ -3918,7 +3916,6 @@ async def test_different_customers_different_invoice_numbers(
         )
         order_1 = await order_service.create_from_checkout_one_time(session, checkout_1)
 
-        # Create checkout and order for customer 2
         checkout_2 = await create_checkout(
             save_fixture,
             products=[product_one_time],
@@ -3927,36 +3924,29 @@ async def test_different_customers_different_invoice_numbers(
         )
         order_2 = await order_service.create_from_checkout_one_time(session, checkout_2)
 
-        # Refresh to get the invoice numbers
         await session.refresh(order_1)
         await session.refresh(order_2)
 
-        # Both customers should have different invoice number sequences
-        # They should both start with the organization prefix but have customer-specific numbering
         assert order_1.invoice_number is not None
         assert order_2.invoice_number is not None
         assert order_1.invoice_number != order_2.invoice_number
 
-        # Verify they both start with organization prefix
         assert order_1.invoice_number.startswith(organization.customer_invoice_prefix)
         assert order_2.invoice_number.startswith(organization.customer_invoice_prefix)
 
-        # Each customer should start at 0001
         assert order_1.invoice_number.endswith("-0001")
         assert order_2.invoice_number.endswith("-0001")
 
-        # Verify invoice numbers contain customer-specific suffix
-        customer_1_suffix = str(customer_1.id).split("-")[0].upper()
-        customer_2_suffix = str(customer_2.id).split("-")[0].upper()
-        assert customer_1_suffix in order_1.invoice_number
-        assert customer_2_suffix in order_2.invoice_number
+        await session.refresh(customer_1)
+        await session.refresh(customer_2)
+        assert customer_1.short_id_str in order_1.invoice_number
+        assert customer_2.short_id_str in order_2.invoice_number
 
-        # Verify format: PREFIX-CUSTOMER_SUFFIX-0001
         assert (
             order_1.invoice_number
-            == f"{organization.customer_invoice_prefix}-{customer_1_suffix}-0001"
+            == f"{organization.customer_invoice_prefix}-{customer_1.short_id_str}-0001"
         )
         assert (
             order_2.invoice_number
-            == f"{organization.customer_invoice_prefix}-{customer_2_suffix}-0001"
+            == f"{organization.customer_invoice_prefix}-{customer_2.short_id_str}-0001"
         )