docker/quay
docker/quay
1
0
Fork 0
mirror of https://github.com/quay/quay.git synced 2025-04-18 10:44:06 +03:00
Code

marketplace: free tier integration for reconciler (PROJQUAY-5698) (#3589)

Browse Source 
free sku integration for reconciliation worker
This commit is contained in:
Marcus Kok 2025-04-08 08:55:37 -04:00 committed by GitHub
parent 1af42c0b5a
commit 5e2fbd986b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 170 additions and 35 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
data/billing.py
View File

@ -372,6 +372,14 @@ PLANS = [
"sku_billing": True,
"plans_page_hidden": True,
},
{
"title": "freetier",
"privateRepos": 0,
"stripeId": "not_a_stripe_plan",
"rh_sku": "MW04192",
"sku_billing": False,
"plans_page_hidden": True,
},
]
RH_SKUS = [plan["rh_sku"] for plan in PLANS if plan.get("rh_sku") is not None]

78
util/marketplace.py
View File

@ -120,7 +120,7 @@ class RedHatSubscriptionApi(object):
if now_ms < end_date:
logger.debug("subscription found for %s", str(skuId))
valid_subscriptions.append(subscription)
return valid_subscriptions
return valid_subscriptions if len(valid_subscriptions) > 0 else None
return None
def extend_subscription(self, subscription_id, endDate):
@ -171,6 +171,7 @@ class RedHatSubscriptionApi(object):
"webCustomerId": customerId,
}
logger.debug("Created entitlement")
try:
r = requests.request(
method="post",
@ -187,6 +188,30 @@ class RedHatSubscriptionApi(object):
return r.status_code
def remove_entitlement(self, subscription_id):
"""
Removes subscription from user given subscription_id
"""
request_url = (
f"{self.marketplace_endpoint}/subscription/v5/terminateSubscription/{subscription_id}"
)
request_headers = {"Content-Type": "application/json"}
logger.debug("Terminating subscription with id %s", subscription_id)
try:
r = requests.request(
method="post",
url=request_url,
cert=self.cert,
headers=request_headers,
verify=True,
timeout=REQUEST_TIMEOUT,
)
except requests.exceptions.ReadTimeout:
logger.info("request to %s timed out", self.marketplace_endpoint)
return 408
return r.status_code
def get_subscription_details(self, subscription_id):
"""
Return the sku and expiration date for a specific subscription
@ -269,7 +294,6 @@ class RedHatSubscriptionApi(object):
# Mocked classes for unit tests
TEST_USER = {
"account_number": 12345,
"email": "subscriptions@devtable.com",
@ -358,6 +382,45 @@ FREE_USER = {
"email": "free_user@test.com",
"username": "free_user",
}
PAID_USER = {
"account_number": 34567,
"email": "paid@test.com",
"username": "paid_user",
"subscriptions": [
{
"id": 12345678,
"masterEndSystemName": "Quay",
"createdEndSystemName": "SUBSCRIPTION",
"createdDate": 1675957362000,
"lastUpdateEndSystemName": "SUBSCRIPTION",
"lastUpdateDate": 1675957362000,
"installBaseStartDate": 1707368400000,
"installBaseEndDate": 1707368399000,
"webCustomerId": 123456,
"subscriptionNumber": "12399889",
"quantity": 1,
"effectiveStartDate": 1707368400000,
"effectiveEndDate": 3813177600000,
}
],
"free_sku": [
{
"id": 56781234,
"masterEndSystemName": "Quay",
"createdEndSystemName": "SUBSCRIPTION",
"createdDate": 1675957362000,
"lastUpdateEndSystemName": "SUBSCRIPTION",
"lastUpdateDate": 1675957362000,
"installBaseStartDate": 1707368400000,
"installBaseEndDate": 1707368399000,
"webCustomerId": 123456,
"subscriptionNumber": "12399889",
"quantity": 1,
"effectiveStartDate": 1707368400000,
"effectiveEndDate": 3813177600000,
}
],
}
class FakeUserApi(RedHatUserApi):
@ -368,6 +431,8 @@ class FakeUserApi(RedHatUserApi):
def lookup_customer_id(self, email):
if email == TEST_USER["email"]:
return [TEST_USER["account_number"]]
if email == PAID_USER["email"]:
return [PAID_USER["account_number"]]
if email == FREE_USER["email"]:
return [FREE_USER["account_number"]]
if email == STRIPE_USER["email"]:
@ -391,11 +456,20 @@ class FakeSubscriptionApi(RedHatSubscriptionApi):
return [TEST_USER["private_subscription"]]
elif customer_id == TEST_USER["account_number"] and sku_id == "MW00584MO":
return [TEST_USER["reconciled_subscription"]]
elif customer_id == PAID_USER["account_number"] and sku_id == "MW02701":
return PAID_USER["subscriptions"]
elif customer_id == PAID_USER["account_number"] and sku_id == "MW04192":
return PAID_USER["free_sku"]
elif customer_id == FREE_USER["account_number"]:
return []
return None
def create_entitlement(self, customer_id, sku_id):
self.subscription_created = True
def remove_entitlement(self, subscription_id):
pass
def extend_subscription(self, subscription_id, end_date):
self.subscription_extended = True

2
util/test/test_marketplace.py
View File

@ -224,6 +224,8 @@ class TestMarketplace(unittest.TestCase):
assert extended_subscription is None
create_subscription_response = subscription_api.create_entitlement(12345, "sku")
assert create_subscription_response == 408
remove_entitlement_response = subscription_api.remove_entitlement(12345)
assert remove_entitlement_response == 408
@patch("requests.request")
def test_user_lookup(self, requests_mock):

100
workers/reconciliationworker.py
View File

@ -7,7 +7,7 @@ from app import app
from app import billing as stripe
from app import marketplace_subscriptions, marketplace_users
from data import model
from data.billing import RECONCILER_SKUS, get_plan
from data.billing import RECONCILER_SKUS, RH_SKUS, get_plan
from data.model import entitlements
from util.locking import GlobalLock, LockNotAcquiredException
from workers.gunicorn_worker import GunicornWorker
@ -24,6 +24,8 @@ MILLISECONDS_IN_SECONDS = 1000
SECONDS_IN_DAYS = 86400
ONE_MONTH = 30 * SECONDS_IN_DAYS * MILLISECONDS_IN_SECONDS
FREE_TIER_SKU = "MW04192"
class ReconciliationWorker(Worker):
def __init__(self):
@ -42,9 +44,9 @@ class ReconciliationWorker(Worker):
users = model.user.get_active_users(include_orgs=True)
stripe_users = [user for user in users if user.stripe_id is not None]
# stripe_users = [user for user in users if user.stripe_id is not None]
for user in stripe_users:
for user in users:
email = user.email
model_customer_ids = entitlements.get_web_customer_ids(user.id)
@ -77,39 +79,77 @@ class ReconciliationWorker(Worker):
for customer_id in model_customer_ids:
if customer_id not in customer_ids:
entitlements.remove_web_customer_id(user, customer_id)
# check for any subscription reconciliations
stripe_customer = None
if user.stripe_id is not None:
try:
stripe_customer = stripe.Customer.retrieve(user.stripe_id)
except stripe.error.APIConnectionError:
logger.error("Cannot connect to Stripe")
continue
except stripe.error.InvalidRequestException:
logger.warn("Invalid request for stripe_id %s", user.stripe_id)
continue
# check if we need to create a subscription for customer in RH marketplace
try:
stripe_customer = stripe.Customer.retrieve(user.stripe_id)
except stripe.error.APIConnectionError:
logger.error("Cannot connect to Stripe")
continue
except stripe.error.InvalidRequestError:
logger.warn("Invalid request for stripe_id %s", user.stripe_id)
continue
try:
subscription = stripe_customer.subscription
except AttributeError:
subscription = None
for sku_id in RECONCILER_SKUS:
if subscription is not None:
plan = get_plan(stripe_customer.subscription.plan.id)
if plan is None:
continue
if plan.get("rh_sku") == sku_id:
for customer_id in customer_ids:
subscription = marketplace_api.lookup_subscription(customer_id, sku_id)
if subscription is None:
logger.debug("Found %s to create for %s", sku_id, user.username)
marketplace_api.create_entitlement(customer_id, sku_id)
break
else:
logger.debug("User %s does not have a stripe subscription", user.username)
self._iterate_over_ids(stripe_customer, customer_ids, marketplace_api, user.username)
logger.debug("Finished work for user %s", user.username)
logger.info("Reconciliation worker is done")
def _iterate_over_ids(self, stripe_customer, customer_ids, marketplace_api, user=None):
"""
Iterate over each customer's web id(s) and perform appropriate reconciliation actions
"""
try:
subscription = stripe_customer.subscription
except AttributeError:
subscription = None
for customer_id in customer_ids:
paying = False
customer_skus = self._prefetch_user_entitlements(customer_id, marketplace_api)
if stripe_customer is not None and subscription is not None:
plan = get_plan(stripe_customer.subscription.plan.id)
if plan is not None:
# check for missing sku
paying = True
plan_sku = plan.get("rh_sku")
if plan_sku not in customer_skus:
logger.debug("Found %s to create for %s", plan_sku, user)
marketplace_api.create_entitlement(customer_id, plan_sku)
# check for free tier sku
else:
# not a stripe customer but we want to check for paying subscriptions for the
# next step
if len(customer_skus) == 1 and customer_skus[0] == FREE_TIER_SKU:
# edge case where there is only one free sku present
paying = False
else:
paying = len(customer_skus) > 0
# check for free-tier reconciliations
if not paying and FREE_TIER_SKU not in customer_skus:
marketplace_api.create_entitlement(customer_id, FREE_TIER_SKU)
elif paying and FREE_TIER_SKU in customer_skus:
free_tier_subscriptions = marketplace_api.lookup_subscription(
customer_id, FREE_TIER_SKU
)
# api returns a list of subscriptions so we want to make sure we remove
# all if there's more than one
for sub in free_tier_subscriptions:
id = sub.get("id")
marketplace_api.remove_entitlement(id)
def _prefetch_user_entitlements(self, customer_id, marketplace_api):
found_skus = []
for sku in RH_SKUS:
subscription = marketplace_api.lookup_subscription(customer_id, sku)
if subscription is not None and len(subscription) > 0:
found_skus.append(sku)
return found_skus
def _reconcile_entitlements(self, skip_lock_for_testing=False):
"""
Performs reconciliation for user entitlements

17
workers/test/test_reconciliationworker.py
View File

@ -19,7 +19,18 @@ def test_skip_free_user(initialized_db):
with patch.object(marketplace_subscriptions, "create_entitlement") as mock:
worker._perform_reconciliation(marketplace_users, marketplace_subscriptions)
mock.assert_not_called()
# adding the free tier
mock.assert_called_with(23456, "MW04192")
def test_remove_free_tier(initialized_db):
# if a user has a sku and also has a free tier, the free tier should be removed
paid_user = model.user.create_user("paid_user", "password", "paid@test.com")
paid_user.save()
marketplace_subscriptions.create_entitlement(12345, "MW04192")
with patch.object(marketplace_subscriptions, "remove_entitlement") as mock:
worker._perform_reconciliation(marketplace_users, marketplace_subscriptions)
mock.assert_called_with(56781234) # fake "free" tier subscription id mocked in marketplace.py
def test_reconcile_org_user(initialized_db):
@ -77,12 +88,12 @@ def test_reconcile_different_ids(initialized_db):
test_user = model.user.create_user("stripe_user", "password", "stripe_user@test.com")
test_user.stripe_id = "cus_" + "".join(random.choices(string.ascii_lowercase, k=14))
test_user.save()
model.entitlements.save_web_customer_id(test_user, 12345)
model.entitlements.save_web_customer_id(test_user, 55555)
worker._perform_reconciliation(marketplace_users, marketplace_subscriptions)
new_id = model.entitlements.get_web_customer_ids(test_user.id)
assert new_id != [12345]
assert new_id != [55555]
assert new_id == marketplace_users.lookup_customer_id(test_user.email)
# make sure it will remove account numbers from db that do not belong