storage.py 52 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431
  1. """实时调控使用的 MySQL 状态与审计存储。"""
  2. from __future__ import annotations
  3. import os
  4. import json
  5. from contextlib import contextmanager
  6. from datetime import date, datetime
  7. from pathlib import Path
  8. from typing import Any, Iterator
  9. import pymysql
  10. import pymysql.cursors
  11. ROOT = Path(__file__).resolve().parent
  12. def _suppress_performance_agency_delivery(cursor) -> None:
  13. """确保 PERFORMANCE 记录对当前及回滚版本的代理商读取链路均不可见。"""
  14. cursor.execute(
  15. """
  16. UPDATE creative_rejection_cleanup_item
  17. SET agency_name='',
  18. agency_notified_at=COALESCE(agency_notified_at, NOW())
  19. WHERE LEFT(cleanup_rule_type, 12)='PERFORMANCE_'
  20. AND (
  21. COALESCE(agency_name, '')<>''
  22. OR agency_notified_at IS NULL
  23. )
  24. """
  25. )
  26. def _backfill_legacy_review_rule_types(cursor) -> None:
  27. """新增 cleanup_rule_type 后回填旧数据中的部分审核语义。"""
  28. cursor.execute(
  29. """
  30. UPDATE creative_rejection_cleanup_item
  31. SET cleanup_rule_type='REVIEW_PARTIAL'
  32. WHERE cleanup_rule_type='REVIEW_DENIED'
  33. AND (
  34. cleanup_action='ALERT_ONLY'
  35. OR pre_state_json LIKE
  36. '%CREATIVE_SET_APPROVAL_STATUS_PARTIAL_NORMAL%'
  37. OR action_reason LIKE '部分投放中%'
  38. )
  39. """
  40. )
  41. def connect() -> pymysql.Connection:
  42. required = ["DB_HOST", "DB_USER", "DB_NAME"]
  43. missing = [key for key in required if not os.getenv(key)]
  44. if missing:
  45. raise RuntimeError(f"Missing database environment variables: {', '.join(missing)}")
  46. return pymysql.connect(
  47. host=os.environ["DB_HOST"],
  48. port=int(os.getenv("DB_PORT", "3306")),
  49. user=os.environ["DB_USER"],
  50. password=os.getenv("DB_PASSWORD", ""),
  51. database=os.environ["DB_NAME"],
  52. charset="utf8mb4",
  53. cursorclass=pymysql.cursors.DictCursor,
  54. autocommit=True,
  55. connect_timeout=int(os.getenv("DB_CONNECT_TIMEOUT", "10")),
  56. read_timeout=int(os.getenv("DB_READ_TIMEOUT", "60")),
  57. write_timeout=int(os.getenv("DB_WRITE_TIMEOUT", "60")),
  58. )
  59. def initialize_schema() -> None:
  60. statements = [
  61. statement.strip()
  62. for statement in (ROOT / "schema.sql").read_text(encoding="utf-8").split(";")
  63. if statement.strip()
  64. ]
  65. connection = connect()
  66. try:
  67. with connection.cursor() as cursor:
  68. for statement in statements:
  69. cursor.execute(statement)
  70. cursor.execute(
  71. """
  72. SELECT COLUMN_NAME
  73. FROM information_schema.COLUMNS
  74. WHERE TABLE_SCHEMA=%s
  75. AND TABLE_NAME='realtime_control_ad_state'
  76. """,
  77. (os.environ["DB_NAME"],),
  78. )
  79. existing_columns = {row["COLUMN_NAME"] for row in cursor.fetchall()}
  80. migrations = {
  81. "initial_base_bid_fen": "INT DEFAULT NULL",
  82. "last_roi_scaled_at": "DATETIME DEFAULT NULL",
  83. "bid_hold": "BOOLEAN NOT NULL DEFAULT FALSE",
  84. "bid_hold_reason": "VARCHAR(255) DEFAULT NULL",
  85. "operator_pause_mode": "VARCHAR(32) DEFAULT NULL",
  86. "operator_resume_at": "DATETIME DEFAULT NULL",
  87. "operator_command_id": "VARCHAR(64) DEFAULT NULL",
  88. "operator_paused_from_status": "VARCHAR(50) DEFAULT NULL",
  89. "operator_paused_at": "DATETIME DEFAULT NULL",
  90. }
  91. for column, definition in migrations.items():
  92. if column not in existing_columns:
  93. cursor.execute(
  94. f"ALTER TABLE realtime_control_ad_state "
  95. f"ADD COLUMN {column} {definition}"
  96. )
  97. cursor.execute(
  98. """
  99. UPDATE realtime_control_ad_state
  100. SET initial_base_bid_fen=base_bid_fen
  101. WHERE initial_base_bid_fen IS NULL
  102. """
  103. )
  104. cursor.execute(
  105. """
  106. SELECT INDEX_NAME
  107. FROM information_schema.STATISTICS
  108. WHERE TABLE_SCHEMA=%s
  109. AND TABLE_NAME='realtime_control_ad_state'
  110. AND INDEX_NAME='idx_operator_pause'
  111. """,
  112. (os.environ["DB_NAME"],),
  113. )
  114. if not cursor.fetchone():
  115. cursor.execute(
  116. """
  117. CREATE INDEX idx_operator_pause
  118. ON realtime_control_ad_state
  119. (operator_pause_mode, operator_resume_at)
  120. """
  121. )
  122. operator_migrations = {
  123. "operator_command": {
  124. "raw_text": "MEDIUMTEXT DEFAULT NULL",
  125. "parse_source": "VARCHAR(32) NOT NULL DEFAULT 'deterministic'",
  126. "intent_json": "MEDIUMTEXT DEFAULT NULL",
  127. "preview_cost_fen": "BIGINT NOT NULL DEFAULT 0",
  128. "preview_impressions": "BIGINT NOT NULL DEFAULT 0",
  129. "preview_clicks": "BIGINT NOT NULL DEFAULT 0",
  130. "preview_conversions": "BIGINT NOT NULL DEFAULT 0",
  131. "previewed_at": "DATETIME DEFAULT NULL",
  132. "resume_at": "DATETIME DEFAULT NULL",
  133. },
  134. "operator_command_item": {
  135. "preview_cost_fen": "BIGINT NOT NULL DEFAULT 0",
  136. "preview_impressions": "BIGINT NOT NULL DEFAULT 0",
  137. "preview_clicks": "BIGINT NOT NULL DEFAULT 0",
  138. "preview_conversions": "BIGINT NOT NULL DEFAULT 0",
  139. "previewed_at": "DATETIME DEFAULT NULL",
  140. "updated_at": (
  141. "TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP "
  142. "ON UPDATE CURRENT_TIMESTAMP"
  143. ),
  144. },
  145. }
  146. for table_name, columns in operator_migrations.items():
  147. cursor.execute(
  148. """
  149. SELECT COLUMN_NAME
  150. FROM information_schema.COLUMNS
  151. WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s
  152. """,
  153. (os.environ["DB_NAME"], table_name),
  154. )
  155. table_columns = {row["COLUMN_NAME"] for row in cursor.fetchall()}
  156. for column, definition in columns.items():
  157. if column not in table_columns:
  158. cursor.execute(
  159. f"ALTER TABLE {table_name} ADD COLUMN {column} {definition}"
  160. )
  161. schema_migrations = {
  162. "roi_metric_run": {
  163. "fission_parameter_version": "VARCHAR(64) DEFAULT NULL",
  164. "fission_cohort_date": "DATE DEFAULT NULL",
  165. },
  166. "roi_entity_snapshot": {
  167. "conversion_goal": "VARCHAR(255) DEFAULT NULL",
  168. "fission_parameter_version": "VARCHAR(64) DEFAULT NULL",
  169. "fission_cohort_date": "VARCHAR(8) DEFAULT NULL",
  170. "fission_multiplier_vs_t0": "DECIMAL(18,8) DEFAULT NULL",
  171. "fission_multiplier_vs_first": "DECIMAL(18,8) DEFAULT NULL",
  172. "fission_match_level": "VARCHAR(64) DEFAULT NULL",
  173. "fission_source": "VARCHAR(512) DEFAULT NULL",
  174. "actual_total_revenue": "DECIMAL(20,4) DEFAULT NULL",
  175. "actual_roi": "DECIMAL(18,8) DEFAULT NULL",
  176. "predicted_tail_revenue": "DECIMAL(20,4) DEFAULT NULL",
  177. "predicted_fission_revenue": "DECIMAL(20,4) DEFAULT NULL",
  178. },
  179. "roi_fission_parameter_value": {
  180. "multiplier_vs_first": "DOUBLE DEFAULT NULL",
  181. },
  182. "roi_action_item": {
  183. "approval_status": "VARCHAR(32) NOT NULL DEFAULT 'PENDING'",
  184. "approval_source": "VARCHAR(32) DEFAULT NULL",
  185. "approved_at": "DATETIME DEFAULT NULL",
  186. "rejected_at": "DATETIME DEFAULT NULL",
  187. "sheet_row_number": "INT DEFAULT NULL",
  188. "result_notified_at": "DATETIME DEFAULT NULL",
  189. "notification_error": "TEXT DEFAULT NULL",
  190. },
  191. "roi_agency_delivery": {
  192. "creative_rows": "INT NOT NULL DEFAULT 0",
  193. "ad_rows": "INT NOT NULL DEFAULT 0",
  194. },
  195. "creative_rejection_cleanup_item": {
  196. "check_date": "DATE DEFAULT NULL",
  197. "agent_name": "VARCHAR(255) DEFAULT NULL",
  198. "cleanup_action": "VARCHAR(32) DEFAULT NULL",
  199. "cleanup_rule_type": (
  200. "VARCHAR(64) NOT NULL DEFAULT 'REVIEW_DENIED'"
  201. ),
  202. "target_component_ids_json": "LONGTEXT DEFAULT NULL",
  203. "target_element_ids_json": "LONGTEXT DEFAULT NULL",
  204. "recent_cost_fen": "BIGINT DEFAULT NULL",
  205. "current_day_cost_fen": "BIGINT DEFAULT NULL",
  206. "cost_start_date": "DATE DEFAULT NULL",
  207. "cost_end_date": "DATE DEFAULT NULL",
  208. "action_reason": "TEXT DEFAULT NULL",
  209. "creative_created_at": "DATETIME DEFAULT NULL",
  210. "creative_age_days": "INT DEFAULT NULL",
  211. "metric_impressions": "BIGINT DEFAULT NULL",
  212. "metric_daily_avg_impressions": (
  213. "DECIMAL(20, 4) DEFAULT NULL"
  214. ),
  215. "metric_window_days": "INT DEFAULT NULL",
  216. "agency_notified_at": "DATETIME DEFAULT NULL",
  217. "operator_notified_at": "DATETIME DEFAULT NULL",
  218. },
  219. "revenue_forecast_result": {
  220. "cost_reserve_date": "DATE DEFAULT NULL",
  221. "channel_costs_json": "TEXT DEFAULT NULL",
  222. "non_miniapp_reserved_cost": (
  223. "DECIMAL(20,4) NOT NULL DEFAULT 0"
  224. ),
  225. "miniapp_target_daily_cost": (
  226. "DECIMAL(20,4) NOT NULL DEFAULT 0"
  227. ),
  228. "forecast_method": (
  229. "VARCHAR(64) NOT NULL DEFAULT 'weighted_30m_speed'"
  230. ),
  231. "parameter_version": "VARCHAR(64) DEFAULT NULL",
  232. "current_cumulative_revenue": (
  233. "DECIMAL(20,4) DEFAULT NULL"
  234. ),
  235. "speed_latest_weight": "DECIMAL(12,8) DEFAULT NULL",
  236. "latest_interval_revenue": "DECIMAL(20,4) DEFAULT NULL",
  237. "previous_interval_revenue": "DECIMAL(20,4) DEFAULT NULL",
  238. "weighted_speed": "DECIMAL(20,4) DEFAULT NULL",
  239. "remaining_multiplier_p10": "DECIMAL(20,8) DEFAULT NULL",
  240. "remaining_multiplier_p50": "DECIMAL(20,8) DEFAULT NULL",
  241. "remaining_multiplier_p90": "DECIMAL(20,8) DEFAULT NULL",
  242. "parameter_sample_count": "INT DEFAULT NULL",
  243. },
  244. "revenue_forecast_observation": {
  245. "source_version": (
  246. "VARCHAR(64) NOT NULL DEFAULT "
  247. "'legacy'"
  248. ),
  249. },
  250. }
  251. for table_name, columns in schema_migrations.items():
  252. cursor.execute(
  253. """
  254. SELECT COLUMN_NAME
  255. FROM information_schema.COLUMNS
  256. WHERE TABLE_SCHEMA=%s AND TABLE_NAME=%s
  257. """,
  258. (os.environ["DB_NAME"], table_name),
  259. )
  260. table_columns = {
  261. row["COLUMN_NAME"] for row in cursor.fetchall()
  262. }
  263. for column, definition in columns.items():
  264. if column not in table_columns:
  265. cursor.execute(
  266. f"ALTER TABLE {table_name} "
  267. f"ADD COLUMN {column} {definition}"
  268. )
  269. _backfill_legacy_review_rule_types(cursor)
  270. _suppress_performance_agency_delivery(cursor)
  271. cursor.execute(
  272. """
  273. UPDATE creative_rejection_cleanup_item
  274. SET check_date=DATE(created_at)
  275. WHERE check_date IS NULL
  276. """
  277. )
  278. cursor.execute(
  279. """
  280. SELECT IS_NULLABLE
  281. FROM information_schema.COLUMNS
  282. WHERE TABLE_SCHEMA=%s
  283. AND TABLE_NAME='creative_rejection_cleanup_item'
  284. AND COLUMN_NAME='check_date'
  285. """,
  286. (os.environ["DB_NAME"],),
  287. )
  288. check_date_column = cursor.fetchone()
  289. if check_date_column and check_date_column["IS_NULLABLE"] == "YES":
  290. cursor.execute(
  291. """
  292. ALTER TABLE creative_rejection_cleanup_item
  293. MODIFY COLUMN check_date DATE NOT NULL
  294. """
  295. )
  296. cursor.execute(
  297. """
  298. SELECT COLUMN_NAME
  299. FROM information_schema.STATISTICS
  300. WHERE TABLE_SCHEMA=%s
  301. AND TABLE_NAME='creative_rejection_cleanup_item'
  302. AND INDEX_NAME='uk_creative_rejection_cleanup'
  303. ORDER BY SEQ_IN_INDEX
  304. """,
  305. (os.environ["DB_NAME"],),
  306. )
  307. cleanup_unique_columns = [
  308. row["COLUMN_NAME"] for row in cursor.fetchall()
  309. ]
  310. expected_cleanup_unique = [
  311. "account_id", "dynamic_creative_id", "check_date"
  312. ]
  313. if cleanup_unique_columns != expected_cleanup_unique:
  314. if cleanup_unique_columns:
  315. cursor.execute(
  316. """
  317. ALTER TABLE creative_rejection_cleanup_item
  318. DROP INDEX uk_creative_rejection_cleanup
  319. """
  320. )
  321. cursor.execute(
  322. """
  323. CREATE UNIQUE INDEX uk_creative_rejection_cleanup
  324. ON creative_rejection_cleanup_item
  325. (account_id, dynamic_creative_id, check_date)
  326. """
  327. )
  328. cursor.execute(
  329. """
  330. SELECT COLUMN_NAME, IS_NULLABLE
  331. FROM information_schema.COLUMNS
  332. WHERE TABLE_SCHEMA=%s
  333. AND TABLE_NAME='revenue_forecast_result'
  334. AND COLUMN_NAME IN (
  335. 'signal_a', 'trend_growths_json', 'signal_a_weight',
  336. 'signal_b_weight', 'confidence_ratio'
  337. )
  338. """,
  339. (os.environ["DB_NAME"],),
  340. )
  341. legacy_required_columns = {
  342. row["COLUMN_NAME"]
  343. for row in cursor.fetchall()
  344. if row["IS_NULLABLE"] == "NO"
  345. }
  346. legacy_definitions = {
  347. "signal_a": "DECIMAL(20,4) DEFAULT NULL",
  348. "trend_growths_json": "TEXT DEFAULT NULL",
  349. "signal_a_weight": "DECIMAL(8,6) DEFAULT NULL",
  350. "signal_b_weight": "DECIMAL(8,6) DEFAULT NULL",
  351. "confidence_ratio": "DECIMAL(8,6) DEFAULT NULL",
  352. }
  353. for column in sorted(legacy_required_columns):
  354. cursor.execute(
  355. f"ALTER TABLE revenue_forecast_result MODIFY COLUMN "
  356. f"{column} {legacy_definitions[column]}"
  357. )
  358. cursor.execute(
  359. """
  360. SELECT COLUMN_NAME, IS_NULLABLE
  361. FROM information_schema.COLUMNS
  362. WHERE TABLE_SCHEMA=%s
  363. AND TABLE_NAME='revenue_forecast_observation'
  364. AND COLUMN_NAME IN (
  365. 'yesterday_same_time_revenue',
  366. 'yesterday_total_revenue'
  367. )
  368. """,
  369. (os.environ["DB_NAME"],),
  370. )
  371. for row in cursor.fetchall():
  372. if row["IS_NULLABLE"] == "NO":
  373. cursor.execute(
  374. f"ALTER TABLE revenue_forecast_observation "
  375. f"MODIFY COLUMN {row['COLUMN_NAME']} "
  376. f"DECIMAL(20,4) DEFAULT NULL"
  377. )
  378. cursor.execute(
  379. """
  380. UPDATE revenue_forecast_observation o
  381. JOIN (
  382. SELECT report_time, MIN(forecast_version) AS forecast_version
  383. FROM revenue_forecast_result
  384. GROUP BY report_time
  385. HAVING COUNT(DISTINCT forecast_version) = 1
  386. ) r ON r.report_time = o.report_time
  387. LEFT JOIN revenue_forecast_observation exact_observation
  388. ON exact_observation.report_time = o.report_time
  389. AND exact_observation.source_version = r.forecast_version
  390. AND exact_observation.id <> o.id
  391. SET o.source_version = r.forecast_version
  392. WHERE o.source_version IN (
  393. 'legacy', 'revenue_forecast_v4_45m'
  394. )
  395. AND exact_observation.id IS NULL
  396. """
  397. )
  398. cursor.execute(
  399. """
  400. SELECT INDEX_NAME
  401. FROM information_schema.STATISTICS
  402. WHERE TABLE_SCHEMA=%s
  403. AND TABLE_NAME='revenue_forecast_observation'
  404. AND INDEX_NAME='uk_revenue_observation_report_time'
  405. LIMIT 1
  406. """,
  407. (os.environ["DB_NAME"],),
  408. )
  409. if cursor.fetchone():
  410. cursor.execute(
  411. """
  412. ALTER TABLE revenue_forecast_observation
  413. DROP INDEX uk_revenue_observation_report_time
  414. """
  415. )
  416. cursor.execute(
  417. """
  418. SELECT INDEX_NAME
  419. FROM information_schema.STATISTICS
  420. WHERE TABLE_SCHEMA=%s
  421. AND TABLE_NAME='revenue_forecast_observation'
  422. AND INDEX_NAME='uk_revenue_observation_version_time'
  423. LIMIT 1
  424. """,
  425. (os.environ["DB_NAME"],),
  426. )
  427. if not cursor.fetchone():
  428. cursor.execute(
  429. """
  430. CREATE UNIQUE INDEX uk_revenue_observation_version_time
  431. ON revenue_forecast_observation
  432. (report_time, source_version)
  433. """
  434. )
  435. cursor.execute(
  436. """
  437. SELECT INDEX_NAME
  438. FROM information_schema.STATISTICS
  439. WHERE TABLE_SCHEMA=%s
  440. AND TABLE_NAME='roi_action_item'
  441. AND INDEX_NAME='idx_roi_action_approval'
  442. """,
  443. (os.environ["DB_NAME"],),
  444. )
  445. if not cursor.fetchone():
  446. cursor.execute(
  447. """
  448. CREATE INDEX idx_roi_action_approval
  449. ON roi_action_item (run_id, approval_status)
  450. """
  451. )
  452. finally:
  453. connection.close()
  454. @contextmanager
  455. def advisory_lock(lock_name: str) -> Iterator[bool]:
  456. connection = connect()
  457. acquired = False
  458. try:
  459. with connection.cursor() as cursor:
  460. cursor.execute("SELECT GET_LOCK(%s, 0) AS acquired", (lock_name,))
  461. acquired = bool((cursor.fetchone() or {}).get("acquired"))
  462. yield acquired
  463. finally:
  464. if acquired:
  465. try:
  466. with connection.cursor() as cursor:
  467. cursor.execute("SELECT RELEASE_LOCK(%s)", (lock_name,))
  468. except Exception:
  469. pass
  470. connection.close()
  471. def load_enabled_accounts() -> list[dict[str, Any]]:
  472. connection = connect()
  473. try:
  474. with connection.cursor() as cursor:
  475. cursor.execute(
  476. """
  477. SELECT c.account_id, c.audience_name, c.bid_scene
  478. FROM ad_creation_account_config c
  479. JOIN account_whitelist w ON w.account_id = c.account_id
  480. WHERE c.enabled = TRUE
  481. AND w.enabled = TRUE
  482. ORDER BY c.account_id
  483. """
  484. )
  485. return list(cursor.fetchall())
  486. finally:
  487. connection.close()
  488. def load_automation_spend_accounts() -> list[dict[str, Any]]:
  489. """读取仍在白名单中的历史自动化账户。"""
  490. connection = connect()
  491. try:
  492. with connection.cursor() as cursor:
  493. cursor.execute(
  494. """
  495. SELECT c.account_id
  496. FROM ad_creation_account_config c
  497. JOIN account_whitelist w ON w.account_id = c.account_id
  498. WHERE w.enabled = TRUE
  499. ORDER BY c.account_id
  500. """
  501. )
  502. return list(cursor.fetchall())
  503. finally:
  504. connection.close()
  505. def load_all_spend_accounts() -> list[dict[str, Any]]:
  506. """读取本地账户白名单中所有启用账户。"""
  507. connection = connect()
  508. try:
  509. with connection.cursor() as cursor:
  510. cursor.execute(
  511. """
  512. SELECT w.account_id
  513. FROM account_whitelist w
  514. WHERE w.enabled = TRUE
  515. ORDER BY w.account_id
  516. """
  517. )
  518. return list(cursor.fetchall())
  519. finally:
  520. connection.close()
  521. def load_realtime_accounts() -> list[dict[str, Any]]:
  522. """读取完整调控的自动化账户与显式额外纳管账户。"""
  523. connection = connect()
  524. try:
  525. with connection.cursor() as cursor:
  526. cursor.execute(
  527. """
  528. SELECT c.account_id, c.audience_name, c.bid_scene,
  529. 'FULL' AS control_mode
  530. FROM ad_creation_account_config c
  531. JOIN account_whitelist w ON w.account_id = c.account_id
  532. WHERE w.enabled = TRUE
  533. ORDER BY c.account_id
  534. """
  535. )
  536. accounts = {
  537. int(row["account_id"]): row for row in cursor.fetchall()
  538. }
  539. cursor.execute(
  540. """
  541. SELECT account_id, audience_name, bid_scene, control_mode
  542. FROM realtime_control_account_scope
  543. WHERE enabled=TRUE
  544. ORDER BY account_id
  545. """
  546. )
  547. for row in cursor.fetchall():
  548. accounts.setdefault(int(row["account_id"]), row)
  549. return [accounts[key] for key in sorted(accounts)]
  550. finally:
  551. connection.close()
  552. def upsert_realtime_account_scope(
  553. *,
  554. account_id: int,
  555. control_mode: str,
  556. audience_name: str,
  557. bid_scene: str | None,
  558. source: str,
  559. note: str,
  560. ) -> None:
  561. if control_mode not in {"FULL", "PAUSE_ONLY"}:
  562. raise ValueError(f"Unsupported real-time control mode: {control_mode}")
  563. connection = connect()
  564. try:
  565. with connection.cursor() as cursor:
  566. cursor.execute(
  567. """
  568. INSERT INTO realtime_control_account_scope
  569. (account_id, control_mode, audience_name, bid_scene,
  570. enabled, source, note)
  571. VALUES (%s,%s,%s,%s,TRUE,%s,%s)
  572. ON DUPLICATE KEY UPDATE
  573. control_mode=VALUES(control_mode),
  574. audience_name=VALUES(audience_name),
  575. bid_scene=VALUES(bid_scene),
  576. enabled=TRUE,
  577. source=VALUES(source),
  578. note=VALUES(note)
  579. """,
  580. (
  581. account_id,
  582. control_mode,
  583. audience_name,
  584. bid_scene,
  585. source,
  586. note,
  587. ),
  588. )
  589. finally:
  590. connection.close()
  591. def load_managed_accounts() -> list[dict[str, Any]]:
  592. """返回账户白名单仍允许纳管的历史自动化账户。"""
  593. connection = connect()
  594. try:
  595. with connection.cursor() as cursor:
  596. cursor.execute(
  597. """
  598. SELECT c.account_id, w.account_name,
  599. c.audience_name, c.bid_scene, c.enabled
  600. FROM ad_creation_account_config c
  601. JOIN account_whitelist w ON w.account_id = c.account_id
  602. WHERE w.enabled = TRUE
  603. ORDER BY c.account_id
  604. """
  605. )
  606. return list(cursor.fetchall())
  607. finally:
  608. connection.close()
  609. def load_daily_state(control_date: date) -> dict[str, Any]:
  610. connection = connect()
  611. try:
  612. with connection.cursor() as cursor:
  613. cursor.execute(
  614. "SELECT * FROM realtime_control_daily_state WHERE control_date=%s",
  615. (control_date,),
  616. )
  617. return cursor.fetchone() or {}
  618. finally:
  619. connection.close()
  620. def save_daily_state(control_date: date, **values: Any) -> None:
  621. allowed = {
  622. "morning_recovery_done",
  623. "cutoff_done",
  624. "last_observed_partition",
  625. "last_observed_cpm",
  626. "last_decision",
  627. "last_inventory_refresh_at",
  628. "last_evaluated_at",
  629. }
  630. unknown = set(values) - allowed
  631. if unknown:
  632. raise ValueError(f"Unsupported daily-state fields: {sorted(unknown)}")
  633. columns = ["control_date", *values]
  634. params = [control_date, *values.values()]
  635. updates = ", ".join(f"{column}=VALUES({column})" for column in values)
  636. placeholders = ", ".join(["%s"] * len(columns))
  637. sql = (
  638. f"INSERT INTO realtime_control_daily_state ({', '.join(columns)}) "
  639. f"VALUES ({placeholders}) ON DUPLICATE KEY UPDATE {updates}"
  640. )
  641. connection = connect()
  642. try:
  643. with connection.cursor() as cursor:
  644. cursor.execute(sql, params)
  645. finally:
  646. connection.close()
  647. def load_ad_states(account_id: int) -> dict[int, dict[str, Any]]:
  648. connection = connect()
  649. try:
  650. with connection.cursor() as cursor:
  651. cursor.execute(
  652. "SELECT * FROM realtime_control_ad_state WHERE account_id=%s",
  653. (account_id,),
  654. )
  655. return {int(row["adgroup_id"]): row for row in cursor.fetchall()}
  656. finally:
  657. connection.close()
  658. def upsert_ad_state(
  659. *,
  660. account_id: int,
  661. adgroup_id: int,
  662. adgroup_name: str,
  663. bid_field: str,
  664. base_bid_fen: int,
  665. boosted_date: date | None,
  666. paused_by_strategy: bool,
  667. pause_reason: str | None,
  668. last_action: str,
  669. action_at: datetime,
  670. ) -> None:
  671. connection = connect()
  672. try:
  673. with connection.cursor() as cursor:
  674. cursor.execute(
  675. """
  676. INSERT INTO realtime_control_ad_state
  677. (account_id, adgroup_id, adgroup_name, bid_field, base_bid_fen,
  678. initial_base_bid_fen,
  679. boosted_date, paused_by_strategy, pause_reason, last_action,
  680. last_action_at, last_seen_at)
  681. VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
  682. ON DUPLICATE KEY UPDATE
  683. adgroup_name=VALUES(adgroup_name),
  684. bid_field=VALUES(bid_field),
  685. initial_base_bid_fen=COALESCE(
  686. initial_base_bid_fen,
  687. VALUES(initial_base_bid_fen)
  688. ),
  689. boosted_date=VALUES(boosted_date),
  690. paused_by_strategy=VALUES(paused_by_strategy),
  691. pause_reason=VALUES(pause_reason),
  692. last_action=VALUES(last_action),
  693. last_action_at=VALUES(last_action_at),
  694. last_seen_at=VALUES(last_seen_at)
  695. """,
  696. (
  697. account_id,
  698. adgroup_id,
  699. adgroup_name,
  700. bid_field,
  701. base_bid_fen,
  702. base_bid_fen,
  703. boosted_date,
  704. paused_by_strategy,
  705. pause_reason,
  706. last_action,
  707. action_at,
  708. action_at,
  709. ),
  710. )
  711. finally:
  712. connection.close()
  713. def sync_roi_base_bid(
  714. *,
  715. account_id: int,
  716. adgroup_id: int,
  717. adgroup_name: str,
  718. bid_field: str,
  719. initial_base_bid_fen: int,
  720. new_base_bid_fen: int,
  721. boosted_date: date | None,
  722. action_at: datetime,
  723. ) -> None:
  724. """持久化永久 ROI 基础出价,同时保留日内扩量状态。"""
  725. connection = connect()
  726. try:
  727. with connection.cursor() as cursor:
  728. cursor.execute(
  729. """
  730. INSERT INTO realtime_control_ad_state
  731. (account_id, adgroup_id, adgroup_name, bid_field,
  732. base_bid_fen, initial_base_bid_fen, boosted_date,
  733. paused_by_strategy, pause_reason, last_action,
  734. last_action_at, last_seen_at, last_roi_scaled_at)
  735. VALUES (%s,%s,%s,%s,%s,%s,%s,FALSE,NULL,'ROI_SCALE',%s,%s,%s)
  736. ON DUPLICATE KEY UPDATE
  737. adgroup_name=VALUES(adgroup_name),
  738. bid_field=VALUES(bid_field),
  739. base_bid_fen=VALUES(base_bid_fen),
  740. initial_base_bid_fen=COALESCE(
  741. initial_base_bid_fen,
  742. VALUES(initial_base_bid_fen)
  743. ),
  744. boosted_date=VALUES(boosted_date),
  745. last_action='ROI_SCALE',
  746. last_action_at=VALUES(last_action_at),
  747. last_seen_at=VALUES(last_seen_at),
  748. last_roi_scaled_at=VALUES(last_roi_scaled_at)
  749. """,
  750. (
  751. account_id,
  752. adgroup_id,
  753. adgroup_name,
  754. bid_field,
  755. new_base_bid_fen,
  756. initial_base_bid_fen,
  757. boosted_date,
  758. action_at,
  759. action_at,
  760. action_at,
  761. ),
  762. )
  763. finally:
  764. connection.close()
  765. def set_bid_experiment_hold(
  766. *,
  767. account_id: int,
  768. adgroup_id: int,
  769. adgroup_name: str,
  770. bid_field: str,
  771. base_bid_fen: int,
  772. action_at: datetime,
  773. reason: str,
  774. ) -> None:
  775. """持久化实验基础出价,并阻止实时调控覆盖该出价。"""
  776. connection = connect()
  777. try:
  778. with connection.cursor() as cursor:
  779. cursor.execute(
  780. """
  781. INSERT INTO realtime_control_ad_state
  782. (account_id, adgroup_id, adgroup_name, bid_field,
  783. base_bid_fen, initial_base_bid_fen, boosted_date,
  784. bid_hold, bid_hold_reason, paused_by_strategy,
  785. last_action, last_action_at, last_seen_at)
  786. VALUES (%s,%s,%s,%s,%s,%s,%s,TRUE,%s,FALSE,
  787. 'BID_EXPERIMENT_HOLD',%s,%s)
  788. ON DUPLICATE KEY UPDATE
  789. adgroup_name=VALUES(adgroup_name),
  790. bid_field=VALUES(bid_field),
  791. base_bid_fen=VALUES(base_bid_fen),
  792. initial_base_bid_fen=COALESCE(
  793. initial_base_bid_fen,
  794. VALUES(initial_base_bid_fen)
  795. ),
  796. boosted_date=VALUES(boosted_date),
  797. bid_hold=TRUE,
  798. bid_hold_reason=VALUES(bid_hold_reason),
  799. last_action='BID_EXPERIMENT_HOLD',
  800. last_action_at=VALUES(last_action_at),
  801. last_seen_at=VALUES(last_seen_at)
  802. """,
  803. (
  804. account_id,
  805. adgroup_id,
  806. adgroup_name,
  807. bid_field,
  808. base_bid_fen,
  809. base_bid_fen,
  810. action_at.date(),
  811. reason,
  812. action_at,
  813. action_at,
  814. ),
  815. )
  816. finally:
  817. connection.close()
  818. def clear_bid_experiment_hold(
  819. account_id: int,
  820. adgroup_id: int,
  821. *,
  822. action_at: datetime,
  823. ) -> None:
  824. connection = connect()
  825. try:
  826. with connection.cursor() as cursor:
  827. cursor.execute(
  828. """
  829. UPDATE realtime_control_ad_state
  830. SET bid_hold=FALSE,
  831. bid_hold_reason=NULL,
  832. last_action='BID_EXPERIMENT_RELEASE',
  833. last_action_at=%s,
  834. last_seen_at=%s
  835. WHERE account_id=%s AND adgroup_id=%s
  836. """,
  837. (action_at, action_at, account_id, adgroup_id),
  838. )
  839. if cursor.rowcount != 1:
  840. raise ValueError(
  841. "未找到实验出价状态: "
  842. f"account={account_id} adgroup={adgroup_id}"
  843. )
  844. finally:
  845. connection.close()
  846. def insert_action_log(record: dict[str, Any]) -> None:
  847. columns = [
  848. "run_id",
  849. "control_date",
  850. "observed_partition",
  851. "observed_cpm",
  852. "decision",
  853. "account_id",
  854. "adgroup_id",
  855. "adgroup_name",
  856. "bid_field",
  857. "base_bid_fen",
  858. "before_bid_fen",
  859. "target_bid_fen",
  860. "before_status",
  861. "target_status",
  862. "apply_mode",
  863. "execution_status",
  864. "error_message",
  865. ]
  866. connection = connect()
  867. try:
  868. with connection.cursor() as cursor:
  869. cursor.execute(
  870. f"INSERT INTO realtime_control_action_log ({', '.join(columns)}) "
  871. f"VALUES ({', '.join(['%s'] * len(columns))})",
  872. [record.get(column) for column in columns],
  873. )
  874. finally:
  875. connection.close()
  876. def create_operator_command(record: dict[str, Any]) -> dict[str, Any]:
  877. connection = connect()
  878. try:
  879. with connection.cursor() as cursor:
  880. try:
  881. cursor.execute(
  882. """
  883. INSERT INTO operator_command
  884. (command_id, source_message_id, chat_id, sender_open_id,
  885. sender_name, action, scope_type, target_account_ids,
  886. status, preview_account_count, preview_ad_count, expires_at)
  887. VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
  888. """,
  889. (
  890. record["command_id"],
  891. record["source_message_id"],
  892. record["chat_id"],
  893. record["sender_open_id"],
  894. record.get("sender_name"),
  895. record["action"],
  896. record["scope_type"],
  897. json.dumps(record["target_account_ids"]),
  898. record["status"],
  899. record.get("preview_account_count", 0),
  900. record.get("preview_ad_count", 0),
  901. record.get("expires_at"),
  902. ),
  903. )
  904. except pymysql.err.IntegrityError:
  905. cursor.execute(
  906. "SELECT * FROM operator_command WHERE source_message_id=%s",
  907. (record["source_message_id"],),
  908. )
  909. existing = cursor.fetchone()
  910. if existing:
  911. existing["target_account_ids"] = json.loads(
  912. existing.get("target_account_ids") or "[]"
  913. )
  914. return existing
  915. raise
  916. return load_operator_command(record["command_id"]) or {}
  917. finally:
  918. connection.close()
  919. def create_operator_command_with_items(
  920. record: dict[str, Any],
  921. items: list[dict[str, Any]],
  922. ) -> dict[str, Any]:
  923. """原子保存不可变的命令预览及其广告快照。"""
  924. connection = connect()
  925. try:
  926. connection.begin()
  927. with connection.cursor() as cursor:
  928. try:
  929. cursor.execute(
  930. """
  931. INSERT INTO operator_command
  932. (command_id, source_message_id, chat_id, sender_open_id,
  933. sender_name, raw_text, parse_source, intent_json,
  934. action, scope_type, target_account_ids, status,
  935. preview_account_count, preview_ad_count,
  936. preview_cost_fen, preview_impressions, preview_clicks,
  937. preview_conversions, previewed_at, resume_at, expires_at)
  938. VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s,%s)
  939. """,
  940. (
  941. record["command_id"], record["source_message_id"],
  942. record["chat_id"], record["sender_open_id"],
  943. record.get("sender_name"), record.get("raw_text"),
  944. record.get("parse_source", "deterministic"),
  945. json.dumps(record.get("intent") or {}, ensure_ascii=False),
  946. record["action"], record["scope_type"],
  947. json.dumps(record["target_account_ids"]), record["status"],
  948. record.get("preview_account_count", 0),
  949. record.get("preview_ad_count", 0),
  950. record.get("preview_cost_fen", 0),
  951. record.get("preview_impressions", 0),
  952. record.get("preview_clicks", 0),
  953. record.get("preview_conversions", 0),
  954. record.get("previewed_at"), record.get("resume_at"),
  955. record.get("expires_at"),
  956. ),
  957. )
  958. except pymysql.err.IntegrityError:
  959. connection.rollback()
  960. existing = load_operator_command_by_source(record["source_message_id"])
  961. if existing:
  962. return existing
  963. raise
  964. columns = [
  965. "command_id", "account_id", "audience_name", "adgroup_id",
  966. "adgroup_name", "before_status", "target_status",
  967. "preview_cost_fen", "preview_impressions", "preview_clicks",
  968. "preview_conversions", "previewed_at", "execution_status",
  969. ]
  970. for item in items:
  971. cursor.execute(
  972. f"INSERT INTO operator_command_item ({', '.join(columns)}) "
  973. f"VALUES ({', '.join(['%s'] * len(columns))})",
  974. [item.get(column) for column in columns],
  975. )
  976. connection.commit()
  977. return load_operator_command(record["command_id"]) or {}
  978. except Exception:
  979. connection.rollback()
  980. raise
  981. finally:
  982. connection.close()
  983. def load_operator_command_by_source(source_message_id: str) -> dict[str, Any] | None:
  984. connection = connect()
  985. try:
  986. with connection.cursor() as cursor:
  987. cursor.execute(
  988. "SELECT command_id FROM operator_command WHERE source_message_id=%s",
  989. (source_message_id,),
  990. )
  991. row = cursor.fetchone()
  992. return load_operator_command(row["command_id"]) if row else None
  993. finally:
  994. connection.close()
  995. def load_operator_command(command_id: str) -> dict[str, Any] | None:
  996. connection = connect()
  997. try:
  998. with connection.cursor() as cursor:
  999. cursor.execute(
  1000. "SELECT * FROM operator_command WHERE command_id=%s",
  1001. (command_id,),
  1002. )
  1003. row = cursor.fetchone()
  1004. if row:
  1005. row["target_account_ids"] = json.loads(
  1006. row.get("target_account_ids") or "[]"
  1007. )
  1008. row["intent"] = json.loads(row.get("intent_json") or "{}")
  1009. return row
  1010. finally:
  1011. connection.close()
  1012. def update_operator_command(command_id: str, status: str, **values: Any) -> None:
  1013. allowed = {"confirmed_at", "executed_at", "error_message"}
  1014. unknown = set(values) - allowed
  1015. if unknown:
  1016. raise ValueError(f"Unsupported operator-command fields: {sorted(unknown)}")
  1017. assignments = ["status=%s"]
  1018. params: list[Any] = [status]
  1019. for column, value in values.items():
  1020. assignments.append(f"{column}=%s")
  1021. params.append(value)
  1022. params.append(command_id)
  1023. connection = connect()
  1024. try:
  1025. with connection.cursor() as cursor:
  1026. cursor.execute(
  1027. f"UPDATE operator_command SET {', '.join(assignments)} "
  1028. "WHERE command_id=%s",
  1029. params,
  1030. )
  1031. finally:
  1032. connection.close()
  1033. def transition_operator_command(
  1034. command_id: str,
  1035. *,
  1036. expected_statuses: set[str],
  1037. target_status: str,
  1038. **values: Any,
  1039. ) -> bool:
  1040. allowed = {"confirmed_at", "executed_at", "error_message"}
  1041. unknown = set(values) - allowed
  1042. if unknown:
  1043. raise ValueError(f"Unsupported operator-command fields: {sorted(unknown)}")
  1044. if not expected_statuses:
  1045. raise ValueError("expected_statuses must not be empty")
  1046. assignments = ["status=%s"]
  1047. params: list[Any] = [target_status]
  1048. for column, value in values.items():
  1049. assignments.append(f"{column}=%s")
  1050. params.append(value)
  1051. placeholders = ", ".join(["%s"] * len(expected_statuses))
  1052. params.extend([command_id, *sorted(expected_statuses)])
  1053. connection = connect()
  1054. try:
  1055. with connection.cursor() as cursor:
  1056. affected = cursor.execute(
  1057. f"""
  1058. UPDATE operator_command
  1059. SET {', '.join(assignments)}
  1060. WHERE command_id=%s
  1061. AND status IN ({placeholders})
  1062. """,
  1063. params,
  1064. )
  1065. return affected == 1
  1066. finally:
  1067. connection.close()
  1068. def insert_operator_command_item(record: dict[str, Any]) -> None:
  1069. columns = [
  1070. "command_id",
  1071. "account_id",
  1072. "audience_name",
  1073. "adgroup_id",
  1074. "adgroup_name",
  1075. "before_status",
  1076. "target_status",
  1077. "readback_status",
  1078. "execution_status",
  1079. "error_message",
  1080. ]
  1081. connection = connect()
  1082. try:
  1083. with connection.cursor() as cursor:
  1084. cursor.execute(
  1085. f"INSERT INTO operator_command_item ({', '.join(columns)}) "
  1086. f"VALUES ({', '.join(['%s'] * len(columns))})",
  1087. [record.get(column) for column in columns],
  1088. )
  1089. finally:
  1090. connection.close()
  1091. def load_operator_command_items(command_id: str) -> list[dict[str, Any]]:
  1092. connection = connect()
  1093. try:
  1094. with connection.cursor() as cursor:
  1095. cursor.execute(
  1096. """
  1097. SELECT * FROM operator_command_item
  1098. WHERE command_id=%s
  1099. ORDER BY account_id, adgroup_id, id
  1100. """,
  1101. (command_id,),
  1102. )
  1103. return list(cursor.fetchall())
  1104. finally:
  1105. connection.close()
  1106. def update_operator_command_item(item_id: int, **values: Any) -> None:
  1107. allowed = {
  1108. "adgroup_id", "adgroup_name", "before_status", "target_status",
  1109. "readback_status", "execution_status", "error_message",
  1110. }
  1111. unknown = set(values) - allowed
  1112. if unknown:
  1113. raise ValueError(f"Unsupported operator-command item fields: {sorted(unknown)}")
  1114. if not values:
  1115. return
  1116. assignments = [f"{column}=%s" for column in values]
  1117. connection = connect()
  1118. try:
  1119. with connection.cursor() as cursor:
  1120. cursor.execute(
  1121. f"UPDATE operator_command_item SET {', '.join(assignments)} WHERE id=%s",
  1122. [*values.values(), item_id],
  1123. )
  1124. finally:
  1125. connection.close()
  1126. def find_pending_command_conflict(
  1127. targets: list[tuple[int, int]],
  1128. *,
  1129. now: datetime,
  1130. ) -> dict[str, Any] | None:
  1131. if not targets:
  1132. return None
  1133. connection = connect()
  1134. try:
  1135. with connection.cursor() as cursor:
  1136. for start in range(0, len(targets), 200):
  1137. chunk = targets[start:start + 200]
  1138. conditions = " OR ".join(
  1139. ["(i.account_id=%s AND i.adgroup_id=%s)"] * len(chunk)
  1140. )
  1141. params: list[Any] = []
  1142. for account_id, adgroup_id in chunk:
  1143. params.extend([account_id, adgroup_id])
  1144. params.append(now.replace(tzinfo=None))
  1145. cursor.execute(
  1146. f"""
  1147. SELECT c.command_id, c.action, i.account_id, i.adgroup_id
  1148. FROM operator_command c
  1149. JOIN operator_command_item i ON i.command_id=c.command_id
  1150. WHERE ({conditions})
  1151. AND (
  1152. (c.status='PENDING_CONFIRMATION' AND c.expires_at >= %s)
  1153. OR c.status='EXECUTING'
  1154. )
  1155. ORDER BY c.created_at
  1156. LIMIT 1
  1157. """,
  1158. params,
  1159. )
  1160. conflict = cursor.fetchone()
  1161. if conflict:
  1162. return conflict
  1163. return None
  1164. finally:
  1165. connection.close()
  1166. def list_pending_operator_commands(
  1167. chat_id: str,
  1168. sender_open_id: str,
  1169. now: datetime,
  1170. ) -> list[dict[str, Any]]:
  1171. connection = connect()
  1172. try:
  1173. with connection.cursor() as cursor:
  1174. cursor.execute(
  1175. """
  1176. SELECT command_id, action, preview_account_count, preview_ad_count,
  1177. preview_cost_fen, expires_at
  1178. FROM operator_command
  1179. WHERE chat_id=%s AND sender_open_id=%s
  1180. AND status='PENDING_CONFIRMATION' AND expires_at >= %s
  1181. ORDER BY created_at
  1182. """,
  1183. (chat_id, sender_open_id, now.replace(tzinfo=None)),
  1184. )
  1185. return list(cursor.fetchall())
  1186. finally:
  1187. connection.close()
  1188. def load_active_operator_draft(
  1189. chat_id: str,
  1190. sender_open_id: str,
  1191. now: datetime,
  1192. ) -> dict[str, Any] | None:
  1193. connection = connect()
  1194. try:
  1195. with connection.cursor() as cursor:
  1196. cursor.execute(
  1197. """
  1198. SELECT * FROM operator_command_draft
  1199. WHERE chat_id=%s AND sender_open_id=%s
  1200. AND status='ACTIVE' AND expires_at >= %s
  1201. """,
  1202. (chat_id, sender_open_id, now.replace(tzinfo=None)),
  1203. )
  1204. row = cursor.fetchone()
  1205. if row:
  1206. for field in ("account_ids", "missing_fields", "source_message_ids"):
  1207. row[field] = json.loads(row.get(field) or "[]")
  1208. return row
  1209. finally:
  1210. connection.close()
  1211. def save_operator_draft(record: dict[str, Any]) -> dict[str, Any]:
  1212. connection = connect()
  1213. try:
  1214. with connection.cursor() as cursor:
  1215. cursor.execute(
  1216. """
  1217. INSERT INTO operator_command_draft
  1218. (draft_id, chat_id, sender_open_id, raw_text, action, scope_type,
  1219. account_ids, missing_fields, source_message_ids, status, expires_at)
  1220. VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s,'ACTIVE',%s)
  1221. ON DUPLICATE KEY UPDATE
  1222. draft_id=VALUES(draft_id), raw_text=VALUES(raw_text),
  1223. action=VALUES(action),
  1224. scope_type=VALUES(scope_type), account_ids=VALUES(account_ids),
  1225. missing_fields=VALUES(missing_fields),
  1226. source_message_ids=VALUES(source_message_ids),
  1227. status='ACTIVE', expires_at=VALUES(expires_at)
  1228. """,
  1229. (
  1230. record["draft_id"], record["chat_id"], record["sender_open_id"],
  1231. record.get("raw_text"), record.get("action"),
  1232. record.get("scope_type", "MISSING"),
  1233. json.dumps(record.get("account_ids") or []),
  1234. json.dumps(record.get("missing_fields") or []),
  1235. json.dumps(record.get("source_message_ids") or []),
  1236. record["expires_at"],
  1237. ),
  1238. )
  1239. return record
  1240. finally:
  1241. connection.close()
  1242. def close_operator_draft(chat_id: str, sender_open_id: str, status: str) -> None:
  1243. if status not in {"COMPLETED", "CANCELLED", "SUPERSEDED", "EXPIRED"}:
  1244. raise ValueError(f"Unsupported draft status: {status}")
  1245. connection = connect()
  1246. try:
  1247. with connection.cursor() as cursor:
  1248. cursor.execute(
  1249. """
  1250. UPDATE operator_command_draft SET status=%s
  1251. WHERE chat_id=%s AND sender_open_id=%s AND status='ACTIVE'
  1252. """,
  1253. (status, chat_id, sender_open_id),
  1254. )
  1255. finally:
  1256. connection.close()
  1257. def set_operator_pause(
  1258. *,
  1259. account_id: int,
  1260. adgroup_id: int,
  1261. mode: str,
  1262. resume_at: datetime | None,
  1263. command_id: str,
  1264. paused_from_status: str,
  1265. paused_at: datetime,
  1266. ) -> None:
  1267. connection = connect()
  1268. try:
  1269. with connection.cursor() as cursor:
  1270. cursor.execute(
  1271. """
  1272. UPDATE realtime_control_ad_state
  1273. SET operator_pause_mode=%s,
  1274. operator_resume_at=%s,
  1275. operator_command_id=%s,
  1276. operator_paused_from_status=%s,
  1277. operator_paused_at=%s,
  1278. last_action='OPERATOR_PAUSE',
  1279. last_action_at=%s,
  1280. last_seen_at=%s
  1281. WHERE account_id=%s AND adgroup_id=%s
  1282. """,
  1283. (
  1284. mode,
  1285. resume_at,
  1286. command_id,
  1287. paused_from_status,
  1288. paused_at,
  1289. paused_at,
  1290. paused_at,
  1291. account_id,
  1292. adgroup_id,
  1293. ),
  1294. )
  1295. finally:
  1296. connection.close()
  1297. def clear_operator_pause(
  1298. account_id: int,
  1299. adgroup_id: int,
  1300. *,
  1301. action: str,
  1302. action_at: datetime,
  1303. ) -> None:
  1304. connection = connect()
  1305. try:
  1306. with connection.cursor() as cursor:
  1307. cursor.execute(
  1308. """
  1309. UPDATE realtime_control_ad_state
  1310. SET operator_pause_mode=NULL,
  1311. operator_resume_at=NULL,
  1312. operator_command_id=NULL,
  1313. operator_paused_from_status=NULL,
  1314. operator_paused_at=NULL,
  1315. last_action=%s,
  1316. last_action_at=%s,
  1317. last_seen_at=%s
  1318. WHERE account_id=%s AND adgroup_id=%s
  1319. """,
  1320. (action, action_at, action_at, account_id, adgroup_id),
  1321. )
  1322. finally:
  1323. connection.close()
  1324. def load_operator_pauses(
  1325. account_ids: list[int] | None = None,
  1326. ) -> list[dict[str, Any]]:
  1327. params: list[Any] = []
  1328. where = "WHERE operator_pause_mode IS NOT NULL"
  1329. if account_ids is not None:
  1330. if not account_ids:
  1331. return []
  1332. where += f" AND account_id IN ({', '.join(['%s'] * len(account_ids))})"
  1333. params.extend(account_ids)
  1334. connection = connect()
  1335. try:
  1336. with connection.cursor() as cursor:
  1337. cursor.execute(
  1338. f"""
  1339. SELECT *
  1340. FROM realtime_control_ad_state
  1341. {where}
  1342. ORDER BY account_id, adgroup_id
  1343. """,
  1344. params,
  1345. )
  1346. return list(cursor.fetchall())
  1347. finally:
  1348. connection.close()