storage.py 50 KB

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