mirror of
https://github.com/OpenHands/OpenHands.git
synced 2026-03-22 13:47:19 +08:00
4a3a42c858
Co-authored-by: openhands <openhands@all-hands.dev> Co-authored-by: OpenHands Bot <contact@all-hands.dev>
127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
from uuid import UUID
|
|
|
|
import stripe
|
|
from server.constants import STRIPE_API_KEY
|
|
from server.logger import logger
|
|
from sqlalchemy import select
|
|
from storage.database import a_session_maker
|
|
from storage.org import Org
|
|
from storage.org_store import OrgStore
|
|
from storage.stripe_customer import StripeCustomer
|
|
|
|
stripe.api_key = STRIPE_API_KEY
|
|
|
|
|
|
async def find_customer_id_by_org_id(org_id: UUID) -> str | None:
|
|
async with a_session_maker() as session:
|
|
stmt = select(StripeCustomer).where(StripeCustomer.org_id == org_id)
|
|
result = await session.execute(stmt)
|
|
stripe_customer = result.scalar_one_or_none()
|
|
if stripe_customer:
|
|
return stripe_customer.stripe_customer_id
|
|
|
|
# If that fails, fallback to stripe
|
|
search_result = await stripe.Customer.search_async(
|
|
query=f"metadata['org_id']:'{str(org_id)}'",
|
|
)
|
|
data = search_result.data
|
|
if not data:
|
|
logger.info(
|
|
'no_customer_for_org_id',
|
|
extra={'org_id': str(org_id)},
|
|
)
|
|
return None
|
|
return data[0].id # type: ignore [attr-defined]
|
|
|
|
|
|
async def find_customer_id_by_user_id(user_id: str) -> str | None:
|
|
# First search our own DB...
|
|
org = await OrgStore.get_current_org_from_keycloak_user_id(user_id)
|
|
if not org:
|
|
logger.warning(f'Org not found for user {user_id}')
|
|
return None
|
|
customer_id = await find_customer_id_by_org_id(org.id)
|
|
return customer_id
|
|
|
|
|
|
async def find_or_create_customer_by_user_id(user_id: str) -> dict | None:
|
|
# Get the current org for the user
|
|
org = await OrgStore.get_current_org_from_keycloak_user_id(user_id)
|
|
if not org:
|
|
logger.warning(f'Org not found for user {user_id}')
|
|
return None
|
|
|
|
customer_id = await find_customer_id_by_org_id(org.id)
|
|
if customer_id:
|
|
return {'customer_id': customer_id, 'org_id': str(org.id)}
|
|
logger.info(
|
|
'creating_customer',
|
|
extra={'user_id': user_id, 'org_id': str(org.id)},
|
|
)
|
|
|
|
# Create the customer in stripe
|
|
customer = await stripe.Customer.create_async(
|
|
email=org.contact_email,
|
|
metadata={'org_id': str(org.id)},
|
|
)
|
|
|
|
# Save the stripe customer in the local db
|
|
async with a_session_maker() as session:
|
|
session.add(
|
|
StripeCustomer(
|
|
keycloak_user_id=user_id,
|
|
org_id=org.id,
|
|
stripe_customer_id=customer.id,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
logger.info(
|
|
'created_customer',
|
|
extra={
|
|
'user_id': user_id,
|
|
'org_id': str(org.id),
|
|
'stripe_customer_id': customer.id,
|
|
},
|
|
)
|
|
return {'customer_id': customer.id, 'org_id': str(org.id)}
|
|
|
|
|
|
async def has_payment_method_by_user_id(user_id: str) -> bool:
|
|
customer_id = await find_customer_id_by_user_id(user_id)
|
|
if customer_id is None:
|
|
return False
|
|
payment_methods = await stripe.Customer.list_payment_methods_async(
|
|
customer_id,
|
|
)
|
|
logger.info(
|
|
f'has_payment_method:{user_id}:{customer_id}:{bool(payment_methods.data)}'
|
|
)
|
|
return bool(payment_methods.data)
|
|
|
|
|
|
async def migrate_customer(user_id: str, org: Org):
|
|
async with a_session_maker() as session:
|
|
result = await session.execute(
|
|
select(StripeCustomer).where(StripeCustomer.keycloak_user_id == user_id)
|
|
)
|
|
stripe_customer = result.scalar_one_or_none()
|
|
if stripe_customer is None:
|
|
return
|
|
stripe_customer.org_id = org.id
|
|
customer = await stripe.Customer.modify_async(
|
|
id=stripe_customer.stripe_customer_id,
|
|
email=org.contact_email,
|
|
metadata={'user_id': '', 'org_id': str(org.id)},
|
|
)
|
|
|
|
logger.info(
|
|
'migrated_customer',
|
|
extra={
|
|
'user_id': user_id,
|
|
'org_id': str(org.id),
|
|
'stripe_customer_id': customer.id,
|
|
},
|
|
)
|
|
await session.commit()
|