test_data_rule_polars_execution.py 98 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542
  1. from __future__ import annotations
  2. import hashlib
  3. import json
  4. import re
  5. from datetime import UTC, datetime, timedelta
  6. from pathlib import Path
  7. import polars as pl
  8. import pytest
  9. from minio import Minio
  10. from sqlalchemy import create_engine, text
  11. from app.core.common.identifiers import new_governance_uid
  12. from app.core.data_rules.contracts import rule_spec_hash, validate_rule_spec
  13. from app.core.data_rules.execution_contracts import canonical_schema_hash
  14. COMPOSE = (
  15. Path(__file__).resolve().parents[2]
  16. / "deploy"
  17. / "docker"
  18. / "docker-compose.yml"
  19. )
  20. def _compose_value(pattern):
  21. source = COMPOSE.read_text(encoding="utf-8")
  22. source = re.sub(r"\$\{[A-Za-z_][A-Za-z0-9_]*:-([^}]+)\}", r"\1", source)
  23. match = re.search(pattern, source, flags=re.DOTALL)
  24. assert match is not None
  25. return match.group(1)
  26. def _schema(schema_ref, fields):
  27. normalized = [
  28. {"name": name, "type": field_type, "nullable": nullable}
  29. for name, field_type, nullable in fields
  30. ]
  31. return {
  32. "id": new_governance_uid(),
  33. "schema_ref": schema_ref,
  34. "schema_hash": canonical_schema_hash(normalized),
  35. "fields": normalized,
  36. "source_revision": "task5:real-cross-source",
  37. }
  38. def _binding(schema, *, source_uid, access_mode, object_ref):
  39. return {
  40. "id": new_governance_uid(),
  41. "data_source_uid": source_uid,
  42. "object_kind": "parquet_artifact",
  43. "object_ref": object_ref,
  44. "schema_snapshot_id": schema["id"],
  45. "access_mode": access_mode,
  46. "dialect": "parquet",
  47. "write_mode": "append",
  48. }
  49. def test_real_postgres_mysql_minio_polars_cross_source_execution(tmp_path):
  50. from app.core.data_rules.compilers.polars import PolarsRuleCompiler
  51. from app.runner.artifacts import ArtifactStore, PostgresArtifactResolver
  52. from app.runner.rule_polars import PolarsRulePlanAdapter
  53. from app.runner.rules import PostgresRulePlanRepository, RulePlanExecutor
  54. source_user = _compose_value(
  55. r"source-postgres:.*?POSTGRES_USER:\s*([^\s]+)"
  56. )
  57. source_password = _compose_value(
  58. r"source-postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
  59. )
  60. platform_user = _compose_value(
  61. r"\n postgres:.*?POSTGRES_USER:\s*([^\s]+)"
  62. )
  63. platform_password = _compose_value(
  64. r"\n postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
  65. )
  66. postgres_port = _compose_value(r'"(25432):5432"')
  67. mysql_port = _compose_value(r'"(23306):3306"')
  68. minio_user = _compose_value(r"MINIO_ROOT_USER:\s*([^\s]+)")
  69. minio_password = _compose_value(r"MINIO_ROOT_PASSWORD:\s*([^\s]+)")
  70. minio_port = _compose_value(r'"(19000):9000"')
  71. platform_port = _compose_value(r'"(15432):5432"')
  72. bucket = _compose_value(r"mc mb --ignore-existing local/([^\s]+)")
  73. postgres = create_engine(
  74. f"postgresql+psycopg2://{source_user}:{source_password}"
  75. f"@127.0.0.1:{postgres_port}/acceptance",
  76. pool_pre_ping=True,
  77. )
  78. mysql = create_engine(
  79. f"mysql+pymysql://{source_user}:{source_password}"
  80. f"@127.0.0.1:{mysql_port}/acceptance",
  81. pool_pre_ping=True,
  82. )
  83. platform = create_engine(
  84. f"postgresql+psycopg2://{platform_user}:{platform_password}"
  85. f"@127.0.0.1:{platform_port}/dataops",
  86. pool_pre_ping=True,
  87. )
  88. minio = Minio(
  89. f"127.0.0.1:{minio_port}",
  90. access_key=minio_user,
  91. secret_key=minio_password,
  92. secure=False,
  93. )
  94. store = ArtifactStore(
  95. minio,
  96. bucket=bucket,
  97. max_artifact_bytes=4 * 1024 * 1024,
  98. max_rows=1_000,
  99. memory_limit_bytes=256 * 1024 * 1024,
  100. max_ttl_seconds=3600,
  101. )
  102. correlation_id = new_governance_uid()
  103. failure_correlation_id = new_governance_uid()
  104. unknown_correlation_id = new_governance_uid()
  105. lease_correlation_id = new_governance_uid()
  106. sample_crash_correlation_id = new_governance_uid()
  107. receipt_correlation_id = new_governance_uid()
  108. failed_receipt_correlation_id = new_governance_uid()
  109. no_run_crash_correlation_id = new_governance_uid()
  110. evidence_crash_correlation_id = new_governance_uid()
  111. active_retry_correlation_id = new_governance_uid()
  112. sql_binding_id = new_governance_uid()
  113. prefix = f"rules/{correlation_id}/"
  114. customer_table = "task5_polars_customers"
  115. segment_table = "task5_polars_segments"
  116. rule_uid = new_governance_uid()
  117. rule_id = new_governance_uid()
  118. downstream_rule_uid = new_governance_uid()
  119. downstream_rule_id = new_governance_uid()
  120. dataflow_uid = new_governance_uid()
  121. dataflow_version_id = new_governance_uid()
  122. deployment_id = new_governance_uid()
  123. component_binding_id = new_governance_uid()
  124. plan_id = new_governance_uid()
  125. downstream_component_binding_id = new_governance_uid()
  126. downstream_plan_id = new_governance_uid()
  127. ledger_jti = None
  128. retry_ledger_jti = None
  129. no_run_crash_jti = None
  130. evidence_crash_jti = None
  131. active_retry_jti = None
  132. gateway_candidate_id = None
  133. try:
  134. with postgres.begin() as connection:
  135. connection.execute(text(f"DROP TABLE IF EXISTS {customer_table}"))
  136. connection.execute(
  137. text(
  138. f"CREATE TABLE {customer_table} ("
  139. "customer_id BIGINT NOT NULL, "
  140. "name VARCHAR(100), mobile VARCHAR(30), "
  141. "segment_code VARCHAR(20), version_no BIGINT NOT NULL)"
  142. )
  143. )
  144. connection.execute(
  145. text(
  146. f"INSERT INTO {customer_table} "
  147. "(customer_id, name, mobile, segment_code, version_no) "
  148. "VALUES "
  149. "(1, ' Alice ', '13800138000', 'A', 1), "
  150. "(1, ' Alice Updated ', '13800138000', 'A', 2), "
  151. "(2, ' Bad ', 'invalid', 'B', 1), "
  152. "(3, ' Carol ', '13900139000', 'C', 1)"
  153. )
  154. )
  155. with mysql.begin() as connection:
  156. connection.execute(text(f"DROP TABLE IF EXISTS {segment_table}"))
  157. connection.execute(
  158. text(
  159. f"CREATE TABLE {segment_table} ("
  160. "code VARCHAR(20) PRIMARY KEY, "
  161. "segment_name VARCHAR(100) NOT NULL)"
  162. )
  163. )
  164. connection.execute(
  165. text(
  166. f"INSERT INTO {segment_table} (code, segment_name) "
  167. "VALUES ('A', 'Gold'), ('B', 'Basic'), ('C', 'Silver')"
  168. )
  169. )
  170. with postgres.connect() as connection:
  171. customer_rows = [
  172. dict(row)
  173. for row in connection.execute(
  174. text(
  175. f"SELECT customer_id, name, mobile, "
  176. f"segment_code, version_no FROM {customer_table}"
  177. )
  178. ).mappings()
  179. ]
  180. with mysql.connect() as connection:
  181. segment_rows = [
  182. dict(row)
  183. for row in connection.execute(
  184. text(
  185. f"SELECT code, segment_name FROM {segment_table}"
  186. )
  187. ).mappings()
  188. ]
  189. input_schema = _schema(
  190. "bd:task5:customer:raw",
  191. [
  192. ("customer_id", "integer", False),
  193. ("name", "string", True),
  194. ("mobile", "string", True),
  195. ("segment_code", "string", True),
  196. ("version_no", "integer", False),
  197. ],
  198. )
  199. lookup_schema = _schema(
  200. "bd:task5:segment:lookup",
  201. [
  202. ("code", "string", False),
  203. ("segment_name", "string", False),
  204. ],
  205. )
  206. output_schema = _schema(
  207. "bd:task5:customer:enriched",
  208. [
  209. ("customer_id", "integer", False),
  210. ("name", "string", True),
  211. ("mobile", "string", True),
  212. ("segment_code", "string", True),
  213. ("version_no", "integer", False),
  214. ("segment_name", "string", True),
  215. ],
  216. )
  217. input_binding = _binding(
  218. input_schema,
  219. source_uid=new_governance_uid(),
  220. access_mode="read",
  221. object_ref="postgres-customer-artifact",
  222. )
  223. lookup_binding = _binding(
  224. lookup_schema,
  225. source_uid=new_governance_uid(),
  226. access_mode="read",
  227. object_ref="mysql-segment-artifact",
  228. )
  229. output_binding = _binding(
  230. output_schema,
  231. source_uid=new_governance_uid(),
  232. access_mode="read_write",
  233. object_ref="polars-output-artifact",
  234. )
  235. downstream_output_binding = _binding(
  236. output_schema,
  237. source_uid=new_governance_uid(),
  238. access_mode="write",
  239. object_ref="polars-downstream-output-artifact",
  240. )
  241. spec = validate_rule_spec(
  242. {
  243. "schema_version": "2.0",
  244. "rule_uid": new_governance_uid(),
  245. "name": "task5_real_cross_source",
  246. "input_schema_ref": input_schema["schema_ref"],
  247. "output_schema_ref": output_schema["schema_ref"],
  248. "steps": [
  249. {
  250. "id": "normalize_name",
  251. "op": "normalize_text",
  252. "column": "name",
  253. "trim": True,
  254. },
  255. {
  256. "id": "join_segment",
  257. "op": "lookup_join",
  258. "lookup": {
  259. "binding_id": lookup_binding["id"],
  260. "left_on": ["segment_code"],
  261. "right_on": ["code"],
  262. "select": {
  263. "segment_name": "segment_name"
  264. },
  265. "how": "left",
  266. },
  267. },
  268. {
  269. "id": "valid_mobile",
  270. "op": "assert",
  271. "expression": "matches(mobile, '^[0-9]{11}$')",
  272. "on_failure": "reject",
  273. "severity": "error",
  274. },
  275. {
  276. "id": "latest_customer",
  277. "op": "deduplicate",
  278. "keys": ["customer_id"],
  279. "order_by": ["version_no"],
  280. "keep": "last",
  281. },
  282. ],
  283. "null_policy": "explicit",
  284. "timezone": "Asia/Shanghai",
  285. }
  286. )
  287. rule = {
  288. "id": rule_id,
  289. "status": "published",
  290. "rule_spec": spec,
  291. "spec_hash": rule_spec_hash(spec),
  292. }
  293. compiled = PolarsRuleCompiler().compile(
  294. rule_version=rule,
  295. input_schema=input_schema,
  296. output_schema=output_schema,
  297. input_binding=input_binding,
  298. output_binding=output_binding,
  299. backend={
  300. "max_rows": 1_000,
  301. "max_artifact_bytes": 4 * 1024 * 1024,
  302. "memory_limit_bytes": 256 * 1024 * 1024,
  303. "masking_policies": {},
  304. "lookup_bindings": {
  305. lookup_binding["id"]: {
  306. "binding": lookup_binding,
  307. "schema": lookup_schema,
  308. }
  309. },
  310. },
  311. )
  312. downstream_spec = validate_rule_spec(
  313. {
  314. "schema_version": "2.0",
  315. "rule_uid": new_governance_uid(),
  316. "name": "task6_real_artifact_handoff",
  317. "input_schema_ref": output_schema["schema_ref"],
  318. "output_schema_ref": output_schema["schema_ref"],
  319. "steps": [
  320. {
  321. "id": "normalize_downstream_name",
  322. "op": "normalize_text",
  323. "column": "name",
  324. "trim": True,
  325. }
  326. ],
  327. "null_policy": "explicit",
  328. "timezone": "Asia/Shanghai",
  329. }
  330. )
  331. downstream_rule = {
  332. "id": downstream_rule_id,
  333. "status": "published",
  334. "rule_spec": downstream_spec,
  335. "spec_hash": rule_spec_hash(downstream_spec),
  336. }
  337. downstream_compiled = PolarsRuleCompiler().compile(
  338. rule_version=downstream_rule,
  339. input_schema=output_schema,
  340. output_schema=output_schema,
  341. input_binding=output_binding,
  342. output_binding=downstream_output_binding,
  343. backend={
  344. "max_rows": 1_000,
  345. "max_artifact_bytes": 4 * 1024 * 1024,
  346. "memory_limit_bytes": 256 * 1024 * 1024,
  347. "masking_policies": {},
  348. "lookup_bindings": {},
  349. },
  350. )
  351. lookup_operation = compiled["plan"]["operations"][1]
  352. schema_hashes = {
  353. "rule_spec_hash": compiled["plan"]["rule_spec_hash"],
  354. "input_schema_snapshot_id": input_schema["id"],
  355. "input_schema_hash": input_schema["schema_hash"],
  356. "output_schema_snapshot_id": output_schema["id"],
  357. "output_schema_hash": output_schema["schema_hash"],
  358. }
  359. with platform.begin() as connection:
  360. for schema in (input_schema, lookup_schema, output_schema):
  361. connection.execute(
  362. text(
  363. """
  364. INSERT INTO public.data_schema_snapshots
  365. (id, schema_ref, schema_hash, fields, source_revision)
  366. VALUES (CAST(:id AS uuid), :schema_ref, :schema_hash,
  367. CAST(:fields AS jsonb), :source_revision)
  368. """
  369. ),
  370. {**schema, "fields": json.dumps(schema["fields"])},
  371. )
  372. connection.execute(
  373. text(
  374. """
  375. INSERT INTO public.data_rules
  376. (id, rule_uid, name, category, status)
  377. VALUES (CAST(:id AS uuid), CAST(:rule_uid AS uuid),
  378. :name, 'general', 'active')
  379. """
  380. ),
  381. {
  382. "id": new_governance_uid(),
  383. "rule_uid": rule_uid,
  384. "name": spec["name"],
  385. },
  386. )
  387. connection.execute(
  388. text(
  389. """
  390. INSERT INTO public.data_rules
  391. (id, rule_uid, name, category, status)
  392. VALUES (CAST(:id AS uuid), CAST(:rule_uid AS uuid),
  393. :name, 'general', 'active')
  394. """
  395. ),
  396. {
  397. "id": new_governance_uid(),
  398. "rule_uid": downstream_rule_uid,
  399. "name": downstream_spec["name"],
  400. },
  401. )
  402. connection.execute(
  403. text(
  404. """
  405. INSERT INTO public.data_rule_versions
  406. (id, rule_uid, version_no, source_text, source_language,
  407. rule_spec, spec_hash, generated_kind, status, published_at)
  408. VALUES (CAST(:id AS uuid), CAST(:rule_uid AS uuid), 1,
  409. :source_text, 'en', CAST(:rule_spec AS jsonb),
  410. :spec_hash, 'polars', 'published',
  411. CURRENT_TIMESTAMP)
  412. """
  413. ),
  414. {
  415. "id": rule_id,
  416. "rule_uid": rule_uid,
  417. "source_text": "Task 5 real cross-source integration",
  418. "rule_spec": json.dumps(spec),
  419. "spec_hash": rule["spec_hash"],
  420. },
  421. )
  422. connection.execute(
  423. text(
  424. """
  425. INSERT INTO public.data_rule_versions
  426. (id, rule_uid, version_no, source_text,
  427. source_language, rule_spec, spec_hash,
  428. generated_kind, status, published_at)
  429. VALUES (
  430. CAST(:id AS uuid), CAST(:rule_uid AS uuid), 1,
  431. :source_text, 'en', CAST(:rule_spec AS jsonb),
  432. :spec_hash, 'polars', 'published',
  433. CURRENT_TIMESTAMP
  434. )
  435. """
  436. ),
  437. {
  438. "id": downstream_rule_id,
  439. "rule_uid": downstream_rule_uid,
  440. "source_text": (
  441. "Task 6 real two-node artifact handoff"
  442. ),
  443. "rule_spec": json.dumps(downstream_spec),
  444. "spec_hash": downstream_rule["spec_hash"],
  445. },
  446. )
  447. connection.execute(
  448. text(
  449. """
  450. INSERT INTO public.dataflow_versions
  451. (id, dataflow_uid, version_no, name, dataflow_spec,
  452. input_schema_hashes, output_schema_hash, status,
  453. released_at)
  454. VALUES (CAST(:id AS uuid), CAST(:dataflow_uid AS uuid), 1,
  455. :name, '{}'::jsonb,
  456. CAST(:input_schema_hashes AS jsonb),
  457. :output_schema_hash, 'released', CURRENT_TIMESTAMP)
  458. """
  459. ),
  460. {
  461. "id": dataflow_version_id,
  462. "dataflow_uid": dataflow_uid,
  463. "name": "Task 5 real cross-source integration",
  464. "input_schema_hashes": json.dumps(
  465. [
  466. input_schema["schema_hash"],
  467. lookup_schema["schema_hash"],
  468. ]
  469. ),
  470. "output_schema_hash": output_schema["schema_hash"],
  471. },
  472. )
  473. connection.execute(
  474. text(
  475. """
  476. INSERT INTO public.dataflow_deployments
  477. (id, dataflow_version_id, environment, deployment_config,
  478. status)
  479. VALUES (CAST(:id AS uuid),
  480. CAST(:dataflow_version_id AS uuid), 'test',
  481. '{}'::jsonb, 'disabled')
  482. """
  483. ),
  484. {
  485. "id": deployment_id,
  486. "dataflow_version_id": dataflow_version_id,
  487. },
  488. )
  489. for logical_ref, binding, binding_hash in (
  490. (
  491. "customers",
  492. input_binding,
  493. compiled["plan"]["input_binding_hash"],
  494. ),
  495. (
  496. "segments",
  497. lookup_binding,
  498. lookup_operation["lookup_binding_hash"],
  499. ),
  500. (
  501. "enriched",
  502. output_binding,
  503. compiled["plan"]["output_binding_hash"],
  504. ),
  505. (
  506. "downstream",
  507. downstream_output_binding,
  508. downstream_compiled["plan"][
  509. "output_binding_hash"
  510. ],
  511. ),
  512. ):
  513. connection.execute(
  514. text(
  515. """
  516. INSERT INTO public.dataflow_dataset_bindings
  517. (id, dataflow_deployment_id, logical_ref,
  518. data_source_uid, object_kind, object_ref,
  519. schema_snapshot_id, dialect, access_mode, write_mode,
  520. binding_hash)
  521. VALUES (CAST(:id AS uuid), CAST(:deployment_id AS uuid),
  522. :logical_ref, CAST(:source_uid AS uuid),
  523. 'parquet_artifact', :object_ref,
  524. CAST(:schema_snapshot_id AS uuid), 'parquet',
  525. :access_mode, 'append', :binding_hash)
  526. """
  527. ),
  528. {
  529. "id": binding["id"],
  530. "deployment_id": deployment_id,
  531. "logical_ref": logical_ref,
  532. "source_uid": binding["data_source_uid"],
  533. "object_ref": binding["object_ref"],
  534. "schema_snapshot_id": binding["schema_snapshot_id"],
  535. "access_mode": binding["access_mode"],
  536. "binding_hash": binding_hash,
  537. },
  538. )
  539. connection.execute(
  540. text(
  541. """
  542. INSERT INTO public.dataflow_component_bindings
  543. (id, dataflow_version_id, component_id, component_kind,
  544. rule_version_id, stage, order_no, idempotency, provenance)
  545. VALUES (CAST(:id AS uuid),
  546. CAST(:dataflow_version_id AS uuid),
  547. 'task5_real_polars', 'rule.apply',
  548. CAST(:rule_version_id AS uuid), 'transform', 0,
  549. CAST(:idempotency AS jsonb), '{}'::jsonb)
  550. """
  551. ),
  552. {
  553. "id": component_binding_id,
  554. "dataflow_version_id": dataflow_version_id,
  555. "rule_version_id": rule_id,
  556. "idempotency": json.dumps(
  557. {
  558. "strategy": "deduplication_key",
  559. "key": "customer_id",
  560. }
  561. ),
  562. },
  563. )
  564. connection.execute(
  565. text(
  566. """
  567. INSERT INTO public.dataflow_component_bindings
  568. (id, dataflow_version_id, component_id,
  569. component_kind, rule_version_id, stage, order_no,
  570. idempotency, provenance)
  571. VALUES (
  572. CAST(:id AS uuid),
  573. CAST(:dataflow_version_id AS uuid),
  574. 'task6_real_handoff', 'rule.apply',
  575. CAST(:rule_version_id AS uuid), 'transform', 1,
  576. CAST(:idempotency AS jsonb), '{}'::jsonb
  577. )
  578. """
  579. ),
  580. {
  581. "id": downstream_component_binding_id,
  582. "dataflow_version_id": dataflow_version_id,
  583. "rule_version_id": downstream_rule_id,
  584. "idempotency": json.dumps(
  585. {
  586. "strategy": "deduplication_key",
  587. "key": "customer_id",
  588. }
  589. ),
  590. },
  591. )
  592. connection.execute(
  593. text(
  594. """
  595. INSERT INTO public.rule_execution_plans
  596. (id, component_binding_id, backend, compiler_version, plan,
  597. plan_hash, schema_hashes, status)
  598. VALUES (CAST(:id AS uuid),
  599. CAST(:component_binding_id AS uuid),
  600. 'polars_batch', :compiler_version,
  601. CAST(:plan AS jsonb), :plan_hash,
  602. CAST(:schema_hashes AS jsonb), 'published')
  603. """
  604. ),
  605. {
  606. "id": plan_id,
  607. "component_binding_id": component_binding_id,
  608. "compiler_version": compiled["compiler_version"],
  609. "plan": json.dumps(compiled["plan"]),
  610. "plan_hash": compiled["plan_hash"],
  611. "schema_hashes": json.dumps(schema_hashes),
  612. },
  613. )
  614. connection.execute(
  615. text(
  616. """
  617. INSERT INTO public.rule_execution_plans
  618. (id, component_binding_id, backend,
  619. compiler_version, plan, plan_hash,
  620. schema_hashes, status)
  621. VALUES (
  622. CAST(:id AS uuid),
  623. CAST(:component_binding_id AS uuid),
  624. 'polars_batch', :compiler_version,
  625. CAST(:plan AS jsonb), :plan_hash,
  626. CAST(:schema_hashes AS jsonb), 'published'
  627. )
  628. """
  629. ),
  630. {
  631. "id": downstream_plan_id,
  632. "component_binding_id": (
  633. downstream_component_binding_id
  634. ),
  635. "compiler_version": downstream_compiled[
  636. "compiler_version"
  637. ],
  638. "plan": json.dumps(downstream_compiled["plan"]),
  639. "plan_hash": downstream_compiled["plan_hash"],
  640. "schema_hashes": json.dumps(
  641. {
  642. "rule_spec_hash": downstream_compiled[
  643. "plan"
  644. ]["rule_spec_hash"],
  645. "input_schema_snapshot_id": output_schema[
  646. "id"
  647. ],
  648. "input_schema_hash": output_schema[
  649. "schema_hash"
  650. ],
  651. "output_schema_snapshot_id": output_schema[
  652. "id"
  653. ],
  654. "output_schema_hash": output_schema[
  655. "schema_hash"
  656. ],
  657. }
  658. ),
  659. },
  660. )
  661. for trusted in (
  662. {
  663. "rule_id": rule_id,
  664. "plan_id": plan_id,
  665. "compiled": compiled,
  666. "input_schema": input_schema,
  667. "output_schema": output_schema,
  668. },
  669. {
  670. "rule_id": downstream_rule_id,
  671. "plan_id": downstream_plan_id,
  672. "compiled": downstream_compiled,
  673. "input_schema": output_schema,
  674. "output_schema": output_schema,
  675. },
  676. ):
  677. trusted_plan = trusted["compiled"]["plan"]
  678. trusted_schema_hashes = {
  679. "input": trusted["input_schema"]["schema_hash"],
  680. "output": trusted["output_schema"]["schema_hash"],
  681. }
  682. trusted_binding_hashes = {
  683. "input": trusted_plan["input_binding_hash"],
  684. "output": trusted_plan["output_binding_hash"],
  685. }
  686. profile_id = new_governance_uid()
  687. logical_plan_id = new_governance_uid()
  688. connection.execute(
  689. text(
  690. """
  691. INSERT INTO public.rule_validation_profiles
  692. (id, rule_version_id, input_schema_snapshot_id,
  693. input_schema_hash, input_fields,
  694. output_schema_snapshot_id, output_schema_hash,
  695. output_fields, input_sample_artifact_ref,
  696. input_sample_artifact_digest, context_hash)
  697. VALUES
  698. (CAST(:id AS uuid), CAST(:rule_id AS uuid),
  699. CAST(:input_snapshot_id AS uuid), :input_hash,
  700. CAST(:input_fields AS jsonb),
  701. CAST(:output_snapshot_id AS uuid), :output_hash,
  702. CAST(:output_fields AS jsonb), :artifact_ref,
  703. :digest, :digest)
  704. """
  705. ),
  706. {
  707. "id": profile_id,
  708. "rule_id": trusted["rule_id"],
  709. "input_snapshot_id": trusted["input_schema"]["id"],
  710. "input_hash": trusted["input_schema"]["schema_hash"],
  711. "input_fields": json.dumps(
  712. trusted["input_schema"]["fields"]
  713. ),
  714. "output_snapshot_id": trusted["output_schema"]["id"],
  715. "output_hash": trusted["output_schema"]["schema_hash"],
  716. "output_fields": json.dumps(
  717. trusted["output_schema"]["fields"]
  718. ),
  719. "artifact_ref": f"test://{trusted['rule_id']}",
  720. "digest": "a" * 64,
  721. },
  722. )
  723. connection.execute(
  724. text(
  725. """
  726. INSERT INTO public.rule_logical_plans
  727. (id, rule_version_id, validation_profile_id,
  728. compiler_version, backend, plan, plan_hash,
  729. schema_hashes, capabilities, status)
  730. VALUES
  731. (CAST(:id AS uuid), CAST(:rule_id AS uuid),
  732. CAST(:profile_id AS uuid), :compiler_version,
  733. 'polars_batch', CAST(:plan AS jsonb), :plan_hash,
  734. CAST(:schema_hashes AS jsonb), '{}'::jsonb,
  735. 'published')
  736. """
  737. ),
  738. {
  739. "id": logical_plan_id,
  740. "rule_id": trusted["rule_id"],
  741. "profile_id": profile_id,
  742. "compiler_version": trusted["compiled"][
  743. "compiler_version"
  744. ],
  745. "plan": json.dumps(trusted_plan),
  746. "plan_hash": trusted["compiled"]["plan_hash"],
  747. "schema_hashes": json.dumps(trusted_schema_hashes),
  748. },
  749. )
  750. connection.execute(
  751. text(
  752. """
  753. INSERT INTO public.rule_logical_compile_evidence
  754. (id, logical_plan_id, compiler_version,
  755. compiler_digest, plan_hash, schema_hashes,
  756. capabilities, status)
  757. VALUES
  758. (CAST(:id AS uuid), CAST(:logical_plan_id AS uuid),
  759. :compiler_version, :digest, :plan_hash,
  760. CAST(:schema_hashes AS jsonb), '{}'::jsonb,
  761. 'success')
  762. """
  763. ),
  764. {
  765. "id": new_governance_uid(),
  766. "logical_plan_id": logical_plan_id,
  767. "compiler_version": trusted["compiled"][
  768. "compiler_version"
  769. ],
  770. "digest": "b" * 64,
  771. "plan_hash": trusted["compiled"]["plan_hash"],
  772. "schema_hashes": json.dumps(trusted_schema_hashes),
  773. },
  774. )
  775. connection.execute(
  776. text(
  777. """
  778. INSERT INTO public.rule_logical_test_evidence
  779. (id, logical_plan_id, test_kind, evidence_hash,
  780. run_id, plan_hash, schema_hashes, evidence, status)
  781. VALUES
  782. (CAST(:id AS uuid), CAST(:logical_plan_id AS uuid),
  783. 'dry_run', :digest, CAST(:run_id AS uuid), :plan_hash,
  784. CAST(:schema_hashes AS jsonb), '{}'::jsonb, 'success')
  785. """
  786. ),
  787. {
  788. "id": new_governance_uid(),
  789. "logical_plan_id": logical_plan_id,
  790. "digest": "c" * 64,
  791. "run_id": new_governance_uid(),
  792. "plan_hash": trusted["compiled"]["plan_hash"],
  793. "schema_hashes": json.dumps(trusted_schema_hashes),
  794. },
  795. )
  796. connection.execute(
  797. text(
  798. """
  799. INSERT INTO public.rule_compile_evidence
  800. (id, rule_execution_plan_id, compiler_version,
  801. compiler_digest, status, evidence, plan_hash,
  802. schema_hashes, binding_hashes, capabilities,
  803. legacy_untrusted)
  804. VALUES
  805. (CAST(:id AS uuid), CAST(:plan_id AS uuid),
  806. :compiler_version, :digest, 'success', '{}'::jsonb,
  807. :plan_hash, CAST(:schema_hashes AS jsonb),
  808. CAST(:binding_hashes AS jsonb),
  809. CAST(:capabilities AS jsonb), FALSE)
  810. """
  811. ),
  812. {
  813. "id": new_governance_uid(),
  814. "plan_id": trusted["plan_id"],
  815. "compiler_version": trusted["compiled"][
  816. "compiler_version"
  817. ],
  818. "digest": "d" * 64,
  819. "plan_hash": trusted["compiled"]["plan_hash"],
  820. "schema_hashes": json.dumps(
  821. {
  822. "rule_spec_hash": trusted_plan[
  823. "rule_spec_hash"
  824. ],
  825. "input_schema_snapshot_id": trusted_plan[
  826. "input_schema_snapshot_id"
  827. ],
  828. "input_schema_hash": trusted_plan[
  829. "input_schema_hash"
  830. ],
  831. "output_schema_snapshot_id": trusted_plan[
  832. "output_schema_snapshot_id"
  833. ],
  834. "output_schema_hash": trusted_plan[
  835. "output_schema_hash"
  836. ],
  837. }
  838. ),
  839. "binding_hashes": json.dumps(
  840. trusted_binding_hashes
  841. ),
  842. "capabilities": json.dumps(
  843. {
  844. "resource_limits": trusted_plan[
  845. "resource_limits"
  846. ]
  847. }
  848. ),
  849. },
  850. )
  851. connection.execute(
  852. text(
  853. """
  854. INSERT INTO public.rule_test_evidence
  855. (id, rule_execution_plan_id, test_kind, evidence_hash,
  856. status, evidence, plan_hash, schema_hashes,
  857. binding_hashes, run_id, legacy_untrusted)
  858. SELECT
  859. CAST(:id AS uuid), p.id, 'sample', :digest,
  860. 'success', '{}'::jsonb, p.plan_hash, p.schema_hashes,
  861. CAST(:binding_hashes AS jsonb),
  862. CAST(:run_id AS uuid), FALSE
  863. FROM public.rule_execution_plans p
  864. WHERE p.id = CAST(:plan_id AS uuid)
  865. """
  866. ),
  867. {
  868. "id": new_governance_uid(),
  869. "plan_id": trusted["plan_id"],
  870. "digest": "e" * 64,
  871. "binding_hashes": json.dumps(
  872. trusted_binding_hashes
  873. ),
  874. "run_id": new_governance_uid(),
  875. },
  876. )
  877. connection.execute(
  878. text(
  879. """
  880. INSERT INTO public.rule_publication_audits
  881. (id, rule_version_id, rule_execution_plan_id,
  882. action, to_status, evidence_hash)
  883. VALUES
  884. (CAST(:id AS uuid), CAST(:rule_id AS uuid),
  885. CAST(:plan_id AS uuid), 'published', 'published',
  886. :plan_hash)
  887. """
  888. ),
  889. {
  890. "id": new_governance_uid(),
  891. "rule_id": trusted["rule_id"],
  892. "plan_id": trusted["plan_id"],
  893. "plan_hash": trusted["compiled"]["plan_hash"],
  894. },
  895. )
  896. customer_path = tmp_path / "customers.parquet"
  897. segment_path = tmp_path / "segments.parquet"
  898. pl.DataFrame(customer_rows).write_parquet(customer_path)
  899. pl.DataFrame(segment_rows).write_parquet(segment_path)
  900. resolver = PostgresArtifactResolver(platform, store)
  901. customer_artifact = resolver.publish_path(
  902. str(customer_path),
  903. binding_id=input_binding["id"],
  904. binding_hash=compiled["plan"]["input_binding_hash"],
  905. correlation_id=correlation_id,
  906. kind="input",
  907. ttl_seconds=900,
  908. schema_fields=input_schema["fields"],
  909. limits=compiled["plan"]["resource_limits"],
  910. )
  911. segment_artifact = resolver.publish_path(
  912. str(segment_path),
  913. binding_id=lookup_binding["id"],
  914. binding_hash=lookup_operation["lookup_binding_hash"],
  915. correlation_id=correlation_id,
  916. kind="lookup",
  917. ttl_seconds=900,
  918. schema_fields=lookup_schema["fields"],
  919. limits=compiled["plan"]["resource_limits"],
  920. )
  921. node = {
  922. "id": "task5_real_polars",
  923. "type": "rule.apply",
  924. "purpose": "write",
  925. "idempotency": {
  926. "strategy": "deduplication_key",
  927. "key": "customer_id",
  928. },
  929. "config": {
  930. "component_binding_id": component_binding_id,
  931. "rule_version_id": rule["id"],
  932. "execution_plan_hash": compiled["plan_hash"],
  933. },
  934. }
  935. from app.core.mcp.gateway import SchedulingGateway
  936. from app.core.mcp.identity import AgentIdentity
  937. from app.core.mcp.persistence import PostgresSchedulingPlanStore
  938. from app.runner.api import create_runner_app
  939. from app.runner.auth import TaskTokenIssuer, TaskTokenVerifier
  940. from app.runner.ledger import PostgresTaskLedger
  941. from app.runner.nodes import NodeRegistry
  942. from app.runner.rule_evidence import PostgresRuleEvidenceWriter
  943. class Audit:
  944. def __init__(self):
  945. self.events = []
  946. def record(self, event):
  947. self.events.append(event)
  948. class CanaryEngine:
  949. def __init__(self):
  950. self.calls = []
  951. def deploy_disabled(self, _definition):
  952. return {"revision": "task6-real-gateway"}
  953. def activate(self, namespace, flow_id):
  954. self.calls.append(("activate", namespace, flow_id))
  955. def execute(self, namespace, flow_id, inputs=None, **_options):
  956. self.calls.append(
  957. ("execute", namespace, flow_id, dict(inputs or {}))
  958. )
  959. return {"id": f"task6-canary-{correlation_id}"}
  960. def deactivate(self, namespace, flow_id):
  961. self.calls.append(("deactivate", namespace, flow_id))
  962. task_secret = "task5-real-http-secret-value-32-bytes"
  963. verifier = TaskTokenVerifier(task_secret)
  964. gateway_engine = CanaryEngine()
  965. gateway_store = PostgresSchedulingPlanStore(platform)
  966. gateway = SchedulingGateway(
  967. plans=gateway_store,
  968. audit=Audit(),
  969. engine=gateway_engine,
  970. token_issuer=TaskTokenIssuer(task_secret),
  971. )
  972. gateway_identity = AgentIdentity(
  973. subject="task6-real-scheduler",
  974. roles=frozenset({"scheduler"}),
  975. business_domains=frozenset({"sales"}),
  976. environments=frozenset({"test"}),
  977. correlation_id=correlation_id,
  978. )
  979. candidate = gateway.create_candidate_plan(
  980. gateway_identity,
  981. business_domain="sales",
  982. environment="test",
  983. workflow_spec={
  984. "schema_version": "1.0",
  985. "dataflow_uid": dataflow_uid,
  986. "name": "Task 6 governed rule canary",
  987. "nodes": [node],
  988. "edges": [],
  989. "parameters": {},
  990. },
  991. schedule_plan={
  992. "schema_version": "1.0",
  993. "timezone": "Asia/Shanghai",
  994. "triggers": [{"type": "manual"}],
  995. "max_concurrency": 1,
  996. "conflict_policy": "skip",
  997. "timeout_seconds": 600,
  998. "retry": {
  999. "max_attempts": 1,
  1000. "delay_seconds": 1,
  1001. },
  1002. "backfill": {"max_days": 1, "max_runs": 1},
  1003. },
  1004. )
  1005. gateway_candidate_id = candidate["candidate_id"]
  1006. with platform.begin() as connection:
  1007. connection.execute(
  1008. text(
  1009. """
  1010. UPDATE public.dataflow_workflow_versions
  1011. SET write_authorized = TRUE
  1012. WHERE id = CAST(:id AS uuid)
  1013. """
  1014. ),
  1015. {"id": gateway_candidate_id},
  1016. )
  1017. deployed = gateway.deploy_disabled_version(
  1018. gateway_identity,
  1019. candidate_id=gateway_candidate_id,
  1020. business_domain="sales",
  1021. environment="test",
  1022. )
  1023. assert deployed["deployment_id"] == deployment_id
  1024. assert deployed["candidate_id"] != deployed["deployment_id"]
  1025. gateway.run_canary(
  1026. gateway_identity,
  1027. candidate_id=gateway_candidate_id,
  1028. business_domain="sales",
  1029. environment="test",
  1030. inputs={},
  1031. )
  1032. execution_call = next(
  1033. call for call in gateway_engine.calls if call[0] == "execute"
  1034. )
  1035. task_token = execution_call[3]["dataops_task_tokens"][node["id"]]
  1036. issued_claims = verifier.verify(task_token, node=node)
  1037. assert issued_claims.deployment_id == deployment_id
  1038. assert issued_claims.workflow_version == 1
  1039. assert issued_claims.environment == "test"
  1040. executor = RulePlanExecutor(
  1041. PostgresRulePlanRepository(platform),
  1042. adapters={
  1043. "polars_batch": PolarsRulePlanAdapter(
  1044. artifact_store=store,
  1045. artifact_resolver=resolver,
  1046. artifact_ttl_seconds=900,
  1047. )
  1048. },
  1049. )
  1050. result = executor.execute(
  1051. node,
  1052. {},
  1053. write_authorized=True,
  1054. correlation_id=correlation_id,
  1055. )
  1056. repeated = executor.execute(
  1057. node,
  1058. {},
  1059. write_authorized=True,
  1060. correlation_id=correlation_id,
  1061. )
  1062. assert result["rows_in"] == 4
  1063. assert result["rows_out"] == 2
  1064. assert result["rows_rejected"] == 1
  1065. assert result["rows_deduplicated"] == 1
  1066. assert result["rows_filtered"] == 0
  1067. assert result["rows_join_dropped"] == 0
  1068. assert result["rows_aggregated"] == 0
  1069. assert result["violation_count"] == 1
  1070. assert result["violations"] == [
  1071. {"step_id": "valid_mobile", "count": 1}
  1072. ]
  1073. output = store.read(
  1074. result["artifact_ref"],
  1075. result["digest"],
  1076. expected_schema_fields=output_schema["fields"],
  1077. limits=compiled["plan"]["resource_limits"],
  1078. ).collect()
  1079. assert output.sort("customer_id").to_dicts() == [
  1080. {
  1081. "customer_id": 1,
  1082. "mobile": "13800138000",
  1083. "name": "Alice Updated",
  1084. "segment_code": "A",
  1085. "segment_name": "Gold",
  1086. "version_no": 2,
  1087. },
  1088. {
  1089. "customer_id": 3,
  1090. "mobile": "13900139000",
  1091. "name": "Carol",
  1092. "segment_code": "C",
  1093. "segment_name": "Silver",
  1094. "version_no": 1,
  1095. },
  1096. ]
  1097. assert all(
  1098. item.object_name.startswith(prefix)
  1099. for item in minio.list_objects(
  1100. bucket, prefix=prefix, recursive=True
  1101. )
  1102. )
  1103. assert repeated["rows_out"] == 2
  1104. assert repeated["artifact_ref"] == result["artifact_ref"]
  1105. assert "schema_fields" not in result
  1106. ledger_jti = verifier.verify(task_token, node=node).jti
  1107. retry_token = TaskTokenIssuer(task_secret).issue(
  1108. task_uid=new_governance_uid(),
  1109. dataflow_uid=dataflow_uid,
  1110. deployment_id=deployment_id,
  1111. environment="test",
  1112. workflow_version=1,
  1113. correlation_id=correlation_id,
  1114. node=node,
  1115. write_authorized=True,
  1116. )
  1117. retry_ledger_jti = verifier.verify(
  1118. retry_token, node=node
  1119. ).jti
  1120. real_evidence_writer = PostgresRuleEvidenceWriter(
  1121. platform,
  1122. store,
  1123. sample_ttl_seconds=900,
  1124. )
  1125. evidenced_executor = RulePlanExecutor(
  1126. PostgresRulePlanRepository(platform),
  1127. adapters={
  1128. "polars_batch": PolarsRulePlanAdapter(
  1129. artifact_store=store,
  1130. artifact_resolver=resolver,
  1131. artifact_ttl_seconds=900,
  1132. )
  1133. },
  1134. evidence_writer=real_evidence_writer,
  1135. )
  1136. real_ledger = PostgresTaskLedger(platform, lease_seconds=30)
  1137. runner_app = create_runner_app(
  1138. verifier=verifier,
  1139. ledger=real_ledger,
  1140. registry=NodeRegistry({"rule.apply": evidenced_executor}),
  1141. )
  1142. no_run_crash_token = TaskTokenIssuer(task_secret).issue(
  1143. task_uid=new_governance_uid(),
  1144. dataflow_uid=dataflow_uid,
  1145. deployment_id=deployment_id,
  1146. environment="test",
  1147. workflow_version=1,
  1148. correlation_id=no_run_crash_correlation_id,
  1149. node=node,
  1150. write_authorized=True,
  1151. )
  1152. no_run_claims = verifier.verify(no_run_crash_token, node=node)
  1153. no_run_crash_jti = no_run_claims.jti
  1154. evidence_crash_token = TaskTokenIssuer(task_secret).issue(
  1155. task_uid=new_governance_uid(),
  1156. dataflow_uid=dataflow_uid,
  1157. deployment_id=deployment_id,
  1158. environment="test",
  1159. workflow_version=1,
  1160. correlation_id=evidence_crash_correlation_id,
  1161. node=node,
  1162. write_authorized=True,
  1163. )
  1164. evidence_crash_claims = verifier.verify(
  1165. evidence_crash_token,
  1166. node=node,
  1167. )
  1168. evidence_crash_jti = evidence_crash_claims.jti
  1169. active_retry_token = TaskTokenIssuer(task_secret).issue(
  1170. task_uid=new_governance_uid(),
  1171. dataflow_uid=dataflow_uid,
  1172. deployment_id=deployment_id,
  1173. environment="test",
  1174. workflow_version=1,
  1175. correlation_id=active_retry_correlation_id,
  1176. node=node,
  1177. write_authorized=True,
  1178. )
  1179. active_retry_claims = verifier.verify(
  1180. active_retry_token,
  1181. node=node,
  1182. )
  1183. active_retry_jti = active_retry_claims.jti
  1184. def task_binding(claims):
  1185. return {
  1186. "task_uid": claims.task_uid,
  1187. "dataflow_uid": claims.dataflow_uid,
  1188. "deployment_id": claims.deployment_id,
  1189. "environment": claims.environment,
  1190. "workflow_version": claims.workflow_version,
  1191. "correlation_id": claims.correlation_id,
  1192. "node_id": claims.node_id,
  1193. "node_type": claims.node_type,
  1194. "data_source_uid": None,
  1195. "idempotency_key": "customer_id",
  1196. }
  1197. assert real_ledger.claim(
  1198. no_run_crash_jti,
  1199. task_binding(no_run_claims),
  1200. expires_at=no_run_claims.expires_at,
  1201. )
  1202. assert real_ledger.claim(
  1203. evidence_crash_jti,
  1204. task_binding(evidence_crash_claims),
  1205. expires_at=evidence_crash_claims.expires_at,
  1206. )
  1207. assert real_ledger.claim(
  1208. active_retry_jti,
  1209. task_binding(active_retry_claims),
  1210. expires_at=active_retry_claims.expires_at,
  1211. )
  1212. real_evidence_writer.start(
  1213. component_binding_id=component_binding_id,
  1214. rule_version_id=rule_id,
  1215. plan_hash=compiled["plan_hash"],
  1216. correlation_id=evidence_crash_correlation_id,
  1217. dataflow_uid=dataflow_uid,
  1218. deployment_id=deployment_id,
  1219. environment="test",
  1220. workflow_version=1,
  1221. node_id=node["id"],
  1222. lease_owner=evidence_crash_jti,
  1223. )
  1224. real_evidence_writer.start(
  1225. component_binding_id=component_binding_id,
  1226. rule_version_id=rule_id,
  1227. plan_hash=compiled["plan_hash"],
  1228. correlation_id=active_retry_correlation_id,
  1229. dataflow_uid=dataflow_uid,
  1230. deployment_id=deployment_id,
  1231. environment="test",
  1232. workflow_version=1,
  1233. node_id=node["id"],
  1234. lease_owner=active_retry_jti,
  1235. )
  1236. with platform.begin() as connection:
  1237. connection.execute(
  1238. text(
  1239. """
  1240. UPDATE public.runner_task_executions
  1241. SET lease_expires_at =
  1242. CURRENT_TIMESTAMP - INTERVAL '1 second'
  1243. WHERE token_jti IN (
  1244. CAST(:no_run_jti AS uuid),
  1245. CAST(:evidence_jti AS uuid)
  1246. )
  1247. """
  1248. ),
  1249. {
  1250. "no_run_jti": no_run_crash_jti,
  1251. "evidence_jti": evidence_crash_jti,
  1252. },
  1253. )
  1254. connection.execute(
  1255. text(
  1256. """
  1257. UPDATE public.rule_runs
  1258. SET lease_expires_at =
  1259. CURRENT_TIMESTAMP - INTERVAL '1 second'
  1260. WHERE lease_owner = CAST(:jti AS uuid)
  1261. """
  1262. ),
  1263. {"jti": evidence_crash_jti},
  1264. )
  1265. with runner_app.test_client() as client:
  1266. http_result = client.post(
  1267. "/v1/tasks/execute",
  1268. json={
  1269. "task_token": task_token,
  1270. "node": node,
  1271. "parameters": {},
  1272. },
  1273. )
  1274. replay = client.post(
  1275. "/v1/tasks/execute",
  1276. json={
  1277. "task_token": task_token,
  1278. "node": node,
  1279. "parameters": {},
  1280. },
  1281. )
  1282. retried = client.post(
  1283. "/v1/tasks/execute",
  1284. json={
  1285. "task_token": retry_token,
  1286. "node": node,
  1287. "parameters": {},
  1288. },
  1289. )
  1290. no_run_crash = client.post(
  1291. "/v1/tasks/execute",
  1292. json={
  1293. "task_token": no_run_crash_token,
  1294. "node": node,
  1295. "parameters": {},
  1296. },
  1297. )
  1298. evidence_crash = client.post(
  1299. "/v1/tasks/execute",
  1300. json={
  1301. "task_token": evidence_crash_token,
  1302. "node": node,
  1303. "parameters": {},
  1304. },
  1305. )
  1306. active_retry = client.post(
  1307. "/v1/tasks/execute",
  1308. json={
  1309. "task_token": active_retry_token,
  1310. "node": node,
  1311. "parameters": {},
  1312. },
  1313. )
  1314. assert http_result.status_code == 200, http_result.get_json()
  1315. assert http_result.get_json()["output_artifact"] == result[
  1316. "artifact_ref"
  1317. ]
  1318. assert http_result.get_json()["result"]["artifact_ref"] == result[
  1319. "artifact_ref"
  1320. ]
  1321. assert replay.status_code == 200
  1322. assert replay.headers["X-Idempotent-Replay"] == "true"
  1323. assert replay.get_json() == http_result.get_json()
  1324. assert retried.status_code == 200
  1325. assert retried.get_json()["output_artifact"] == result[
  1326. "artifact_ref"
  1327. ]
  1328. assert no_run_crash.status_code == 409
  1329. assert no_run_crash.get_json() == {
  1330. "error": "task execution outcome is unknown"
  1331. }
  1332. assert evidence_crash.status_code == 409
  1333. assert evidence_crash.get_json() == {
  1334. "error": "task execution outcome is unknown"
  1335. }
  1336. assert active_retry.status_code == 202
  1337. assert active_retry.headers["Retry-After"] == "2"
  1338. assert real_ledger.get(active_retry_jti).status == "running"
  1339. assert real_ledger.get(active_retry_jti).replay_body is None
  1340. assert real_ledger.get(no_run_crash_jti).status == "unknown"
  1341. assert real_ledger.get(evidence_crash_jti).status == "unknown"
  1342. with platform.connect() as connection:
  1343. assert connection.execute(
  1344. text(
  1345. """
  1346. SELECT COUNT(*)
  1347. FROM public.rule_runs
  1348. WHERE correlation_id =
  1349. CAST(:correlation_id AS uuid)
  1350. """
  1351. ),
  1352. {"correlation_id": no_run_crash_correlation_id},
  1353. ).scalar_one() == 0
  1354. crash_evidence = connection.execute(
  1355. text(
  1356. """
  1357. SELECT status, commit_outcome
  1358. FROM public.rule_runs
  1359. WHERE correlation_id =
  1360. CAST(:correlation_id AS uuid)
  1361. """
  1362. ),
  1363. {"correlation_id": evidence_crash_correlation_id},
  1364. ).mappings().one()
  1365. assert dict(crash_evidence) == {
  1366. "status": "unknown",
  1367. "commit_outcome": "unknown",
  1368. }
  1369. terminal_replay = real_evidence_writer.reconcile_expired_lease(
  1370. lease_owner=ledger_jti,
  1371. deployment_id=deployment_id,
  1372. correlation_id=correlation_id,
  1373. component_binding_id=component_binding_id,
  1374. rule_version_id=rule_id,
  1375. plan_hash=compiled["plan_hash"],
  1376. )
  1377. assert terminal_replay["state"] == "terminal"
  1378. assert terminal_replay["status"] == "success"
  1379. assert terminal_replay["result_digest"] == hashlib.sha256(
  1380. json.dumps(
  1381. terminal_replay["result"],
  1382. sort_keys=True,
  1383. separators=(",", ":"),
  1384. ensure_ascii=False,
  1385. ).encode("utf-8")
  1386. ).hexdigest()
  1387. assert re.fullmatch(
  1388. r"[0-9a-f]{64}",
  1389. terminal_replay["evidence_digest"],
  1390. )
  1391. ledger_record = PostgresTaskLedger(platform).get(ledger_jti)
  1392. assert ledger_record is not None
  1393. assert ledger_record.status == "success"
  1394. assert ledger_record.commit_outcome == "committed"
  1395. with platform.connect() as connection:
  1396. run_evidence = connection.execute(
  1397. text(
  1398. """
  1399. SELECT status, commit_outcome, rows_in, rows_out,
  1400. rows_rejected, rows_quarantined, public_result
  1401. FROM public.rule_runs
  1402. WHERE correlation_id = CAST(:correlation_id AS uuid)
  1403. AND component_binding_id =
  1404. CAST(:component_binding_id AS uuid)
  1405. """
  1406. ),
  1407. {
  1408. "correlation_id": correlation_id,
  1409. "component_binding_id": component_binding_id,
  1410. },
  1411. ).mappings().one()
  1412. sample_evidence = connection.execute(
  1413. text(
  1414. """
  1415. SELECT s.artifact_ref, s.artifact_digest,
  1416. s.sample_count, s.redaction_policy,
  1417. s.expires_at, s.handoff_status
  1418. FROM public.rule_violation_samples s
  1419. JOIN public.rule_runs r ON r.id = s.rule_run_id
  1420. WHERE r.correlation_id =
  1421. CAST(:correlation_id AS uuid)
  1422. """
  1423. ),
  1424. {"correlation_id": correlation_id},
  1425. ).mappings().one()
  1426. assert run_evidence["status"] == "success"
  1427. assert run_evidence["commit_outcome"] == "committed"
  1428. assert run_evidence["rows_in"] == 4
  1429. assert run_evidence["rows_out"] == 2
  1430. assert run_evidence["rows_rejected"] == 1
  1431. assert run_evidence["rows_quarantined"] == 0
  1432. assert run_evidence["public_result"]["output_artifact"] == result[
  1433. "artifact_ref"
  1434. ]
  1435. assert sample_evidence["artifact_digest"]
  1436. assert sample_evidence["sample_count"] == 1
  1437. assert (
  1438. sample_evidence["redaction_policy"]
  1439. == "rule-violation-default-v1"
  1440. )
  1441. assert sample_evidence["handoff_status"] == "ready"
  1442. assert store.read(
  1443. sample_evidence["artifact_ref"],
  1444. sample_evidence["artifact_digest"],
  1445. expected_schema_fields=[
  1446. {
  1447. "name": name,
  1448. "type": "string",
  1449. "nullable": True,
  1450. }
  1451. for name in (
  1452. "customer_id",
  1453. "mobile",
  1454. "name",
  1455. "segment_code",
  1456. "segment_name",
  1457. "version_no",
  1458. )
  1459. ],
  1460. ).collect().to_dicts() == [
  1461. {
  1462. "customer_id": "[REDACTED]",
  1463. "mobile": "[REDACTED]",
  1464. "name": "[REDACTED]",
  1465. "segment_code": "[REDACTED]",
  1466. "segment_name": "[REDACTED]",
  1467. "version_no": "[REDACTED]",
  1468. }
  1469. ]
  1470. orphan = store.write(
  1471. pl.DataFrame({"value": ["orphan"]}),
  1472. correlation_id,
  1473. 900,
  1474. schema_fields=[
  1475. {
  1476. "name": "value",
  1477. "type": "string",
  1478. "nullable": True,
  1479. }
  1480. ],
  1481. )
  1482. original_clock = store.clock
  1483. store.clock = lambda: datetime.now(UTC) + timedelta(seconds=60)
  1484. try:
  1485. reconciliation = resolver.reconcile(
  1486. limit=20,
  1487. grace_seconds=30,
  1488. )
  1489. finally:
  1490. store.clock = original_clock
  1491. assert reconciliation["orphans_deleted"] >= 1
  1492. assert store.describe_optional(orphan["artifact_ref"]) is None
  1493. assert store.describe(sample_evidence["artifact_ref"])[
  1494. "digest"
  1495. ] == sample_evidence["artifact_digest"]
  1496. downstream_node = {
  1497. "id": "task6_real_handoff",
  1498. "type": "rule.apply",
  1499. "purpose": "write",
  1500. "idempotency": {
  1501. "strategy": "deduplication_key",
  1502. "key": "customer_id",
  1503. },
  1504. "config": {
  1505. "component_binding_id": (
  1506. downstream_component_binding_id
  1507. ),
  1508. "rule_version_id": downstream_rule_id,
  1509. "execution_plan_hash": downstream_compiled[
  1510. "plan_hash"
  1511. ],
  1512. },
  1513. }
  1514. downstream_result = evidenced_executor.execute(
  1515. downstream_node,
  1516. {"input_artifact": result["artifact_ref"]},
  1517. write_authorized=True,
  1518. correlation_id=correlation_id,
  1519. dataflow_uid=dataflow_uid,
  1520. deployment_id=deployment_id,
  1521. environment="test",
  1522. workflow_version=1,
  1523. node_id="task6_real_handoff",
  1524. task_jti=new_governance_uid(),
  1525. )
  1526. assert downstream_result["rows_in"] == 2
  1527. assert downstream_result["rows_out"] == 2
  1528. assert downstream_result["output_artifact"] != result[
  1529. "artifact_ref"
  1530. ]
  1531. assert store.read(
  1532. downstream_result["artifact_ref"],
  1533. downstream_result["digest"],
  1534. expected_schema_fields=output_schema["fields"],
  1535. limits=downstream_compiled["plan"]["resource_limits"],
  1536. ).collect().sort("customer_id").to_dicts() == (
  1537. output.sort("customer_id").to_dicts()
  1538. )
  1539. replayed_downstream = evidenced_executor.execute(
  1540. downstream_node,
  1541. {"input_artifact": result["artifact_ref"]},
  1542. write_authorized=True,
  1543. correlation_id=correlation_id,
  1544. dataflow_uid=dataflow_uid,
  1545. deployment_id=deployment_id,
  1546. environment="test",
  1547. workflow_version=1,
  1548. node_id="task6_real_handoff",
  1549. task_jti=new_governance_uid(),
  1550. )
  1551. assert replayed_downstream["artifact_ref"] == downstream_result[
  1552. "artifact_ref"
  1553. ]
  1554. with platform.connect() as connection:
  1555. assert connection.execute(
  1556. text(
  1557. """
  1558. SELECT COUNT(*)
  1559. FROM public.rule_runs
  1560. WHERE correlation_id =
  1561. CAST(:correlation_id AS uuid)
  1562. """
  1563. ),
  1564. {"correlation_id": correlation_id},
  1565. ).scalar_one() == 2
  1566. from app.runner.nodes import NodeExecutionError
  1567. class FailingAdapter:
  1568. def execute(self, **_kwargs):
  1569. raise NodeExecutionError(
  1570. "safe downstream failure",
  1571. commit_outcome="not_committed",
  1572. )
  1573. failed_executor = RulePlanExecutor(
  1574. PostgresRulePlanRepository(platform),
  1575. adapters={"polars_batch": FailingAdapter()},
  1576. evidence_writer=PostgresRuleEvidenceWriter(
  1577. platform,
  1578. store,
  1579. sample_ttl_seconds=900,
  1580. ),
  1581. )
  1582. with pytest.raises(NodeExecutionError, match="safe downstream"):
  1583. failed_executor.execute(
  1584. node,
  1585. {},
  1586. write_authorized=True,
  1587. correlation_id=failure_correlation_id,
  1588. dataflow_uid=dataflow_uid,
  1589. deployment_id=deployment_id,
  1590. environment="test",
  1591. workflow_version=1,
  1592. node_id="task5_real_polars",
  1593. task_jti=new_governance_uid(),
  1594. )
  1595. with platform.connect() as connection:
  1596. failed_evidence = connection.execute(
  1597. text(
  1598. """
  1599. SELECT status, commit_outcome
  1600. FROM public.rule_runs
  1601. WHERE correlation_id =
  1602. CAST(:correlation_id AS uuid)
  1603. """
  1604. ),
  1605. {"correlation_id": failure_correlation_id},
  1606. ).mappings().one()
  1607. assert dict(failed_evidence) == {
  1608. "status": "failed",
  1609. "commit_outcome": "not_committed",
  1610. }
  1611. class UnknownAdapter:
  1612. def execute(self, **_kwargs):
  1613. raise NodeExecutionError(
  1614. "safe uncertain commit",
  1615. commit_outcome="unknown",
  1616. )
  1617. unknown_executor = RulePlanExecutor(
  1618. PostgresRulePlanRepository(platform),
  1619. adapters={"polars_batch": UnknownAdapter()},
  1620. evidence_writer=PostgresRuleEvidenceWriter(
  1621. platform,
  1622. store,
  1623. sample_ttl_seconds=900,
  1624. ),
  1625. )
  1626. with pytest.raises(NodeExecutionError, match="uncertain commit"):
  1627. unknown_executor.execute(
  1628. node,
  1629. {},
  1630. write_authorized=True,
  1631. correlation_id=unknown_correlation_id,
  1632. dataflow_uid=dataflow_uid,
  1633. deployment_id=deployment_id,
  1634. environment="test",
  1635. workflow_version=1,
  1636. node_id="task5_real_polars",
  1637. task_jti=new_governance_uid(),
  1638. )
  1639. with platform.connect() as connection:
  1640. unknown_evidence = connection.execute(
  1641. text(
  1642. """
  1643. SELECT status, commit_outcome
  1644. FROM public.rule_runs
  1645. WHERE correlation_id =
  1646. CAST(:correlation_id AS uuid)
  1647. """
  1648. ),
  1649. {"correlation_id": unknown_correlation_id},
  1650. ).mappings().one()
  1651. assert dict(unknown_evidence) == {
  1652. "status": "unknown",
  1653. "commit_outcome": "unknown",
  1654. }
  1655. evidence_writer = PostgresRuleEvidenceWriter(
  1656. platform,
  1657. store,
  1658. sample_ttl_seconds=900,
  1659. lease_seconds=30,
  1660. )
  1661. lease_owner = new_governance_uid()
  1662. lease_run_id = evidence_writer.start(
  1663. component_binding_id=component_binding_id,
  1664. rule_version_id=rule_id,
  1665. plan_hash=compiled["plan_hash"],
  1666. correlation_id=lease_correlation_id,
  1667. dataflow_uid=dataflow_uid,
  1668. deployment_id=deployment_id,
  1669. environment="test",
  1670. workflow_version=1,
  1671. node_id="task5_real_polars",
  1672. lease_owner=lease_owner,
  1673. )
  1674. with pytest.raises(ValueError, match="not owned"):
  1675. evidence_writer.heartbeat(
  1676. lease_run_id,
  1677. new_governance_uid(),
  1678. )
  1679. evidence_writer.heartbeat(lease_run_id, lease_owner)
  1680. with platform.begin() as connection:
  1681. connection.execute(
  1682. text(
  1683. """
  1684. UPDATE public.rule_runs
  1685. SET lease_expires_at =
  1686. CURRENT_TIMESTAMP - INTERVAL '1 second'
  1687. WHERE id = CAST(:id AS uuid)
  1688. """
  1689. ),
  1690. {"id": lease_run_id},
  1691. )
  1692. assert evidence_writer.start(
  1693. component_binding_id=component_binding_id,
  1694. rule_version_id=rule_id,
  1695. plan_hash=compiled["plan_hash"],
  1696. correlation_id=lease_correlation_id,
  1697. dataflow_uid=dataflow_uid,
  1698. deployment_id=deployment_id,
  1699. environment="test",
  1700. workflow_version=1,
  1701. node_id="task5_real_polars",
  1702. lease_owner=new_governance_uid(),
  1703. ) == lease_run_id
  1704. assert evidence_writer.replay(lease_run_id)["status"] == "unknown"
  1705. class FailOnceConnection:
  1706. def __init__(self, connection, state):
  1707. self.connection = connection
  1708. self.state = state
  1709. def execute(self, statement, parameters=None):
  1710. sql = str(statement)
  1711. if (
  1712. self.state["armed"]
  1713. and "UPDATE public.rule_runs" in sql
  1714. and "rows_in = :rows_in" in sql
  1715. ):
  1716. self.state["armed"] = False
  1717. raise RuntimeError("simulated response loss")
  1718. return self.connection.execute(statement, parameters)
  1719. class FailOnceBegin:
  1720. def __init__(self, context, state):
  1721. self.context = context
  1722. self.state = state
  1723. def __enter__(self):
  1724. return FailOnceConnection(
  1725. self.context.__enter__(),
  1726. self.state,
  1727. )
  1728. def __exit__(self, *args):
  1729. return self.context.__exit__(*args)
  1730. class FailOnceEngine:
  1731. def __init__(self, engine):
  1732. self.engine = engine
  1733. self.state = {"armed": True}
  1734. def connect(self):
  1735. return self.engine.connect()
  1736. def begin(self):
  1737. return FailOnceBegin(
  1738. self.engine.begin(),
  1739. self.state,
  1740. )
  1741. sample_crash_writer = PostgresRuleEvidenceWriter(
  1742. FailOnceEngine(platform),
  1743. store,
  1744. sample_ttl_seconds=900,
  1745. )
  1746. sample_crash_run_id = sample_crash_writer.start(
  1747. component_binding_id=component_binding_id,
  1748. rule_version_id=rule_id,
  1749. plan_hash=compiled["plan_hash"],
  1750. correlation_id=sample_crash_correlation_id,
  1751. dataflow_uid=dataflow_uid,
  1752. deployment_id=deployment_id,
  1753. environment="test",
  1754. workflow_version=1,
  1755. node_id="task5_real_polars",
  1756. lease_owner=new_governance_uid(),
  1757. )
  1758. sample_crash_evidence = {
  1759. "status": "success",
  1760. "rows_in": 1,
  1761. "rows_out": 0,
  1762. "rows_rejected": 1,
  1763. "rows_quarantined": 0,
  1764. "commit_outcome": "committed",
  1765. "timings": {"duration_ms": 1},
  1766. "violation_sample": [{"mobile": "[REDACTED]"}],
  1767. "sample_count": 1,
  1768. "redaction_policy": "rule-violation-default-v1",
  1769. }
  1770. with pytest.raises(RuntimeError, match="outcome is unknown"):
  1771. sample_crash_writer.finish(
  1772. sample_crash_run_id,
  1773. sample_crash_evidence,
  1774. )
  1775. with platform.connect() as connection:
  1776. assert connection.execute(
  1777. text(
  1778. """
  1779. SELECT handoff_status
  1780. FROM public.rule_violation_samples
  1781. WHERE rule_run_id = CAST(:id AS uuid)
  1782. """
  1783. ),
  1784. {"id": sample_crash_run_id},
  1785. ).scalar_one() == "ready"
  1786. with platform.begin() as connection:
  1787. connection.execute(
  1788. text(
  1789. """
  1790. UPDATE public.rule_violation_samples
  1791. SET handoff_status = 'unknown',
  1792. cleanup_claim = CAST(:cleanup_claim AS uuid),
  1793. cleanup_claim_expires_at =
  1794. CURRENT_TIMESTAMP + INTERVAL '5 minutes'
  1795. WHERE rule_run_id = CAST(:id AS uuid)
  1796. """
  1797. ),
  1798. {
  1799. "id": sample_crash_run_id,
  1800. "cleanup_claim": new_governance_uid(),
  1801. },
  1802. )
  1803. before_claim_expiry = evidence_writer.reconcile_samples(
  1804. limit=10
  1805. )
  1806. assert before_claim_expiry["claimed"] == 0
  1807. with platform.begin() as connection:
  1808. connection.execute(
  1809. text(
  1810. """
  1811. UPDATE public.rule_violation_samples
  1812. SET cleanup_claim_expires_at =
  1813. CURRENT_TIMESTAMP - INTERVAL '1 second'
  1814. WHERE rule_run_id = CAST(:id AS uuid)
  1815. """
  1816. ),
  1817. {"id": sample_crash_run_id},
  1818. )
  1819. sample_reconciliation = evidence_writer.reconcile_samples(
  1820. limit=10
  1821. )
  1822. assert sample_reconciliation["claimed"] == 1
  1823. assert sample_reconciliation["ready"] >= 1
  1824. evidence_writer.finish(
  1825. sample_crash_run_id,
  1826. sample_crash_evidence,
  1827. )
  1828. assert evidence_writer.replay(sample_crash_run_id)[
  1829. "status"
  1830. ] == "success"
  1831. with platform.connect() as connection:
  1832. missing_sample_ref = connection.execute(
  1833. text(
  1834. """
  1835. SELECT artifact_ref
  1836. FROM public.rule_violation_samples
  1837. WHERE rule_run_id = CAST(:id AS uuid)
  1838. """
  1839. ),
  1840. {"id": sample_crash_run_id},
  1841. ).scalar_one()
  1842. store.delete(missing_sample_ref)
  1843. with platform.begin() as connection:
  1844. connection.execute(
  1845. text(
  1846. """
  1847. UPDATE public.rule_violation_samples
  1848. SET handoff_status = 'unknown',
  1849. cleanup_claim = NULL,
  1850. cleanup_claim_expires_at = NULL
  1851. WHERE rule_run_id = CAST(:id AS uuid)
  1852. """
  1853. ),
  1854. {"id": sample_crash_run_id},
  1855. )
  1856. missing_reconciliation = evidence_writer.reconcile_samples(
  1857. limit=10
  1858. )
  1859. assert missing_reconciliation["failed"] == 1
  1860. with platform.connect() as connection:
  1861. assert connection.execute(
  1862. text(
  1863. """
  1864. SELECT handoff_status
  1865. FROM public.rule_violation_samples
  1866. WHERE rule_run_id = CAST(:id AS uuid)
  1867. """
  1868. ),
  1869. {"id": sample_crash_run_id},
  1870. ).scalar_one() == "failed"
  1871. class TransientArtifactStore:
  1872. def __init__(self, wrapped):
  1873. self.wrapped = wrapped
  1874. def __getattr__(self, name):
  1875. return getattr(self.wrapped, name)
  1876. def describe_optional(self, _ref):
  1877. raise TimeoutError("temporary MinIO timeout")
  1878. with platform.begin() as connection:
  1879. transient_sample_id = connection.execute(
  1880. text(
  1881. """
  1882. UPDATE public.rule_violation_samples s
  1883. SET handoff_status = 'unknown',
  1884. cleanup_claim = NULL,
  1885. cleanup_claim_expires_at = NULL
  1886. FROM public.rule_runs r
  1887. WHERE s.rule_run_id = r.id
  1888. AND r.correlation_id =
  1889. CAST(:correlation_id AS uuid)
  1890. RETURNING s.id::text
  1891. """
  1892. ),
  1893. {"correlation_id": correlation_id},
  1894. ).scalar_one()
  1895. transient_writer = PostgresRuleEvidenceWriter(
  1896. platform,
  1897. TransientArtifactStore(store),
  1898. sample_ttl_seconds=900,
  1899. )
  1900. transient_result = transient_writer.reconcile_samples(
  1901. limit=10
  1902. )
  1903. assert transient_result["claimed"] == 1
  1904. assert transient_result["ready"] == 0
  1905. assert transient_result["failed"] == 0
  1906. with platform.connect() as connection:
  1907. transient_state = connection.execute(
  1908. text(
  1909. """
  1910. SELECT handoff_status,
  1911. cleanup_claim,
  1912. cleanup_claim_expires_at
  1913. FROM public.rule_violation_samples
  1914. WHERE id = CAST(:id AS uuid)
  1915. """
  1916. ),
  1917. {"id": transient_sample_id},
  1918. ).mappings().one()
  1919. assert transient_state["handoff_status"] == "unknown"
  1920. assert transient_state["cleanup_claim"] is None
  1921. assert transient_state["cleanup_claim_expires_at"] is None
  1922. assert evidence_writer.reconcile_samples(limit=10)["ready"] == 1
  1923. with platform.begin() as connection:
  1924. connection.execute(
  1925. text(
  1926. """
  1927. INSERT INTO public.dataflow_dataset_bindings
  1928. (id, dataflow_deployment_id, logical_ref,
  1929. data_source_uid, object_kind, object_ref,
  1930. schema_snapshot_id, dialect, access_mode, write_mode,
  1931. binding_hash)
  1932. VALUES (
  1933. CAST(:id AS uuid), CAST(:deployment_id AS uuid),
  1934. 'sql_receipt', CAST(:source_uid AS uuid), 'table',
  1935. 'public.task6_sql_receipt',
  1936. CAST(:schema_snapshot_id AS uuid), 'postgresql',
  1937. 'read_write', 'upsert', :binding_hash
  1938. )
  1939. """
  1940. ),
  1941. {
  1942. "id": sql_binding_id,
  1943. "deployment_id": deployment_id,
  1944. "source_uid": output_binding["data_source_uid"],
  1945. "schema_snapshot_id": output_schema["id"],
  1946. "binding_hash": "b" * 64,
  1947. },
  1948. )
  1949. receipt_run_id = evidence_writer.start(
  1950. component_binding_id=component_binding_id,
  1951. rule_version_id=rule_id,
  1952. plan_hash=compiled["plan_hash"],
  1953. correlation_id=receipt_correlation_id,
  1954. dataflow_uid=dataflow_uid,
  1955. deployment_id=deployment_id,
  1956. environment="test",
  1957. workflow_version=1,
  1958. node_id="task5_real_polars",
  1959. lease_owner=new_governance_uid(),
  1960. )
  1961. receipt = evidence_writer.stage_sql_output(
  1962. receipt_run_id,
  1963. output_binding_id=sql_binding_id,
  1964. )
  1965. assert re.fullmatch(
  1966. r"dataops-staging://[0-9a-f-]{36}",
  1967. receipt,
  1968. )
  1969. with pytest.raises(ValueError, match="not executable"):
  1970. evidence_writer.resolve_sql_staging(
  1971. receipt,
  1972. deployment_id=deployment_id,
  1973. correlation_id=receipt_correlation_id,
  1974. input_binding_id=sql_binding_id,
  1975. )
  1976. receipt_evidence = {
  1977. "status": "success",
  1978. "rows_in": 2,
  1979. "rows_out": 2,
  1980. "rows_rejected": 0,
  1981. "rows_quarantined": 0,
  1982. "commit_outcome": "committed",
  1983. "timings": {"duration_ms": 1},
  1984. "public_result": {
  1985. "rows_in": 2,
  1986. "rows_out": 2,
  1987. "rows_rejected": 0,
  1988. "rows_quarantined": 0,
  1989. "commit_outcome": "committed",
  1990. },
  1991. }
  1992. evidence_writer.finish(receipt_run_id, receipt_evidence)
  1993. evidence_writer.finish(receipt_run_id, receipt_evidence)
  1994. with pytest.raises(ValueError, match="immutable"):
  1995. evidence_writer.finish(
  1996. receipt_run_id,
  1997. {
  1998. **receipt_evidence,
  1999. "rows_out": 1,
  2000. },
  2001. )
  2002. resolved_receipt = evidence_writer.resolve_sql_staging(
  2003. receipt,
  2004. deployment_id=deployment_id,
  2005. correlation_id=receipt_correlation_id,
  2006. input_binding_id=sql_binding_id,
  2007. )
  2008. assert resolved_receipt["relation_ref"] == (
  2009. "public.task6_sql_receipt"
  2010. )
  2011. assert len(resolved_receipt["relation_digest"]) == 64
  2012. with pytest.raises(ValueError, match="not executable"):
  2013. evidence_writer.resolve_sql_staging(
  2014. receipt,
  2015. deployment_id=deployment_id,
  2016. correlation_id=new_governance_uid(),
  2017. input_binding_id=sql_binding_id,
  2018. )
  2019. with pytest.raises(ValueError, match="invalid"):
  2020. evidence_writer.resolve_sql_staging(
  2021. "dataops-staging://public.task6_sql_receipt",
  2022. deployment_id=deployment_id,
  2023. correlation_id=receipt_correlation_id,
  2024. input_binding_id=sql_binding_id,
  2025. )
  2026. failed_receipt_run_id = evidence_writer.start(
  2027. component_binding_id=component_binding_id,
  2028. rule_version_id=rule_id,
  2029. plan_hash=compiled["plan_hash"],
  2030. correlation_id=failed_receipt_correlation_id,
  2031. dataflow_uid=dataflow_uid,
  2032. deployment_id=deployment_id,
  2033. environment="test",
  2034. workflow_version=1,
  2035. node_id="task5_real_polars",
  2036. lease_owner=new_governance_uid(),
  2037. )
  2038. failed_receipt = evidence_writer.stage_sql_output(
  2039. failed_receipt_run_id,
  2040. output_binding_id=sql_binding_id,
  2041. )
  2042. evidence_writer.finish(
  2043. failed_receipt_run_id,
  2044. {
  2045. "status": "failed",
  2046. "commit_outcome": "not_committed",
  2047. "timings": {"duration_ms": 1},
  2048. },
  2049. )
  2050. with pytest.raises(ValueError, match="not executable"):
  2051. evidence_writer.resolve_sql_staging(
  2052. failed_receipt,
  2053. deployment_id=deployment_id,
  2054. correlation_id=failed_receipt_correlation_id,
  2055. input_binding_id=sql_binding_id,
  2056. )
  2057. with platform.begin() as connection:
  2058. connection.execute(
  2059. text(
  2060. """
  2061. UPDATE public.rule_sql_staging_receipts
  2062. SET expires_at =
  2063. CURRENT_TIMESTAMP - INTERVAL '1 second',
  2064. cleanup_claim = CASE
  2065. WHEN correlation_id =
  2066. CAST(:ready AS uuid)
  2067. THEN CAST(:cleanup_claim AS uuid)
  2068. ELSE NULL
  2069. END,
  2070. cleanup_claim_expires_at = CASE
  2071. WHEN correlation_id =
  2072. CAST(:ready AS uuid)
  2073. THEN CURRENT_TIMESTAMP
  2074. + INTERVAL '5 minutes'
  2075. ELSE NULL
  2076. END
  2077. WHERE correlation_id IN (
  2078. CAST(:ready AS uuid),
  2079. CAST(:failed AS uuid)
  2080. )
  2081. """
  2082. ),
  2083. {
  2084. "ready": receipt_correlation_id,
  2085. "failed": failed_receipt_correlation_id,
  2086. "cleanup_claim": new_governance_uid(),
  2087. },
  2088. )
  2089. assert evidence_writer.cleanup_sql_staging(limit=10) == 1
  2090. with platform.connect() as connection:
  2091. assert connection.execute(
  2092. text(
  2093. """
  2094. SELECT status
  2095. FROM public.rule_sql_staging_receipts
  2096. WHERE correlation_id = CAST(:ready AS uuid)
  2097. """
  2098. ),
  2099. {"ready": receipt_correlation_id},
  2100. ).scalar_one() == "ready"
  2101. with platform.begin() as connection:
  2102. connection.execute(
  2103. text(
  2104. """
  2105. UPDATE public.rule_sql_staging_receipts
  2106. SET cleanup_claim_expires_at =
  2107. CURRENT_TIMESTAMP - INTERVAL '1 second'
  2108. WHERE correlation_id = CAST(:ready AS uuid)
  2109. """
  2110. ),
  2111. {"ready": receipt_correlation_id},
  2112. )
  2113. assert evidence_writer.cleanup_sql_staging(limit=10) == 1
  2114. with platform.connect() as connection:
  2115. assert connection.execute(
  2116. text(
  2117. """
  2118. SELECT COUNT(*)
  2119. FROM public.rule_sql_staging_receipts
  2120. WHERE correlation_id IN (
  2121. CAST(:ready AS uuid),
  2122. CAST(:failed AS uuid)
  2123. )
  2124. AND status = 'expired'
  2125. """
  2126. ),
  2127. {
  2128. "ready": receipt_correlation_id,
  2129. "failed": failed_receipt_correlation_id,
  2130. },
  2131. ).scalar_one() == 2
  2132. conflict_path = tmp_path / "conflict.parquet"
  2133. pl.DataFrame(
  2134. {
  2135. "customer_id": [99],
  2136. "mobile": ["13800138000"],
  2137. "name": ["Conflict"],
  2138. "segment_code": ["A"],
  2139. "segment_name": ["Gold"],
  2140. "version_no": [1],
  2141. }
  2142. ).write_parquet(conflict_path)
  2143. object_count_before_conflict = len(
  2144. list(minio.list_objects(bucket, prefix=prefix, recursive=True))
  2145. )
  2146. with pytest.raises(ValueError, match="immutable|digest"):
  2147. resolver.publish_path(
  2148. str(conflict_path),
  2149. binding_id=output_binding["id"],
  2150. binding_hash=compiled["plan"]["output_binding_hash"],
  2151. correlation_id=correlation_id,
  2152. kind="output",
  2153. ttl_seconds=900,
  2154. schema_fields=output_schema["fields"],
  2155. limits=compiled["plan"]["resource_limits"],
  2156. )
  2157. assert len(
  2158. list(minio.list_objects(bucket, prefix=prefix, recursive=True))
  2159. ) == object_count_before_conflict
  2160. assert customer_artifact["digest"]
  2161. assert segment_artifact["digest"]
  2162. with platform.connect() as connection:
  2163. catalog_rows = connection.execute(
  2164. text(
  2165. """
  2166. SELECT artifact_kind, artifact_ref, handoff_status,
  2167. binding_hash
  2168. FROM public.rule_run_artifacts
  2169. WHERE correlation_id = CAST(:correlation_id AS uuid)
  2170. ORDER BY artifact_kind
  2171. """
  2172. ),
  2173. {"correlation_id": correlation_id},
  2174. ).mappings().all()
  2175. # The repeated deterministic output has the same digest and is
  2176. # idempotently retained as one stable catalog handoff.
  2177. assert len(catalog_rows) == 4
  2178. assert all(row["handoff_status"] == "ready" for row in catalog_rows)
  2179. assert all(len(row["binding_hash"]) == 64 for row in catalog_rows)
  2180. assert {
  2181. row["artifact_ref"]
  2182. for row in catalog_rows
  2183. if row["artifact_kind"] == "output"
  2184. } == {
  2185. result["artifact_ref"],
  2186. downstream_result["artifact_ref"],
  2187. }
  2188. assert len(
  2189. list(minio.list_objects(bucket, prefix=prefix, recursive=True))
  2190. ) == 5
  2191. finally:
  2192. for item in list(
  2193. minio.list_objects(bucket, prefix=prefix, recursive=True)
  2194. ):
  2195. minio.remove_object(bucket, item.object_name)
  2196. assert list(
  2197. minio.list_objects(bucket, prefix=prefix, recursive=True)
  2198. ) == []
  2199. with postgres.begin() as connection:
  2200. connection.execute(text(f"DROP TABLE IF EXISTS {customer_table}"))
  2201. with mysql.begin() as connection:
  2202. connection.execute(text(f"DROP TABLE IF EXISTS {segment_table}"))
  2203. with platform.begin() as connection:
  2204. if ledger_jti is not None:
  2205. connection.execute(
  2206. text(
  2207. "DELETE FROM public.runner_task_executions "
  2208. "WHERE token_jti = CAST(:jti AS uuid)"
  2209. ),
  2210. {"jti": ledger_jti},
  2211. )
  2212. if retry_ledger_jti is not None:
  2213. connection.execute(
  2214. text(
  2215. "DELETE FROM public.runner_task_executions "
  2216. "WHERE token_jti = CAST(:jti AS uuid)"
  2217. ),
  2218. {"jti": retry_ledger_jti},
  2219. )
  2220. for crash_jti in (
  2221. no_run_crash_jti,
  2222. evidence_crash_jti,
  2223. active_retry_jti,
  2224. ):
  2225. if crash_jti is not None:
  2226. connection.execute(
  2227. text(
  2228. "DELETE FROM public.runner_task_executions "
  2229. "WHERE token_jti = CAST(:jti AS uuid)"
  2230. ),
  2231. {"jti": crash_jti},
  2232. )
  2233. connection.execute(
  2234. text(
  2235. """
  2236. DELETE FROM public.rule_sql_staging_receipts
  2237. WHERE correlation_id IN (
  2238. CAST(:receipt_correlation_id AS uuid),
  2239. CAST(:failed_receipt_correlation_id AS uuid)
  2240. )
  2241. """
  2242. ),
  2243. {
  2244. "receipt_correlation_id": receipt_correlation_id,
  2245. "failed_receipt_correlation_id": (
  2246. failed_receipt_correlation_id
  2247. ),
  2248. },
  2249. )
  2250. connection.execute(
  2251. text(
  2252. """
  2253. DELETE FROM public.rule_violation_samples s
  2254. USING public.rule_runs r
  2255. WHERE s.rule_run_id = r.id
  2256. AND r.correlation_id IN (
  2257. CAST(:correlation_id AS uuid),
  2258. CAST(:sample_crash_correlation_id AS uuid)
  2259. )
  2260. """
  2261. ),
  2262. {
  2263. "correlation_id": correlation_id,
  2264. "sample_crash_correlation_id": (
  2265. sample_crash_correlation_id
  2266. ),
  2267. },
  2268. )
  2269. connection.execute(
  2270. text(
  2271. "DELETE FROM public.rule_runs "
  2272. "WHERE correlation_id IN ("
  2273. "CAST(:correlation_id AS uuid), "
  2274. "CAST(:failure_correlation_id AS uuid), "
  2275. "CAST(:unknown_correlation_id AS uuid), "
  2276. "CAST(:lease_correlation_id AS uuid), "
  2277. "CAST(:sample_crash_correlation_id AS uuid), "
  2278. "CAST(:receipt_correlation_id AS uuid), "
  2279. "CAST(:failed_receipt_correlation_id AS uuid), "
  2280. "CAST(:no_run_crash_correlation_id AS uuid), "
  2281. "CAST(:evidence_crash_correlation_id AS uuid), "
  2282. "CAST(:active_retry_correlation_id AS uuid))"
  2283. ),
  2284. {
  2285. "correlation_id": correlation_id,
  2286. "failure_correlation_id": failure_correlation_id,
  2287. "unknown_correlation_id": unknown_correlation_id,
  2288. "lease_correlation_id": lease_correlation_id,
  2289. "sample_crash_correlation_id": (
  2290. sample_crash_correlation_id
  2291. ),
  2292. "receipt_correlation_id": receipt_correlation_id,
  2293. "failed_receipt_correlation_id": (
  2294. failed_receipt_correlation_id
  2295. ),
  2296. "no_run_crash_correlation_id": (
  2297. no_run_crash_correlation_id
  2298. ),
  2299. "evidence_crash_correlation_id": (
  2300. evidence_crash_correlation_id
  2301. ),
  2302. "active_retry_correlation_id": (
  2303. active_retry_correlation_id
  2304. ),
  2305. },
  2306. )
  2307. connection.execute(
  2308. text(
  2309. "DELETE FROM public.dataflow_dataset_bindings "
  2310. "WHERE id = CAST(:id AS uuid)"
  2311. ),
  2312. {"id": sql_binding_id},
  2313. )
  2314. connection.execute(
  2315. text(
  2316. "DELETE FROM public.rule_run_artifacts "
  2317. "WHERE correlation_id = CAST(:correlation_id AS uuid)"
  2318. ),
  2319. {"correlation_id": correlation_id},
  2320. )
  2321. connection.execute(
  2322. text(
  2323. "DELETE FROM public.rule_publication_audits "
  2324. "WHERE rule_execution_plan_id IN (CAST(:id AS uuid), "
  2325. "CAST(:downstream_id AS uuid))"
  2326. ),
  2327. {
  2328. "id": plan_id,
  2329. "downstream_id": downstream_plan_id,
  2330. },
  2331. )
  2332. for evidence_table in (
  2333. "rule_test_evidence",
  2334. "rule_compile_evidence",
  2335. ):
  2336. connection.execute(
  2337. text(
  2338. f"DELETE FROM public.{evidence_table} "
  2339. "WHERE rule_execution_plan_id IN "
  2340. "(CAST(:id AS uuid), CAST(:downstream_id AS uuid))"
  2341. ),
  2342. {
  2343. "id": plan_id,
  2344. "downstream_id": downstream_plan_id,
  2345. },
  2346. )
  2347. connection.execute(
  2348. text(
  2349. "DELETE FROM public.rule_execution_plans "
  2350. "WHERE id IN (CAST(:id AS uuid), "
  2351. "CAST(:downstream_id AS uuid))"
  2352. ),
  2353. {
  2354. "id": plan_id,
  2355. "downstream_id": downstream_plan_id,
  2356. },
  2357. )
  2358. connection.execute(
  2359. text(
  2360. "DELETE FROM public.dataflow_dataset_bindings "
  2361. "WHERE dataflow_deployment_id = CAST(:id AS uuid)"
  2362. ),
  2363. {"id": deployment_id},
  2364. )
  2365. connection.execute(
  2366. text(
  2367. "DELETE FROM public.dataflow_component_bindings "
  2368. "WHERE id IN (CAST(:id AS uuid), "
  2369. "CAST(:downstream_id AS uuid))"
  2370. ),
  2371. {
  2372. "id": component_binding_id,
  2373. "downstream_id": downstream_component_binding_id,
  2374. },
  2375. )
  2376. connection.execute(
  2377. text(
  2378. "DELETE FROM public.dataflow_deployments "
  2379. "WHERE id = CAST(:id AS uuid)"
  2380. ),
  2381. {"id": deployment_id},
  2382. )
  2383. if gateway_candidate_id is not None:
  2384. connection.execute(
  2385. text(
  2386. "DELETE FROM public.workflow_canary_evidence "
  2387. "WHERE workflow_version_id = CAST(:id AS uuid)"
  2388. ),
  2389. {"id": gateway_candidate_id},
  2390. )
  2391. connection.execute(
  2392. text(
  2393. "DELETE FROM public.workflow_schedules "
  2394. "WHERE workflow_version_id = CAST(:id AS uuid)"
  2395. ),
  2396. {"id": gateway_candidate_id},
  2397. )
  2398. connection.execute(
  2399. text(
  2400. "DELETE FROM public.workflow_engine_bindings "
  2401. "WHERE workflow_version_id = CAST(:id AS uuid)"
  2402. ),
  2403. {"id": gateway_candidate_id},
  2404. )
  2405. connection.execute(
  2406. text(
  2407. "DELETE FROM public.dataflow_workflow_versions "
  2408. "WHERE id = CAST(:id AS uuid)"
  2409. ),
  2410. {"id": gateway_candidate_id},
  2411. )
  2412. connection.execute(
  2413. text(
  2414. "DELETE FROM public.dataflow_versions "
  2415. "WHERE id = CAST(:id AS uuid)"
  2416. ),
  2417. {"id": dataflow_version_id},
  2418. )
  2419. connection.execute(
  2420. text(
  2421. "DELETE FROM public.rule_logical_test_evidence "
  2422. "WHERE logical_plan_id IN (SELECT id FROM "
  2423. "public.rule_logical_plans WHERE rule_version_id IN "
  2424. "(CAST(:id AS uuid), CAST(:downstream_id AS uuid)))"
  2425. ),
  2426. {
  2427. "id": rule_id,
  2428. "downstream_id": downstream_rule_id,
  2429. },
  2430. )
  2431. connection.execute(
  2432. text(
  2433. "DELETE FROM public.rule_logical_compile_evidence "
  2434. "WHERE logical_plan_id IN (SELECT id FROM "
  2435. "public.rule_logical_plans WHERE rule_version_id IN "
  2436. "(CAST(:id AS uuid), CAST(:downstream_id AS uuid)))"
  2437. ),
  2438. {
  2439. "id": rule_id,
  2440. "downstream_id": downstream_rule_id,
  2441. },
  2442. )
  2443. connection.execute(
  2444. text(
  2445. "DELETE FROM public.rule_logical_plans "
  2446. "WHERE rule_version_id IN "
  2447. "(CAST(:id AS uuid), CAST(:downstream_id AS uuid))"
  2448. ),
  2449. {
  2450. "id": rule_id,
  2451. "downstream_id": downstream_rule_id,
  2452. },
  2453. )
  2454. connection.execute(
  2455. text(
  2456. "DELETE FROM public.rule_validation_profiles "
  2457. "WHERE rule_version_id IN "
  2458. "(CAST(:id AS uuid), CAST(:downstream_id AS uuid))"
  2459. ),
  2460. {
  2461. "id": rule_id,
  2462. "downstream_id": downstream_rule_id,
  2463. },
  2464. )
  2465. connection.execute(
  2466. text(
  2467. "DELETE FROM public.data_rule_versions "
  2468. "WHERE id IN (CAST(:id AS uuid), "
  2469. "CAST(:downstream_id AS uuid))"
  2470. ),
  2471. {
  2472. "id": rule_id,
  2473. "downstream_id": downstream_rule_id,
  2474. },
  2475. )
  2476. connection.execute(
  2477. text(
  2478. "DELETE FROM public.data_rules "
  2479. "WHERE rule_uid IN (CAST(:rule_uid AS uuid), "
  2480. "CAST(:downstream_rule_uid AS uuid))"
  2481. ),
  2482. {
  2483. "rule_uid": rule_uid,
  2484. "downstream_rule_uid": downstream_rule_uid,
  2485. },
  2486. )
  2487. for schema in (input_schema, lookup_schema, output_schema):
  2488. connection.execute(
  2489. text(
  2490. "DELETE FROM public.data_schema_snapshots "
  2491. "WHERE id = CAST(:id AS uuid)"
  2492. ),
  2493. {"id": schema["id"]},
  2494. )
  2495. postgres.dispose()
  2496. mysql.dispose()
  2497. platform.dispose()