Skip to content

Commit

Permalink
feature: implement support for free trials
Browse files Browse the repository at this point in the history 
* Add `tenants.free_trial_start` column to indicate the day that the free trial began
* Refactor billing report to support generating daily line items
* Use these daily line items to calculate the free trial credit
  • Loading branch information
@jshearer
jshearer committed Sep 5, 2023
1 parent 9f1328a commit cabb00d
Show file tree
Hide file tree
Showing 2 changed files with 385 additions and 0 deletions.
300 changes: 300 additions & 0 deletions supabase/migrations/26_free_trial.sql
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,300 @@

-- Always use a transaction, y'all.
begin;

alter table tenants add column free_trial_start timestamptz;

create or replace function internal.compute_incremental_line_items(item_name text, item_unit text, item_unit_plural text, single_usage numeric, tiers integer[], running_usage_sum numeric)
returns jsonb as $$
declare
line_items jsonb = '[]';

-- Calculating tiered usage.
tier_rate integer;
tier_pivot integer;
tier_count numeric;

running_tier_pivot integer = 0;
begin
-- Walk up the tiers
for tier_idx in 1..array_length(tiers,1) by 2 loop
tier_rate = tiers[tier_idx];
tier_pivot = tiers[tier_idx+1];
if tier_pivot is null then
-- No limits here, roll all of the remaining usage into this tier
tier_count = single_usage;
running_usage_sum = running_usage_sum + tier_count;
line_items = line_items || jsonb_build_object(
'description', format(
'%s (at %s/%s)',
item_name,
(tier_rate / 100.0)::money,
item_unit
),
'count', tier_count,
'rate', tier_rate,
'subtotal_frac', tier_count * tier_rate
);
elsif tier_pivot > running_usage_sum then
running_tier_pivot = running_tier_pivot + tier_pivot;
-- We haven't already surpassed this tier's pivot
-- Calculate how much more usage we'd need to surpass this tier
tier_count = least(single_usage, running_tier_pivot - running_usage_sum);
single_usage = single_usage - tier_count;
running_usage_sum = running_usage_sum + tier_count;
line_items = line_items || jsonb_build_object(
'description', format(
case
when tier_idx = 1 then '%s (first %s%s at %s/%s)'
else '%s (next %s%s at %s/%s)'
end,
item_name,
tier_pivot,
item_unit_plural,
(tier_rate / 100.0)::money,
item_unit
),
'count', tier_count,
'rate', tier_rate,
'subtotal_frac', tier_count * tier_rate
);
end if;
end loop;

return jsonb_build_object(
'line_items', line_items,
'running_usage_sum', running_usage_sum
);
end
$$ language plpgsql;

create or replace function internal.incremental_usage_report(requested_grain text, billed_prefix catalog_prefix, billed_range tstzrange)
returns jsonb as $$
declare
-- Retrieved from tenants table.
data_tiers integer[];
usage_tiers integer[];

granules jsonb = '[]';
returned_data_line_items jsonb = '{}';
returned_hours_line_items jsonb = '{}';
combined_line_items jsonb;

-- We can't round these subtotals yet because we might want to add them up later.
-- Rounding before adding will result in inconsistent subtotals between daily and monthly
-- granularities, since the monthly granularity by definition sums up all fractional
-- values, which we then round later on. In reality the discrepancies are tiny since
-- the maximum error per grain is 1 ($0.01), but better to be entirely exact when dealing with money.
subtotal_frac numeric;

running_gb_sum numeric = 0;
running_hour_sum numeric = 0;
line_items jsonb = '[]';
begin
-- Because usage tiers reset at the beginning of every month, the logic defined here
-- is only correct when operating on at most a whole month.
if upper(billed_range) > date_trunc('month', lower(billed_range)) + '1 month' then
raise 'Invalid input range, must span at most one month: %', billed_range;
end if;

select into data_tiers, usage_tiers
t.data_tiers,
t.usage_tiers
from tenants t
where billed_prefix ^@ t.tenant;

-- Get all stats records for the selected time period at the selected granularity
select into granules
(select json_agg(res.obj) from (
select jsonb_build_object(
'processed_data_gb', sum((bytes_written_by_me + bytes_read_by_me)) / (1024.0 * 1024 * 1024),
'task_usage_hours', sum(usage_seconds) / (60.0 * 60),
'ts', ts
) as obj
from catalog_stats
where catalog_name ^@ billed_prefix
and grain = requested_grain
and billed_range @> ts
group by ts
) as res)
;

if granules is not null then
for idx in 0..jsonb_array_length(granules)-1 loop
returned_data_line_items = internal.compute_incremental_line_items('Data processing', 'GB', 'GB', (granules->idx->'processed_data_gb')::numeric, data_tiers, running_gb_sum);
running_gb_sum = (returned_data_line_items->'running_usage_sum')::numeric;

returned_hours_line_items = internal.compute_incremental_line_items('Task usage', 'hour', ' hours', (granules->idx->'task_usage_hours')::numeric, usage_tiers, running_hour_sum);
running_hour_sum = (returned_hours_line_items->'running_usage_sum')::numeric;

combined_line_items = (returned_data_line_items->'line_items')::jsonb || (returned_hours_line_items->'line_items')::jsonb;

select into subtotal_frac sum((item->'subtotal_frac')::numeric) from jsonb_array_elements(combined_line_items) as item;

line_items = line_items || jsonb_build_object(
'line_items', combined_line_items,
'subtotal_frac', subtotal_frac,
'processed_data_gb', (granules->idx->'processed_data_gb')::numeric,
'task_usage_hours', (granules->idx->'task_usage_hours')::numeric,
'ts', granules->idx->'ts'
);
end loop;
end if;

return line_items;
end
$$ language plpgsql;

-- Billing report which is effective August 2023.
create or replace function billing_report_202308(billed_prefix catalog_prefix, billed_month timestamptz)
returns jsonb as $$
#variable_conflict use_variable
declare
-- Auth checks
has_admin_grant boolean;
has_bypassrls boolean;

-- Computed
recurring_usd_cents integer;
free_trial_range tstzrange;
billed_range tstzrange;
free_trial_overlap tstzrange;

free_trial_credit numeric;

-- Temporary line items holders for free trial calculations
task_usage_line_items jsonb = '[]';
data_usage_line_items jsonb = '[]';

-- Calculating adjustments.
adjustment internal.billing_adjustments;

-- Aggregated outputs.
line_items jsonb = '[]';
subtotal_usd_cents integer;
processed_data_gb numeric;
task_usage_hours numeric;

-- Free trial outputs
free_trial_gb numeric;
free_trial_hours numeric;
begin

-- Ensure `billed_month` is the truncated start of the billed month.
billed_month = date_trunc('month', billed_month);
billed_range = tstzrange(billed_month, billed_month + '1 month', '[)');

-- Verify that the user has an admin grant for the requested `billed_prefix`.
perform 1 from auth_roles('admin') as r where billed_prefix ^@ r.role_prefix;
has_admin_grant = found;

-- Check whether the real active role has bypassrls flag set.
-- Because this function is SECURITY DEFINER, both `current_user` and `current_role`
-- will be `postgres`, which does have bypassrls set. Instead we want the
-- role of the caller, which can be accessed like so according to:
-- https://www.postgresql.org/message-id/13906.1141711109%40sss.pgh.pa.us
perform * from pg_roles where rolname = current_setting('role') and rolbypassrls = true;
has_bypassrls = found;

if not has_admin_grant and not has_bypassrls then
-- errcode 28000 causes PostgREST to return an HTTP 403
-- see: https://www.postgresql.org/docs/current/errcodes-appendix.html
-- and: https://postgrest.org/en/stable/errors.html#status-codes
raise 'You are not authorized for the billed prefix %', billed_prefix using errcode = 28000;
end if;

-- Fetch data & usage tiers for `billed_prefix`'s tenant.
select into free_trial_range
case
when t.free_trial_start is null then 'empty'::tstzrange
-- Inclusive start, exclusive end
else tstzrange(date_trunc('day', t.free_trial_start), date_trunc('day', t.free_trial_start) + '1 month', '[)')
end
from tenants t
where billed_prefix ^@ t.tenant
;
-- Reveal contract costs only when the computing tenant-level billing.
select into recurring_usd_cents t.recurring_usd_cents
from tenants t
where billed_prefix = t.tenant
;

-- Apply a recurring service cost, if defined.
if recurring_usd_cents != 0 then
line_items = line_items || jsonb_build_object(
'description', 'Recurring service charge',
'count', 1,
'rate', recurring_usd_cents,
'subtotal', recurring_usd_cents
);
end if;

-- Transform from `{"subtotal_frac": 1.98}` into `{"subtotal": 2}`
-- We can comfortably round here because we're loading the monthly granularity
-- meaning that summing has already happened.
select into line_items, processed_data_gb, task_usage_hours
line_items || (
select json_agg(
(item - 'subtotal_frac') ||
jsonb_build_object(
'subtotal', round((item->'subtotal_frac')::numeric)
)
)::jsonb
from jsonb_array_elements(report->0->'line_items') as item
),
(report->0->'processed_data_gb')::numeric,
(report->0->'task_usage_hours')::numeric
from internal.incremental_usage_report('monthly', billed_prefix, billed_range) as report;

-- Does the free trial range overlap the month in question?
if not isempty(free_trial_range) and (free_trial_range && billed_range) then
free_trial_overlap = billed_range * free_trial_range;
-- Sum up the fractional subtotals for each day in the portion of this
-- month covered by the free trial. Note that we don't want to round yet
-- since these are exact fractional values. Only after summing do we round.
select into
free_trial_credit coalesce(sum((line_item->>'subtotal_frac')::numeric), 0)
from
jsonb_array_elements(
internal.incremental_usage_report('daily', billed_prefix, free_trial_overlap)
) as line_item;

line_items = line_items || jsonb_build_object(
'description', 'Free trial credit',
'count', 1,
'rate', round(free_trial_credit) * -1,
'subtotal', round(free_trial_credit) * -1
);
end if;

-- Apply any billing adjustments.
for adjustment in select * from internal.billing_adjustments a
where a.billed_month = billed_month and a.tenant = billed_prefix
loop
line_items = line_items || jsonb_build_object(
'description', adjustment.detail,
'count', 1,
'rate', adjustment.usd_cents,
'subtotal', adjustment.usd_cents
);
end loop;

-- Roll up the final subtotal.
select into subtotal_usd_cents sum((l->>'subtotal')::numeric)
from jsonb_array_elements(line_items) l;

return jsonb_build_object(
'billed_month', billed_month,
'billed_prefix', billed_prefix,
'line_items', line_items,
'processed_data_gb', processed_data_gb,
'recurring_fee', coalesce(recurring_usd_cents, 0),
'subtotal', subtotal_usd_cents,
'task_usage_hours', task_usage_hours
);

end
$$ language plpgsql volatile security definer;

commit;

85 changes: 85 additions & 0 deletions supabase/tests/billing.test.sql
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,6 +160,8 @@ begin
) values
('aliceCo/aa/hello', 'monthly', '2022-08-01T00:00:00Z', '{}', 5.125 * 1024 * 1024 * 1024, 0, 3600 * 720),
('aliceCo/aa/big', 'monthly', '2022-08-01T00:00:00Z', '{}', 7::bigint * 1024 * 1024 * 1024, 9::bigint * 1024 * 1024 * 1024, 0),
('aliceCo/aa/big', 'daily', '2022-08-01T00:00:00Z', '{}', 6::bigint * 1024 * 1024 * 1024, 7::bigint * 1024 * 1024 * 1024, 0),
('aliceCo/aa/big', 'daily', '2022-08-30T00:00:00Z', '{}', 1::bigint * 1024 * 1024 * 1024, 2::bigint * 1024 * 1024 * 1024, 0),
('aliceCo/bb/world', 'monthly', '2022-08-01T00:00:00Z', '{}', 0, 22::bigint * 1024 * 1024 * 1024, 3600 * 18.375)
;

Expand Down Expand Up @@ -303,5 +305,88 @@ begin
'Attempting to fetch a report for aliceCo/ as Bob fails'
);

set role postgres;
-- Switch tiers so usage spills over
-- and set trial so half of August is covered
update tenants set
free_trial_start='2022-08-15',
data_tiers = '{50, 5, 20}'
where tenant = 'aliceCo/';

perform set_authenticated_context('11111111-1111-1111-1111-111111111111');

-- aliceCo/aa has a free trial set, and has free trial usage, so let's check that
return query select is(billing_report_202308('aliceCo/aa/', '2022-08-29T13:00:00Z'), '{
"billed_month": "2022-08-01T00:00:00+00:00",
"billed_prefix": "aliceCo/aa/",
"line_items": [
{
"count": 5,
"description": "Data processing (first 5GB at $0.50/GB)",
"rate": 50,
"subtotal": 250
},
{
"count": 16.125,
"description": "Data processing (at $0.20/GB)",
"rate": 20,
"subtotal": 323
},
{
"count": 720,
"description": "Task usage (at $0.15/hour)",
"rate": 15,
"subtotal": 10800
},
{
"count": 1,
"description": "Free trial credit",
"rate": -150,
"subtotal": -150
}
],
"processed_data_gb": 21.125,
"recurring_fee": 0,
"subtotal": 11223,
"task_usage_hours": 720
}'::jsonb);

-- aliceCo/bb has a free trial set, but has no usage in the trial period
-- this should result in a free trial credit line item of $0
return query select is(billing_report_202308('aliceCo/bb/', '2022-08-29T13:00:00Z'), '{
"billed_month": "2022-08-01T00:00:00+00:00",
"billed_prefix": "aliceCo/bb/",
"line_items": [
{
"count": 5,
"description": "Data processing (first 5GB at $0.50/GB)",
"rate": 50,
"subtotal": 250
},
{
"count": 17,
"description": "Data processing (at $0.20/GB)",
"rate": 20,
"subtotal": 340
},
{
"count": 18.375,
"description": "Task usage (at $0.15/hour)",
"rate": 15,
"subtotal": 276
},
{
"count": 1,
"description": "Free trial credit",
"rate": 0,
"subtotal": 0
}
],
"processed_data_gb": 22,
"recurring_fee": 0,
"subtotal": 866,
"task_usage_hours": 18.375
}'::jsonb);

end
$$ language plpgsql;

0 comments on commit cabb00d

Please sign in to comment.