storage.py 52 KB

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