"""monthly_cloud_billing_revrec — Python transform job (ported from the BigQuery dbt python model models/warehouse/finance/monthly_cloud_billing_revrec.py). It reads two curated Databend products — ``salesforce_org_monthly_revrec`` (monthly billing per org) and ``salesforce_order_daily_view`` (the daily order snapshot) — runs the same pandas revenue-recognition algorithm as the BigQuery original, and writes two products back to Databend: * ``monthly_revrec_collections`` — one row per (org, bill month): billing + collection info. * ``monthly_revrec_orders`` — one row per order with recognised/adjusted revenue. The revenue-recognition logic (``process`` and its helpers) is a faithful port of the BigQuery model and is engine-agnostic pure pandas; only the I/O is swapped from BigQuery to Databend (via ``engine/databend_io.py``). dbt-databend-cloud cannot run python models, so this is a standalone ad-hoc script for now — structured so an orchestrator (e.g. Airflow) can later import ``run()``. Usage (pandas/numpy only — no dbt needed): uv run --extra pyjobs --env-file .env -- \ python transform/python/jobs/monthly_cloud_billing_revrec.py [--billing-month YYYY-MM] [--order-date YYYY-MM-DD] [--revenue recognised|predicted] [--dry-run] [--verbose] """ from __future__ import annotations import argparse import json import logging import sys from calendar import monthrange from datetime import date, datetime from pathlib import Path import numpy as np import pandas as pd sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from engine.databend_io import connect, read_df, replace_table # noqa: E402 logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) class NoOutputProductsError(RuntimeError): pass # Define revenue categories in one place for easy maintenance REVENUE_CATEGORIES = ['serverless', 'dedicated', 'support', 'starter', 'essential', 'premium', 'lake'] CREDIT_TYPES = ['program', 'redeem', 'customer_success'] COLLECTION_OUTPUT_COLUMNS = [ "bill_month", "adjust_tenant_id", "start_date", "end_date", "account_surrogate_key", "account_id", "account_name", "account_region", "payment_type", "marketplace_amount", "stripe_amount", "wire_transfer_amount", "total_amount", "stripe_tax_rate", "stripe_amount_including_tax", "stripe_invoice_number", "stripe_invoice", "stripe_invoice_paid", "stripe_invoice_finalized_date", "stripe_invoice_paid_date", "invoice_amount", "actual_revenue_recognition_date", ] ORDER_OUTPUT_BASE_COLUMNS = [ "tenant_id", "bill_month", "is_new_order", "is_new_salesforce_order", "account_surrogate_key", "account_id", "account_name", "account_region", "payment_type", "stripe_tax_rate", "collections_from", "order_number", "order_start_date", "order_end_date", "order_billing_start_date", "order_billing_end_date", "order_deferred_revenue_usd", "order_termination_date", "order_total_fees_usd", "order_accumulated_revenue_before_adjust", "order_accumulated_credit_amount", "revenue_amount", "credit_amount", "adjusted_revenue_amount", "adjusted_shortfall_amount", "is_order_closing", "billing_allocation_amount", "end_of_contract_shortfall_amount", "actual_revenue_recognition_date", ] # Generate order_columns dynamically based on revenue categories def generate_order_columns(): base_columns = [ 'accumulate_revenue_to_org_tenant_id', 'order_name', 'order_service_start_date', 'order_service_end_date', 'order_is_terminated', 'order_termination_date', 'real_order_termination_date', 'need_new_order', 'bill_month', 'reason', 'revenue_amount', 'credit_amount', 'accumulated_revenue', 'accumulated_revenue_before_adjust', 'accumulated_credit_amount', 'adjusted_revenue_amount', 'revenue_breakdown_ratio', ] # Add category-specific columns category_columns = [] for category in REVENUE_CATEGORIES: category_columns.extend([ f'{category}_amount', f'adjusted_{category}_amount', ]) # Add fixed columns fixed_columns = [ 'adjusted_shortfall_amount', 'order_net_fees', 'deferred_revenue', 'new_deferred_revenue', 'is_consider_credit_during_adj_rev_breakdown', 'is_order_closing', 'billing_allocation_amount', 'end_of_contract_shortfall_amount', ] return base_columns + category_columns + fixed_columns def generate_order_output_columns(): revenue_columns = [] for category in REVENUE_CATEGORIES: revenue_columns.extend([ f"{category}_revenue_amount", f"adjusted_{category}_amount", ]) return ORDER_OUTPUT_BASE_COLUMNS + revenue_columns order_columns = generate_order_columns() ORDER_OUTPUT_COLUMNS = generate_order_output_columns() def _is_missing_scalar(val): if val is None: return True if isinstance(val, float) and np.isnan(val): return True if val is pd.NaT: return True try: return bool(pd.isna(val)) except (TypeError, ValueError): return False def coalesce(val, val_to_return): if _is_missing_scalar(val): return val_to_return return val def format_optional_date_str(val): """ISO date strings; downstream expects string (e.g. Go bigquery.NullString).""" if _is_missing_scalar(val): return None if isinstance(val, pd.Timestamp): return val.strftime('%Y-%m-%d') if isinstance(val, datetime): return val.date().isoformat() if isinstance(val, date): return val.isoformat() if isinstance(val, str): s = val.strip() return s if s else None if isinstance(val, (int, np.integer)): s = str(int(val)).zfill(8) if len(s) == 8 and s.isdigit(): return f'{s[:4]}-{s[4:6]}-{s[6:8]}' return str(int(val)) return str(val) def create_on_demand_order( tenant_id, bill_month, start_date, end_date, **category_amounts): """ Create an on-demand order with dynamic revenue category amounts. category_amounts should contain keys like: serverless_amount, dedicated_amount, etc. """ # Get total amount from category_amounts or default parameter amount = category_amounts.pop('amount', 0) o = process_order({}) o['need_new_order'] = True o['accumulate_revenue_to_org_tenant_id'] = tenant_id o['order_name'] = '__new_order__' o['order_service_start_date'] = start_date o['order_service_end_date'] = end_date o['bill_month'] = bill_month o['revenue_amount'] = amount o['order_net_fees'] = amount o['accumulated_revenue_before_adjust'] = 0.0 # Set category amounts dynamically for category in REVENUE_CATEGORIES: amount_key = f'{category}_amount' o[amount_key] = category_amounts.get(amount_key, 0.0) # Create logging message dynamically category_breakdown = ' '.join([f'{cat}:{category_amounts.get(f"{cat}_amount", 0.0)}' for cat in REVENUE_CATEGORIES]) logger.info('creating on demand order for {} at {}: {} with revenue breakdowns: {}'.format( tenant_id, bill_month, amount, category_breakdown)) return o def append_json_string(existing_json_str, new_json_str): if _is_missing_scalar(existing_json_str) or existing_json_str == '': return None if _is_missing_scalar(new_json_str) else new_json_str if _is_missing_scalar(new_json_str) or new_json_str == '': return existing_json_str existing_list = json.loads(existing_json_str) new_list = json.loads(new_json_str) combined_list = existing_list + new_list return json.dumps(combined_list) def prioritise_orders(orders, billing_start_date, billing_end_date): # We need to prioritise orders based on the following rules: # 1. Orders that are terminated this month will have highest priority: order_is_terminated # 2. Orders that end naturally this month will have higher priority: order_service_end_date between billing_start_date and billing_end_date # 3. Then, orders with smaller start date will have higher priority: order_service_start_date # 4. Then, orders with a reason field will have higher priority: reason # We will sort the orders based on these rules if isinstance(orders, pd.DataFrame): # Create sorting keys for DataFrame # First key: order_is_terminated (True first) # Second key: ends this month (True first) # Third key: order_service_start_date (ascending) # Fourth key: has reason (True first) orders = orders.copy() orders['_ends_this_month'] = ( (orders['order_service_end_date'] >= billing_start_date) & (orders['order_service_end_date'] <= billing_end_date) ) orders['_has_reason'] = orders['reason'].notna() & (orders['reason'] != '') # Sort the DataFrame sorted_orders = orders.sort_values( by=['order_is_terminated', '_ends_this_month', 'order_service_start_date', '_has_reason'], ascending=[False, False, True, False] ) # Drop temporary columns sorted_orders = sorted_orders.drop(columns=['_ends_this_month', '_has_reason']) return sorted_orders elif isinstance(orders, list): # For handling lists of dictionaries (like in tests) def sort_key(order): is_terminated = bool(order.get('order_is_terminated') or False) end_date = order.get('order_service_end_date') start_date = order.get('order_service_start_date') or date.min has_reason = bool(order.get('reason', '')) ends_this_month = ( end_date and billing_start_date and billing_end_date and end_date >= billing_start_date and end_date <= billing_end_date ) # Return tuple for sorting: (terminated, ends_this_month, start_date, has_reason) # Negate first two booleans because True sorts after False by default return (-int(is_terminated), -int(ends_this_month), start_date, -int(has_reason)) return sorted(orders, key=sort_key) return orders def break_billing_segments(df): # Here we fetch all rows that have # segment_a_start_date, segment_a_end_date # segment_b_start_date, segment_b_end_date # segment_c_start_date, segment_c_end_date # that are not empty. # We want to split this row into multiple rows (depending on how many segment data we have). # Changes: # start_date = segment_start_date # end_date = segment_end_date # total_amount = segment_amount # Rest of rows remain the same. segments = ['a', 'b', 'c'] dfs = [] for segment in segments: mask = ( pd.notna(df[f'segment_{segment}_start_date']) & (df[f'segment_{segment}_start_date'] != '') & pd.notna(df[f'segment_{segment}_end_date']) & (df[f'segment_{segment}_end_date'] != '') ) segment_df = df[mask].copy() segment_ratio = segment_df[f'segment_{segment}_amount'] / segment_df['total_amount'] segment_df['billing_start_date'] = pd.to_datetime(segment_df[f'segment_{segment}_start_date']).dt.date segment_df['billing_end_date'] = pd.to_datetime(segment_df[f'segment_{segment}_end_date']).dt.date # We also need to set the breakdown columns by calculating the ratio of each segment amount to the total amount for category in REVENUE_CATEGORIES: raw_col = f'raw_{category}_amount' if raw_col in segment_df.columns: segment_df[raw_col] = segment_df[raw_col] * segment_ratio segment_df['total_amount'] = segment_df[f'segment_{segment}_amount'] dfs.append(segment_df) df_original = df[~df.index.isin(pd.concat(dfs).index)] df_combined = pd.concat([df_original] + dfs, ignore_index=True) return df_combined def process_order(order): o = {} for col in order_columns: if col in order: o[col] = order[col] else: o[col] = None return o def process(df_orders, df_billing): # TODO: TEMP: break billing by segments into individual rows df_billing = break_billing_segments(df_billing) # Filter orders so that only activated tidb cloud orders are considered df_orders = df_orders[ (df_orders['order_status'] == 'Activated') & (df_orders['order_product_family'] == 'TiDB Cloud') ] # For each billing recorrd, we need to know: # 1. How much billing occurred for that billing period # 2. What is the deferred revenue of order(s) for that billing period # 3. Whether existing order(s) end/terminate during that billing period all_orders = list() org_info = dict() closed_orders = set() for billing in df_billing.to_dict('records'): logger.info("processing row {}".format(str(billing['tenant_id']))) tenant_id = billing['accumulate_revenue_to_org_tenant_id'] org_billing = billing['total_amount'] bill_month = billing['billing_month'] bill_account_id = billing['account_id'] year, month = map(int, bill_month.split('-')) month_first_day = datetime(year, month, 1).date() month_last_day = datetime(year, month, monthrange(year, month)[1]).date() billing_start_date = billing['billing_start_date'] billing_end_date = billing['billing_end_date'] remaining_billing = org_billing def stripe_stuff(o, b, key, col_name): return append_json_string(o.get(key, {}).get(col_name, None), b[col_name]) # Collect org info org_info[(tenant_id, bill_month)] = { 'month_first_day': month_first_day, 'month_last_day': month_last_day, 'org_total_billing': org_info.get((tenant_id, bill_month), {}).get('org_total_billing', 0.0) + org_billing, 'account_surrogate_key': billing['customer_surrogate_key'], 'account_number': billing['account_number_for_netsuite'], 'account_name': billing['customer_name'], 'region': billing['region'], 'payment_type': billing['payment_type'], 'wire_transfer_amount': billing['wire_transfer_amount'], # to sum 'marketplace_amount': billing['marketplace_amount'], 'stripe_amount': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_amount'), 'stripe_tax_rate': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_tax_rate'), 'stripe_amount_including_tax': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_amount_including_tax'), 'stripe_invoice_number': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_number'), 'stripe_invoice': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice'), 'stripe_invoice_paid': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_paid'), 'stripe_invoice_finalized_date': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_finalized_date'), 'stripe_invoice_paid_date': stripe_stuff(org_info, billing, (tenant_id, bill_month), 'stripe_invoice_paid_date'), 'collections_from': '' if (billing['stripe_amount'] is None or billing['stripe_amount'] == '') else 'Stripe', 'invoice_amount': billing['invoice_amount'], 'is_salesforce_customer': billing['is_salesforce_customer'], 'actual_revenue_recognition_date': billing.get('actual_revenue_recognition_date'), } # Calculate total credit amount dynamically all_credit_fields = [] for category in REVENUE_CATEGORIES: for credit_type in CREDIT_TYPES: all_credit_fields.append(f'raw_{credit_type}_credit_amount_{category}') org_info[(tenant_id, bill_month)]['credit_amount'] = ( org_info.get((tenant_id, bill_month), {}).get('credit_amount', 0.0) + sum([coalesce(billing.get(field), 0.0) for field in all_credit_fields]) ) # Calculate category-specific credit amounts dynamically for category in REVENUE_CATEGORIES: category_credit_fields = [f'raw_{credit_type}_credit_amount_{category}' for credit_type in CREDIT_TYPES] credit_key = f'{category}_credit_amount' org_info[(tenant_id, bill_month)][credit_key] = ( org_info.get((tenant_id, bill_month), {}).get(credit_key, 0.0) + sum([coalesce(billing.get(field), 0.0) for field in category_credit_fields]) ) logger.info("processing billing record for org: {} at {}".format(tenant_id, bill_month)) bill_orders = [] # We only examine orders that are active between the billing start/end dates # We also filter out orders that are not terminated this month _df_orders = df_orders[ (df_orders['order_service_start_date'] <= billing_end_date) & (df_orders['order_service_end_date'] >= billing_start_date) & (df_orders['accumulate_revenue_to_org_tenant_id'] == tenant_id) & ( (df_orders['order_is_terminated'] == False) | ( (df_orders['order_is_terminated'] == True) & (df_orders['order_termination_date'] > billing_end_date) ) ) ] _df_orders = prioritise_orders(_df_orders, billing_start_date, billing_end_date) # Then, we loop through all orders for the tenant, and assign correct revenue to each order # The priority of orders which billing is applied to is as below: # 1. Orders that finish this month (order_service_end_date) # 2. Orders that were created earlier (created_date) # last. newly created on demand order(s) for _, order in _df_orders.iterrows(): logger.info("order net fees: {}, accumulated revenue b4 adjust: {}, accumulated_revenue: {}".format(order['order_net_fees'], order['accumulated_revenue_before_adjust'], order['accumulated_revenue'])) order['bill_month'] = bill_month order['order_net_fees'] = float(round(order['order_net_fees'], 4)) deferred_revenue = order['order_net_fees'] - order['accumulated_revenue_before_adjust'] if deferred_revenue < 0.02: logger.warning('Order {} does not have deferred revenue, skipping..'.format(order['order_name'])) continue # If order ends this month, we set the order_end flag if (order['order_service_end_date'] >= billing_start_date and order['order_service_end_date'] <= billing_end_date) or (order['order_is_terminated'] == True and order['real_order_termination_date'] == billing_end_date): logger.info('order {} ends this month.'.format(order['order_name'])) closed_orders.add((order['order_name'], tenant_id, bill_month)) if True: if deferred_revenue > remaining_billing: order['revenue_amount'] = remaining_billing if org_billing == 0: # Set all category amounts to 0 for category in REVENUE_CATEGORIES: order[f'{category}_amount'] = 0 else: # Proportionally assign breakdown amounts for category in REVENUE_CATEGORIES: raw_col = f'raw_{category}_amount' category_col = f'{category}_amount' order[category_col] = float(remaining_billing / org_billing) * billing.get(raw_col, 0.0) remaining_billing = 0 else: order['revenue_amount'] = deferred_revenue # Proportionally assign breakdown amounts for category in REVENUE_CATEGORIES: raw_col = f'raw_{category}_amount' category_col = f'{category}_amount' order[category_col] = float(deferred_revenue / org_billing) * billing.get(raw_col, 0.0) remaining_billing -= deferred_revenue # Create logging message dynamically category_breakdown = ' '.join([f'{cat}: {order.get(f"{cat}_amount", 0)}' for cat in REVENUE_CATEGORIES]) logger.info('order {} assigning revenue to this month. Revenue: {}, breakdowns: {}'.format( order['order_name'], order['revenue_amount'], category_breakdown)) if remaining_billing < 0: remaining_billing = 0 order['order_net_fees'] = float(round(order['order_net_fees'], 2)) bill_orders.append(process_order(order)) if remaining_billing <= 0: logger.info('all billing has been allocated for org: {} at {}'.format(tenant_id, bill_month)) break logger.info('total billing: {}'.format(str(org_billing))) logger.info('remaining billing: {}'.format(str(remaining_billing))) # If there are still money to be reconciled, then we need to create an on-demand order, for the bill, for that month if remaining_billing > 0.02: logger.info('on demand order created for bill {} at {}'.format(tenant_id, bill_month)) # Calculate category amounts dynamically category_amounts = {} for category in REVENUE_CATEGORIES: raw_col = f'raw_{category}_amount' amount_col = f'{category}_amount' category_amounts[amount_col] = float(remaining_billing / org_billing) * billing.get(raw_col, 0.0) category_amounts['amount'] = remaining_billing bill_orders.append(create_on_demand_order(tenant_id, bill_month, month_first_day, month_last_day, **category_amounts)) all_orders.extend(bill_orders) merged_orders = [] # Finally, we loop through all orders and combine order records that have been billed across multiple bills # We use (order_name, tenant_id, bill_month) as keys to merge for tenant_id, bill_month in list(org_info.keys()): _orders = prioritise_orders( [x for x in all_orders if x['accumulate_revenue_to_org_tenant_id'] == tenant_id and x['bill_month'] == bill_month], org_info[(tenant_id, bill_month)]['month_first_day'], org_info[(tenant_id, bill_month)]['month_last_day'] ) n_o_key = ('__new_order__', tenant_id, bill_month) merged = {} excess = 0.0 for order in _orders: key = (order['order_name'], tenant_id, bill_month) logger.info('processing order {}'.format(str(key))) if key not in merged: logger.info('adding order {}...'.format(str(key))) merged[key] = order else: logger.info('order {} is duplicated, merging...'.format(str(key))) merged_order = merged[key] merged_rev = merged_order['revenue_amount'] + order['revenue_amount'] # Calculate merged category amounts dynamically merged_category_amounts = {} for category in REVENUE_CATEGORIES: category_col = f'{category}_amount' merged_category_amounts[category_col] = merged_order.get(category_col, 0) + order.get(category_col, 0) deferred_revenue = merged_order['order_net_fees'] - merged_order['accumulated_revenue_before_adjust'] if merged_order['need_new_order']: merged_order['revenue_amount'] = merged_rev # Update all category amounts for category in REVENUE_CATEGORIES: category_col = f'{category}_amount' merged_order[category_col] = merged_category_amounts[category_col] merged_order['order_net_fees'] += order['order_net_fees'] # For existing orders, revenue is capped at deferred_revenue # Excess revenue amounts will be moved to new orders elif merged_rev > deferred_revenue: excess += merged_rev - deferred_revenue merged_order['revenue_amount'] = deferred_revenue # Proportionally adjust category amounts for category in REVENUE_CATEGORIES: category_col = f'{category}_amount' merged_order[category_col] = float(deferred_revenue / merged_rev) * merged_category_amounts[category_col] else: merged_order['revenue_amount'] = merged_rev # Update all category amounts for category in REVENUE_CATEGORIES: category_col = f'{category}_amount' merged_order[category_col] = merged_category_amounts[category_col] # If an order has been completed - either by being terminated / ends this month, or by being fully billed # We do not care about credit calculations for these orders if abs(merged[key]['order_net_fees'] - merged[key]['revenue_amount'] + merged[key]['accumulated_revenue_before_adjust']) <= 0.01 or key in closed_orders: merged[key]['is_consider_credit_during_adj_rev_breakdown'] = False else: merged[key]['is_consider_credit_during_adj_rev_breakdown'] = True if excess > 0: if n_o_key in merged: logger.info('moving excess revenue to new order...') merged[n_o_key]['revenue_amount'] += excess merged[n_o_key]['order_net_fees'] += excess # Update category amounts proportionally for category in REVENUE_CATEGORIES: category_col = f'{category}_amount' merged[n_o_key][category_col] += float(excess / merged_rev) * merged_category_amounts[category_col] logger.info(excess) else: logger.info('creating new order for excess...') # Calculate category amounts for new order excess_category_amounts = {} for category in REVENUE_CATEGORIES: category_col = f'{category}_amount' excess_category_amounts[category_col] = float(excess / merged_rev) * merged_category_amounts[category_col] excess_category_amounts['amount'] = excess _new_order = create_on_demand_order(tenant_id, bill_month, month_first_day, month_last_day, **excess_category_amounts) merged[n_o_key] = _new_order # We do not care about credit for on demand orders if key == n_o_key: merged[key]['credit_amount'] = 0.0 # Set all category credit amounts to 0 for category in REVENUE_CATEGORIES: merged[key][f'{category}_credit_amount'] = 0.0 merged[key]['is_consider_credit_during_adj_rev_breakdown'] = False else: merged[key]['credit_amount'] = org_info[(tenant_id, bill_month)]['credit_amount'] # Set category credit amounts for category in REVENUE_CATEGORIES: credit_key = f'{category}_credit_amount' merged[key][credit_key] = org_info[(tenant_id, bill_month)].get(credit_key, 0.0) # Set credit amount for org to 0 to avoid double counting org_info[(tenant_id, bill_month)]['credit_amount'] = 0.0 for category in REVENUE_CATEGORIES: credit_key = f'{category}_credit_amount' org_info[(tenant_id, bill_month)][credit_key] = 0.0 # Snapshot billing allocation before end-of-contract bump for key in merged.keys(): merged[key]['is_order_closing'] = key in closed_orders merged[key]['billing_allocation_amount'] = merged[key]['revenue_amount'] merged[key]['end_of_contract_shortfall_amount'] = 0.0 # Then, we loop through orders that end in that billing month, and set their revenue to be fully recognised if not already for key in merged.keys(): if key in closed_orders: _o = merged[key] deferred_revenue = _o['order_net_fees'] - _o['accumulated_revenue_before_adjust'] if _o['revenue_amount'] < deferred_revenue: _o['revenue_amount'] = deferred_revenue _o['end_of_contract_shortfall_amount'] = deferred_revenue - _o['billing_allocation_amount'] logger.info(f"order {_o['order_name']} ends this month, revenue adjusted to deferred revenue: {deferred_revenue}, shortfall: {_o['end_of_contract_shortfall_amount']}") merged_orders.extend(merged.values()) # Once revenue has been calculated for each order, we can now calculate adjusted values for the month. # TODO: This may need to be moved across scopes for order in merged_orders: revenue_amount = coalesce(order['revenue_amount'], 0) credit_amount = coalesce(order['credit_amount'], 0) acc_credit = coalesce(order['accumulated_credit_amount'], 0) acc_revenue = coalesce(order['accumulated_revenue_before_adjust'], 0) acc_adjusted_revenue = coalesce(order['accumulated_revenue'], 0) net_fees = coalesce(order['order_net_fees'], 0) deferred_revenue = net_fees - acc_revenue adjusted_deferred_revenue = net_fees - acc_adjusted_revenue is_consider_credit_during_adj_rev_breakdown = order['is_consider_credit_during_adj_rev_breakdown'] # Get category amounts and credit amounts category_amounts = {} category_credit_amounts = {} for category in REVENUE_CATEGORIES: category_amounts[category] = coalesce(order.get(f'{category}_amount'), 0) category_credit_amounts[category] = coalesce(order.get(f'{category}_credit_amount'), 0) logger.debug('order {}: rev: {}, credit: {}, acc_revenue_before_adjust: {}, acc_credit: {}, net_fees: {}, acc_rev: {}, deferred_revenue: {}, adj_deferred_revenue: {}'.format( order['order_name'], revenue_amount, credit_amount, acc_revenue, acc_credit, net_fees, acc_revenue, deferred_revenue, adjusted_deferred_revenue )) adjusted_revenue_amount = (((acc_revenue + revenue_amount + credit_amount + acc_credit) * net_fees) / (net_fees + credit_amount + acc_credit)) - net_fees + adjusted_deferred_revenue # Calculate other adjusted values def adjusted_values(val, credit_val): if is_consider_credit_during_adj_rev_breakdown: if revenue_amount + credit_amount == 0: return 0.0 return (val + credit_val) * adjusted_revenue_amount / (revenue_amount + credit_amount) else: if revenue_amount == 0: return 0.0 return val * adjusted_revenue_amount / revenue_amount # Calculate adjusted amounts for each category dynamically adjusted_category_amounts = {} for category in REVENUE_CATEGORIES: adjusted_amount = adjusted_values(category_amounts[category], category_credit_amounts[category]) adjusted_category_amounts[category] = adjusted_amount order[f'adjusted_{category}_amount'] = adjusted_amount # Calculate shortfall amount total_adjusted_categories = sum(adjusted_category_amounts.values()) adjusted_shortfall_amount = adjusted_revenue_amount - total_adjusted_categories # Calculate deferred revenues order['deferred_revenue'] = net_fees - acc_adjusted_revenue order['new_deferred_revenue'] = net_fees - (adjusted_revenue_amount + acc_adjusted_revenue) # Assign adjusted amounts order['adjusted_revenue_amount'] = adjusted_revenue_amount order['adjusted_shortfall_amount'] = adjusted_shortfall_amount # Create dynamic logging messages adjusted_breakdown = ' '.join([f'{cat}: {adjusted_category_amounts[cat]}' for cat in REVENUE_CATEGORIES]) original_breakdown = ' '.join([f'{cat}: {category_amounts[cat]}' for cat in REVENUE_CATEGORIES]) logger.info('{} {} adjusted: rev: {}, {}, shortfall: {}'.format( order['accumulate_revenue_to_org_tenant_id'], order['order_name'], adjusted_revenue_amount, adjusted_breakdown, adjusted_shortfall_amount)) logger.info('{} {} original: rev: {}, {}'.format( order['accumulate_revenue_to_org_tenant_id'], order['order_name'], revenue_amount, original_breakdown)) # Temp: We reorganise the records into existing row format for integration with downstream code. org_data = [] order_data = [] for key in list(org_info.keys()): tenant_id, bill_month = key org = org_info[key] cur_orders = [o for o in merged_orders if o['accumulate_revenue_to_org_tenant_id'] == tenant_id and o['bill_month'] == bill_month] if len(cur_orders) == 0: logger.warning("{} is missing matchable orders.".format(str(key))) continue _org = { "bill_month": bill_month, "adjust_tenant_id": tenant_id, "start_date": org['month_first_day'], "end_date": org['month_last_day'], "account_surrogate_key": org['account_surrogate_key'], "account_id": org['account_number'], "account_name": org['account_name'], "account_region": org['region'], "payment_type": org['payment_type'], "marketplace_amount": org['marketplace_amount'], "stripe_amount": org['stripe_amount'], "wire_transfer_amount": org['wire_transfer_amount'], "total_amount": org['org_total_billing'], "stripe_tax_rate": org['stripe_tax_rate'], "stripe_amount_including_tax": org['stripe_amount_including_tax'], "stripe_invoice_number": org['stripe_invoice_number'], "stripe_invoice": org['stripe_invoice'], "stripe_invoice_paid": org['stripe_invoice_paid'], "stripe_invoice_finalized_date": org['stripe_invoice_finalized_date'], "stripe_invoice_paid_date": org['stripe_invoice_paid_date'], "invoice_amount": org['invoice_amount'], "actual_revenue_recognition_date": format_optional_date_str(org.get('actual_revenue_recognition_date')), } org_data.append(_org) for order in cur_orders: tenant_id = order['accumulate_revenue_to_org_tenant_id'] bill_month = order['bill_month'] need_new_order = bool(coalesce(order.get('need_new_order'), False)) is_new_salesforce_order = bool(coalesce(org.get('is_salesforce_customer'), False) and need_new_order) if order.get('order_name', None) == '__new_order__': order_number = bill_month.replace('-', '_') + '__ON_DEMAND__' + str(tenant_id).replace('-', '_') else: order_number = order.get('order_name', None) _order = { "tenant_id": tenant_id, "bill_month": bill_month, "is_new_order": need_new_order, "is_new_salesforce_order": is_new_salesforce_order, "account_surrogate_key": org['account_surrogate_key'], "account_id": org['account_number'], "account_name": org['account_name'], "account_region": org['region'], "payment_type": org['payment_type'], "stripe_tax_rate": org['stripe_tax_rate'], "collections_from": org['collections_from'], "order_number": order_number, "order_start_date": order.get('order_service_start_date', None), "order_end_date": order.get('order_service_end_date', None), "order_billing_start_date": org['month_first_day'], "order_billing_end_date": org['month_last_day'], "order_deferred_revenue_usd": order.get('deferred_revenue'), "order_termination_date": format_optional_date_str( order.get('order_termination_date', None) ), "order_total_fees_usd": order.get('order_net_fees', None), "order_accumulated_revenue_before_adjust": order.get('accumulated_revenue_before_adjust'), "order_accumulated_credit_amount": coalesce(order.get('accumulated_credit_amount'), 0.0), "revenue_amount": order.get('revenue_amount', None), "credit_amount": order.get('credit_amount', None), "adjusted_revenue_amount": order.get('adjusted_revenue_amount', None), "adjusted_shortfall_amount": order.get('adjusted_shortfall_amount', None), "is_order_closing": order.get('is_order_closing', False), "billing_allocation_amount": order.get('billing_allocation_amount', None), "end_of_contract_shortfall_amount": order.get('end_of_contract_shortfall_amount', None), "actual_revenue_recognition_date": format_optional_date_str(org.get('actual_revenue_recognition_date')), } # Dynamically add revenue breakdown amounts for category in REVENUE_CATEGORIES: _order[f"{category}_revenue_amount"] = order.get(f'{category}_amount', None) _order[f"adjusted_{category}_amount"] = order.get(f'adjusted_{category}_amount', None) order_data.append(_order) df_org = pd.DataFrame.from_records(org_data, columns=COLLECTION_OUTPUT_COLUMNS) df_order = pd.DataFrame.from_records(order_data, columns=ORDER_OUTPUT_COLUMNS) # All-NULL columns otherwise autodetect as INT64; downstream expects STRING. if 'order_termination_date' in df_order.columns: df_order['order_termination_date'] = df_order['order_termination_date'].astype(pd.StringDtype()) return df_org, df_order # --- Databend I/O (replaces the BigQuery local_run of the original model) ------------------- # Billing numeric columns the algorithm does arithmetic on. Many are loaded VARCHAR (the seed is # all-VARCHAR) and must be numified; stripe_* are deliberately excluded — they are JSON strings # concatenated via append_json_string, not numbers. _BILLING_NUMERIC = ( {'total_amount', 'invoice_amount', 'marketplace_amount', 'wire_transfer_amount'} | {f'raw_{c}_amount' for c in REVENUE_CATEGORIES} | {f'segment_{s}_amount' for s in ('a', 'b', 'c')} | {f'raw_{ct}_credit_amount_{c}' for ct in CREDIT_TYPES for c in REVENUE_CATEGORIES} ) _BILLING_DATE = {'billing_start_date', 'billing_end_date'} # Order amounts are DECIMAL (auto-numified by the driver type); only the TINYINT flag needs forcing. _ORDER_BOOLEAN = {'order_is_terminated'} def run(billing_month=None, order_date=None, revenue='recognised', input_db='transform_dbt_dwh', output_db='transform_python_dwh', dry_run=False): """Read the two curated products from Databend, run the revrec algorithm, and write the two output products back. Returns (df_org, df_order).""" cur = connect() def scalar(sql): cur.execute(sql) rows = [list(r.values()) if hasattr(r, "values") else list(r) for r in cur.fetchall()] return rows[0][0] if rows and rows[0] else None if billing_month is None: billing_month = scalar(f"SELECT max(billing_month) FROM {input_db}.salesforce_org_monthly_revrec") if order_date is None: order_date = scalar(f"SELECT max(date_day) FROM {input_db}.salesforce_order_daily_view") or date.today() if isinstance(order_date, str): order_date = pd.to_datetime(order_date).date() logger.info("billing_month=%s order_date=%s revenue=%s in=%s out=%s", billing_month, order_date, revenue, input_db, output_db) df_billing = read_df( cur, f""" SELECT * FROM {input_db}.salesforce_org_monthly_revrec WHERE billing_month = '{billing_month}' ORDER BY billing_month, _row """, numeric=_BILLING_NUMERIC, date=_BILLING_DATE, ) df_orders = read_df( cur, f""" SELECT * FROM {input_db}.salesforce_order_daily_view WHERE date_day = '{order_date.isoformat()}' ORDER BY date_day, order_id """, boolean=_ORDER_BOOLEAN, ) # Match the BigQuery model's pre-process filtering. df_orders = df_orders[ (df_orders['date_day'] == order_date) & (df_orders['order_service_start_date'] <= order_date) ] if revenue == 'recognised': df_billing = df_billing[df_billing['actual_revenue_recognition_date'].notna() & (df_billing['actual_revenue_recognition_date'] != '')] else: df_billing = df_billing[df_billing['actual_revenue_recognition_date'].isna() | (df_billing['actual_revenue_recognition_date'] == '')] print(f"input: {len(df_billing)} billing rows ({revenue}), {len(df_orders)} order rows") if df_billing.empty: raise NoOutputProductsError("monthly_cloud_billing_revrec found no billing rows after filter") df_org, df_order = process(df_orders=df_orders, df_billing=df_billing) print(f"output: {len(df_org)} collection rows, {len(df_order)} order rows") if dry_run: print("--dry-run: not writing to Databend.") return df_org, df_order cur.execute(f"CREATE DATABASE IF NOT EXISTS {output_db}") n1 = replace_table(cur, f"{output_db}.monthly_revrec_collections", df_org) n2 = replace_table(cur, f"{output_db}.monthly_revrec_orders", df_order) print(f"wrote {output_db}.monthly_revrec_collections ({n1} rows), {output_db}.monthly_revrec_orders ({n2} rows)") return df_org, df_order def main(): p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) p.add_argument('--billing-month', help="billing_month to process (YYYY-MM); default = latest in the table") p.add_argument('--order-date', help="order snapshot date_day (YYYY-MM-DD); default = latest in the daily view") p.add_argument('--revenue', choices=['recognised', 'predicted'], default='recognised', help="recognised = rows with actual_revenue_recognition_date set (default); predicted = those without") p.add_argument('--input-db', default='transform_dbt_dwh') p.add_argument('--output-db', default='transform_python_dwh') p.add_argument('--dry-run', action='store_true', help="compute and report row counts but do not write") p.add_argument('--verbose', action='store_true', help="emit the per-order INFO logs") args = p.parse_args() if args.verbose: logging.basicConfig(level=logging.INFO, format='%(levelname)s %(message)s') run(billing_month=args.billing_month, order_date=args.order_date, revenue=args.revenue, input_db=args.input_db, output_db=args.output_db, dry_run=args.dry_run) if __name__ == '__main__': main()