|
|
@@ -0,0 +1,370 @@
|
|
|
+"""MySQL persistence for revenue observations and forecasts."""
|
|
|
+
|
|
|
+from __future__ import annotations
|
|
|
+
|
|
|
+import json
|
|
|
+from datetime import datetime
|
|
|
+from decimal import Decimal
|
|
|
+from typing import Any
|
|
|
+
|
|
|
+from revenue_forecast import RevenueForecast, RevenueObservation
|
|
|
+from revenue_speed_forecast import RevenueSpeedParameter, RevenueSpeedSample
|
|
|
+from storage import connect
|
|
|
+
|
|
|
+
|
|
|
+def _db_datetime(value: datetime | None) -> datetime | None:
|
|
|
+ if value is None:
|
|
|
+ return None
|
|
|
+ return value.replace(tzinfo=None)
|
|
|
+
|
|
|
+
|
|
|
+def _aware_datetime(value: datetime, timezone: Any) -> datetime:
|
|
|
+ if value.tzinfo is None:
|
|
|
+ return value.replace(tzinfo=timezone)
|
|
|
+ return value.astimezone(timezone)
|
|
|
+
|
|
|
+
|
|
|
+def load_previous_observation(
|
|
|
+ report_time: datetime,
|
|
|
+ source_version: str,
|
|
|
+) -> RevenueObservation | None:
|
|
|
+ connection = connect()
|
|
|
+ try:
|
|
|
+ with connection.cursor() as cursor:
|
|
|
+ cursor.execute(
|
|
|
+ """
|
|
|
+ SELECT *
|
|
|
+ FROM revenue_forecast_observation
|
|
|
+ WHERE report_time < %s
|
|
|
+ AND report_time >= DATE(%s)
|
|
|
+ AND source_version = %s
|
|
|
+ AND quality_status = 'VALID'
|
|
|
+ ORDER BY report_time DESC
|
|
|
+ LIMIT 1
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ _db_datetime(report_time),
|
|
|
+ _db_datetime(report_time),
|
|
|
+ source_version,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ row = cursor.fetchone()
|
|
|
+ return _row_to_observation(row, report_time.tzinfo) if row else None
|
|
|
+ finally:
|
|
|
+ connection.close()
|
|
|
+
|
|
|
+
|
|
|
+def upsert_observation(
|
|
|
+ observation: RevenueObservation,
|
|
|
+ *,
|
|
|
+ source_version: str,
|
|
|
+ lag_seconds: int,
|
|
|
+ quality_status: str,
|
|
|
+ quality_message: str | None,
|
|
|
+) -> None:
|
|
|
+ connection = connect()
|
|
|
+ try:
|
|
|
+ with connection.cursor() as cursor:
|
|
|
+ cursor.execute(
|
|
|
+ """
|
|
|
+ INSERT INTO revenue_forecast_observation
|
|
|
+ (source_partition, source_version, report_time, today_revenue,
|
|
|
+ overall_cpm, impressions, dau, fill_rate,
|
|
|
+ revenue_change_pct, cpm_change_pct,
|
|
|
+ exposure_change_pct, dau_change_pct,
|
|
|
+ fill_rate_change_pct,
|
|
|
+ source_modified_at, lag_seconds, quality_status,
|
|
|
+ quality_message)
|
|
|
+ VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
|
|
|
+ ON DUPLICATE KEY UPDATE
|
|
|
+ source_partition=VALUES(source_partition),
|
|
|
+ source_version=VALUES(source_version),
|
|
|
+ today_revenue=VALUES(today_revenue),
|
|
|
+ overall_cpm=VALUES(overall_cpm),
|
|
|
+ impressions=VALUES(impressions),
|
|
|
+ dau=VALUES(dau),
|
|
|
+ fill_rate=VALUES(fill_rate),
|
|
|
+ revenue_change_pct=VALUES(revenue_change_pct),
|
|
|
+ cpm_change_pct=VALUES(cpm_change_pct),
|
|
|
+ exposure_change_pct=VALUES(exposure_change_pct),
|
|
|
+ dau_change_pct=VALUES(dau_change_pct),
|
|
|
+ fill_rate_change_pct=VALUES(fill_rate_change_pct),
|
|
|
+ source_modified_at=VALUES(source_modified_at),
|
|
|
+ lag_seconds=VALUES(lag_seconds),
|
|
|
+ quality_status=VALUES(quality_status),
|
|
|
+ quality_message=VALUES(quality_message)
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ observation.partition,
|
|
|
+ source_version,
|
|
|
+ _db_datetime(observation.report_time),
|
|
|
+ observation.today_revenue,
|
|
|
+ observation.overall_cpm,
|
|
|
+ observation.impressions,
|
|
|
+ observation.dau,
|
|
|
+ observation.fill_rate,
|
|
|
+ observation.revenue_change_pct,
|
|
|
+ observation.cpm_change_pct,
|
|
|
+ observation.exposure_change_pct,
|
|
|
+ observation.dau_change_pct,
|
|
|
+ observation.fill_rate_change_pct,
|
|
|
+ _db_datetime(observation.source_modified_at),
|
|
|
+ lag_seconds,
|
|
|
+ quality_status,
|
|
|
+ quality_message,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ finally:
|
|
|
+ connection.close()
|
|
|
+
|
|
|
+
|
|
|
+def save_forecast(forecast: RevenueForecast) -> None:
|
|
|
+ channel_costs_json = json.dumps(
|
|
|
+ {channel: str(cost) for channel, cost in forecast.channel_costs},
|
|
|
+ ensure_ascii=False,
|
|
|
+ sort_keys=True,
|
|
|
+ )
|
|
|
+ connection = connect()
|
|
|
+ try:
|
|
|
+ with connection.cursor() as cursor:
|
|
|
+ cursor.execute(
|
|
|
+ """
|
|
|
+ INSERT INTO revenue_forecast_result
|
|
|
+ (report_time, forecast_version, forecast_method,
|
|
|
+ parameter_version, current_cumulative_revenue,
|
|
|
+ speed_latest_weight, latest_interval_revenue,
|
|
|
+ previous_interval_revenue, weighted_speed,
|
|
|
+ remaining_multiplier_p10, remaining_multiplier_p50,
|
|
|
+ remaining_multiplier_p90, parameter_sample_count,
|
|
|
+ forecast_revenue, forecast_lower, forecast_upper,
|
|
|
+ target_cost_ratio, target_daily_cost, cost_reserve_date,
|
|
|
+ channel_costs_json, non_miniapp_reserved_cost,
|
|
|
+ miniapp_target_daily_cost, forecast_status)
|
|
|
+ VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
|
|
|
+ ON DUPLICATE KEY UPDATE
|
|
|
+ forecast_method=VALUES(forecast_method),
|
|
|
+ parameter_version=VALUES(parameter_version),
|
|
|
+ current_cumulative_revenue=VALUES(current_cumulative_revenue),
|
|
|
+ speed_latest_weight=VALUES(speed_latest_weight),
|
|
|
+ latest_interval_revenue=VALUES(latest_interval_revenue),
|
|
|
+ previous_interval_revenue=VALUES(previous_interval_revenue),
|
|
|
+ weighted_speed=VALUES(weighted_speed),
|
|
|
+ remaining_multiplier_p10=VALUES(remaining_multiplier_p10),
|
|
|
+ remaining_multiplier_p50=VALUES(remaining_multiplier_p50),
|
|
|
+ remaining_multiplier_p90=VALUES(remaining_multiplier_p90),
|
|
|
+ parameter_sample_count=VALUES(parameter_sample_count),
|
|
|
+ forecast_revenue=VALUES(forecast_revenue),
|
|
|
+ forecast_lower=VALUES(forecast_lower),
|
|
|
+ forecast_upper=VALUES(forecast_upper),
|
|
|
+ target_cost_ratio=VALUES(target_cost_ratio),
|
|
|
+ target_daily_cost=VALUES(target_daily_cost),
|
|
|
+ cost_reserve_date=VALUES(cost_reserve_date),
|
|
|
+ channel_costs_json=VALUES(channel_costs_json),
|
|
|
+ non_miniapp_reserved_cost=VALUES(non_miniapp_reserved_cost),
|
|
|
+ miniapp_target_daily_cost=VALUES(miniapp_target_daily_cost),
|
|
|
+ forecast_status=VALUES(forecast_status)
|
|
|
+ """,
|
|
|
+ (
|
|
|
+ _db_datetime(forecast.report_time),
|
|
|
+ forecast.forecast_version,
|
|
|
+ forecast.forecast_method,
|
|
|
+ forecast.parameter_version,
|
|
|
+ forecast.current_cumulative_revenue,
|
|
|
+ forecast.speed_latest_weight,
|
|
|
+ forecast.latest_interval_revenue,
|
|
|
+ forecast.previous_interval_revenue,
|
|
|
+ forecast.weighted_speed,
|
|
|
+ forecast.remaining_multiplier_p10,
|
|
|
+ forecast.remaining_multiplier_p50,
|
|
|
+ forecast.remaining_multiplier_p90,
|
|
|
+ forecast.parameter_sample_count,
|
|
|
+ forecast.forecast_revenue,
|
|
|
+ forecast.forecast_lower,
|
|
|
+ forecast.forecast_upper,
|
|
|
+ forecast.target_cost_ratio,
|
|
|
+ forecast.target_daily_cost,
|
|
|
+ forecast.cost_reserve_date,
|
|
|
+ channel_costs_json,
|
|
|
+ forecast.non_miniapp_reserved_cost,
|
|
|
+ forecast.miniapp_target_daily_cost,
|
|
|
+ forecast.status,
|
|
|
+ ),
|
|
|
+ )
|
|
|
+ finally:
|
|
|
+ connection.close()
|
|
|
+
|
|
|
+
|
|
|
+def load_speed_parameter(
|
|
|
+ report_time: datetime,
|
|
|
+ parameter_version: str,
|
|
|
+) -> RevenueSpeedParameter | None:
|
|
|
+ connection = connect()
|
|
|
+ try:
|
|
|
+ with connection.cursor() as cursor:
|
|
|
+ cursor.execute(
|
|
|
+ """
|
|
|
+ SELECT *, TIME_FORMAT(time_slot, '%%H:%%i:%%s') AS time_slot_text
|
|
|
+ FROM revenue_forecast_speed_parameter
|
|
|
+ WHERE parameter_version = %s
|
|
|
+ AND time_slot = TIME(%s)
|
|
|
+ LIMIT 1
|
|
|
+ """,
|
|
|
+ (parameter_version, _db_datetime(report_time)),
|
|
|
+ )
|
|
|
+ row = cursor.fetchone()
|
|
|
+ finally:
|
|
|
+ connection.close()
|
|
|
+ if not row:
|
|
|
+ return None
|
|
|
+ return RevenueSpeedParameter(
|
|
|
+ parameter_version=str(row["parameter_version"]),
|
|
|
+ time_slot=datetime.strptime(row["time_slot_text"], "%H:%M:%S").time(),
|
|
|
+ training_start=row["training_start"],
|
|
|
+ training_end=row["training_end"],
|
|
|
+ sample_count=int(row["sample_count"]),
|
|
|
+ latest_weight=Decimal(str(row["latest_weight"])),
|
|
|
+ multiplier_p10=Decimal(str(row["multiplier_p10"])),
|
|
|
+ multiplier_p50=Decimal(str(row["multiplier_p50"])),
|
|
|
+ multiplier_p90=Decimal(str(row["multiplier_p90"])),
|
|
|
+ multiplier_mad=Decimal(str(row["multiplier_mad"])),
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
+def publish_speed_parameter_release(
|
|
|
+ *,
|
|
|
+ parameter_version: str,
|
|
|
+ samples: list[RevenueSpeedSample],
|
|
|
+ parameters: list[RevenueSpeedParameter],
|
|
|
+) -> None:
|
|
|
+ if not samples or not parameters:
|
|
|
+ raise ValueError("Cannot publish an empty speed parameter release")
|
|
|
+ if any(item.parameter_version != parameter_version for item in samples):
|
|
|
+ raise ValueError("Sample parameter version mismatch")
|
|
|
+ if any(item.parameter_version != parameter_version for item in parameters):
|
|
|
+ raise ValueError("Aggregate parameter version mismatch")
|
|
|
+
|
|
|
+ connection = connect()
|
|
|
+ try:
|
|
|
+ connection.autocommit(False)
|
|
|
+ with connection.cursor() as cursor:
|
|
|
+ cursor.execute(
|
|
|
+ "SELECT COUNT(*) AS row_count "
|
|
|
+ "FROM revenue_forecast_speed_parameter "
|
|
|
+ "WHERE parameter_version=%s",
|
|
|
+ (parameter_version,),
|
|
|
+ )
|
|
|
+ if int(cursor.fetchone()["row_count"]) > 0:
|
|
|
+ raise ValueError(
|
|
|
+ f"Parameter version already published: {parameter_version}"
|
|
|
+ )
|
|
|
+ cursor.executemany(
|
|
|
+ """
|
|
|
+ INSERT INTO revenue_forecast_speed_sample
|
|
|
+ (data_date, time_slot, parameter_version,
|
|
|
+ cumulative_revenue, final_revenue,
|
|
|
+ latest_interval_revenue, previous_interval_revenue,
|
|
|
+ latest_weight, weighted_speed, remaining_revenue,
|
|
|
+ remaining_multiplier)
|
|
|
+ VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
|
|
|
+ """,
|
|
|
+ [
|
|
|
+ (
|
|
|
+ item.data_date,
|
|
|
+ item.report_time.time(),
|
|
|
+ item.parameter_version,
|
|
|
+ item.cumulative_revenue,
|
|
|
+ item.final_revenue,
|
|
|
+ item.latest_interval_revenue,
|
|
|
+ item.previous_interval_revenue,
|
|
|
+ item.latest_weight,
|
|
|
+ item.weighted_speed,
|
|
|
+ item.remaining_revenue,
|
|
|
+ item.remaining_multiplier,
|
|
|
+ )
|
|
|
+ for item in samples
|
|
|
+ ],
|
|
|
+ )
|
|
|
+ cursor.executemany(
|
|
|
+ """
|
|
|
+ INSERT INTO revenue_forecast_speed_parameter
|
|
|
+ (parameter_version, time_slot, training_start,
|
|
|
+ training_end, sample_count, latest_weight,
|
|
|
+ multiplier_p10, multiplier_p50, multiplier_p90,
|
|
|
+ multiplier_mad)
|
|
|
+ VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
|
|
|
+ """,
|
|
|
+ [
|
|
|
+ (
|
|
|
+ item.parameter_version,
|
|
|
+ item.time_slot,
|
|
|
+ item.training_start,
|
|
|
+ item.training_end,
|
|
|
+ item.sample_count,
|
|
|
+ item.latest_weight,
|
|
|
+ item.multiplier_p10,
|
|
|
+ item.multiplier_p50,
|
|
|
+ item.multiplier_p90,
|
|
|
+ item.multiplier_mad,
|
|
|
+ )
|
|
|
+ for item in parameters
|
|
|
+ ],
|
|
|
+ )
|
|
|
+ connection.commit()
|
|
|
+ except Exception:
|
|
|
+ connection.rollback()
|
|
|
+ raise
|
|
|
+ finally:
|
|
|
+ connection.autocommit(True)
|
|
|
+ connection.close()
|
|
|
+
|
|
|
+
|
|
|
+def _row_to_observation(row: dict[str, Any], timezone: Any) -> RevenueObservation:
|
|
|
+ return RevenueObservation(
|
|
|
+ partition=str(row["source_partition"]),
|
|
|
+ report_time=_aware_datetime(row["report_time"], timezone),
|
|
|
+ today_revenue=Decimal(str(row["today_revenue"])),
|
|
|
+ overall_cpm=(
|
|
|
+ Decimal(str(row["overall_cpm"]))
|
|
|
+ if row.get("overall_cpm") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ impressions=(
|
|
|
+ int(row["impressions"]) if row.get("impressions") is not None else None
|
|
|
+ ),
|
|
|
+ dau=int(row["dau"]) if row.get("dau") is not None else None,
|
|
|
+ fill_rate=(
|
|
|
+ Decimal(str(row["fill_rate"]))
|
|
|
+ if row.get("fill_rate") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ revenue_change_pct=(
|
|
|
+ Decimal(str(row["revenue_change_pct"]))
|
|
|
+ if row.get("revenue_change_pct") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ cpm_change_pct=(
|
|
|
+ Decimal(str(row["cpm_change_pct"]))
|
|
|
+ if row.get("cpm_change_pct") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ exposure_change_pct=(
|
|
|
+ Decimal(str(row["exposure_change_pct"]))
|
|
|
+ if row.get("exposure_change_pct") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ dau_change_pct=(
|
|
|
+ Decimal(str(row["dau_change_pct"]))
|
|
|
+ if row.get("dau_change_pct") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ fill_rate_change_pct=(
|
|
|
+ Decimal(str(row["fill_rate_change_pct"]))
|
|
|
+ if row.get("fill_rate_change_pct") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ source_modified_at=(
|
|
|
+ _aware_datetime(row["source_modified_at"], timezone)
|
|
|
+ if row.get("source_modified_at") is not None
|
|
|
+ else None
|
|
|
+ ),
|
|
|
+ )
|