quay/endpoints/api/subscribe.py

""" Subscribe to plans. """
import logging
import stripe
import features
from app import billing
from endpoints.api import request_error, log_action
from data.billing import PLANS
from endpoints.api.subscribe_models_pre_oci import data_model as model
from endpoints.exception import NotFound


logger = logging.getLogger(__name__)


def check_repository_usage(user_or_org, plan_found):
  private_repos = model.get_private_repo_count(user_or_org.username)
  if plan_found is None:
    repos_allowed = 0
  else:
    repos_allowed = plan_found['privateRepos']

  if private_repos > repos_allowed:
    model.create_unique_notification('over_private_usage', user_or_org.username, {'namespace': user_or_org.username})
  else:
    model.delete_notifications_by_kind(user_or_org.username, 'over_private_usage')


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

def connection_response(exc):
  return {'message': 'Could not contact Stripe. Please try again.'}, 503


def subscription_view(stripe_subscription, used_repos):
  view = {
    'hasSubscription': True,
    'isExistingCustomer': True,
    'currentPeriodStart': stripe_subscription.current_period_start,
    'currentPeriodEnd': stripe_subscription.current_period_end,
    'plan': stripe_subscription.plan.id,
    'usedPrivateRepos': used_repos,
    'trialStart': stripe_subscription.trial_start,
    'trialEnd': stripe_subscription.trial_end
  }

  return view


def subscribe(user, plan, token, require_business_plan):
  if not features.BILLING:
    return

  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 = billing.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.error.CardError as e:
        return carderror_response(e)
      except stripe.error.APIConnectionError as e:
        return connection_response(e)

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

  else:
    # Change the plan
    try:
      cus = billing.Customer.retrieve(user.stripe_id)
    except stripe.error.APIConnectionError as e:
      return connection_response(e)

    if plan_found['price'] == 0:
      if cus.subscription is not None:
        # We only have to cancel the subscription if they actually have one
        try:
          cus.subscription.delete()
        except stripe.error.APIConnectionError as e:
          return connection_response(e)

        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.error.CardError as e:
        return carderror_response(e)
      except stripe.error.APIConnectionError as e:
        return connection_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