quay/endpoints/api/subscribe.py

import logging
import stripe

from endpoints.api import request_error, log_action, NotFound
from endpoints.common import check_repository_usage
from data import model
from data.plans import PLANS


logger = logging.getLogger(__name__)


def carderror_response(exc):
  return {'carderror': exc.message}, 402


def subscription_view(stripe_subscription, used_repos):
  return {
    'currentPeriodStart': stripe_subscription.current_period_start,
    'currentPeriodEnd': stripe_subscription.current_period_end,
    'plan': stripe_subscription.plan.id,
    'usedPrivateRepos': used_repos,
  }


def subscribe(user, plan, token, require_business_plan):
  plan_found = None
  for plan_obj in PLANS:
    if plan_obj['stripeId'] == plan:
      plan_found = plan_obj

  if not plan_found or plan_found['deprecated']:
    logger.warning('Plan not found or deprecated: %s', plan)
    raise NotFound()

  if (require_business_plan and not plan_found['bus_features'] and not
                                  plan_found['price'] == 0):
    logger.warning('Business attempting to subscribe to personal plan: %s',
                   user.username)
    raise request_error(message='No matching plan found')

  private_repos = model.get_private_repo_count(user.username)

  # This is the default response
  response_json = {
    'plan': plan,
    'usedPrivateRepos': private_repos,
  }
  status_code = 200

  if not user.stripe_id:
    # Check if a non-paying user is trying to subscribe to a free plan
    if not plan_found['price'] == 0:
      # They want a real paying plan, create the customer and plan
      # simultaneously
      card = token
      
      try:
        cus = stripe.Customer.create(email=user.email, plan=plan, card=card)
        user.stripe_id = cus.id
        user.save()
        check_repository_usage(user, plan_found)
        log_action('account_change_plan', user.username, {'plan': plan})
      except stripe.CardError as e:
        return carderror_response(e)

      response_json = subscription_view(cus.subscription, private_repos)
      status_code = 201

  else:
    # Change the plan
    cus = stripe.Customer.retrieve(user.stripe_id)

    if plan_found['price'] == 0:
      if cus.subscription is not None:
        # We only have to cancel the subscription if they actually have one
        cus.cancel_subscription()
        cus.save()
        check_repository_usage(user, plan_found)
        log_action('account_change_plan', user.username, {'plan': plan})

    else:
      # User may have been a previous customer who is resubscribing
      if token:
        cus.card = token

      cus.plan = plan

      try:
        cus.save()
      except stripe.CardError as e:
        return carderror_response(e)
          
      response_json = subscription_view(cus.subscription, private_repos)
      check_repository_usage(user, plan_found)
      log_action('account_change_plan', user.username, {'plan': plan})

  return response_json, status_code