mirror of
https://github.com/All-Hands-AI/OpenHands.git
synced 2026-01-09 14:57:59 -05:00
94 lines
2.9 KiB
Python
94 lines
2.9 KiB
Python
"""Remove duplicates from stripe. This is a non standard alembic migration for non sql resources.
|
|
|
|
Revision ID: 019
|
|
Revises: 018
|
|
Create Date: 2025-03-20 16:30:00.000
|
|
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
from collections import defaultdict
|
|
from typing import Union
|
|
|
|
import sqlalchemy as sa
|
|
import stripe
|
|
from alembic import op
|
|
from sqlalchemy.sql import text
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = '019'
|
|
down_revision: Union[str, None] = '018'
|
|
branch_labels: Union[str, sa.Sequence[str], None] = None
|
|
depends_on: Union[str, sa.Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# Skip migration if STRIPE_API_KEY is not set
|
|
if 'STRIPE_API_KEY' not in os.environ:
|
|
print('Skipping migration 019: STRIPE_API_KEY not set')
|
|
return
|
|
|
|
stripe.api_key = os.environ['STRIPE_API_KEY']
|
|
|
|
# Get all users from stripe
|
|
user_id_to_customer_ids = defaultdict(list)
|
|
customers = stripe.Customer.list()
|
|
for customer in customers.auto_paging_iter():
|
|
user_id = customer.metadata.get('user_id')
|
|
if user_id:
|
|
user_id_to_customer_ids[user_id].append(customer.id)
|
|
|
|
# Canonical
|
|
stripe_customers = {
|
|
row[0]: row[1]
|
|
for row in op.get_bind().execute(
|
|
text('SELECT keycloak_user_id, stripe_customer_id FROM stripe_customers')
|
|
)
|
|
}
|
|
|
|
to_delete = []
|
|
for user_id, customer_ids in user_id_to_customer_ids.items():
|
|
if len(customer_ids) == 1:
|
|
continue
|
|
canonical_customer_id = stripe_customers.get(user_id)
|
|
if canonical_customer_id:
|
|
for customer_id in customer_ids:
|
|
if customer_id != canonical_customer_id:
|
|
to_delete.append({'user_id': user_id, 'customer_id': customer_id})
|
|
else:
|
|
# Prioritize deletion of items that don't have payment methods
|
|
to_delete_for_customer = []
|
|
for customer_id in customer_ids:
|
|
payment_methods = stripe.Customer.list_payment_methods(customer_id)
|
|
to_delete_for_customer.append(
|
|
{
|
|
'user_id': user_id,
|
|
'customer_id': customer_id,
|
|
'num_payment_methods': len(payment_methods),
|
|
}
|
|
)
|
|
to_delete_for_customer.sort(
|
|
key=lambda c: c['num_payment_methods'], reverse=True
|
|
)
|
|
to_delete.extend(to_delete_for_customer[1:])
|
|
|
|
for item in to_delete:
|
|
op.get_bind().execute(
|
|
text(
|
|
'INSERT INTO script_results (revision, data) VALUES (:revision, :data)'
|
|
),
|
|
{
|
|
'revision': revision,
|
|
'data': json.dumps(item),
|
|
},
|
|
)
|
|
stripe.Customer.delete(item['customer_id'])
|
|
|
|
|
|
def downgrade() -> None:
|
|
op.get_bind().execute(
|
|
text('DELETE FROM script_results WHERE revision=:revision'),
|
|
{'revision': revision},
|
|
)
|