test_data_rule_polars_execution.py 48 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278
  1. from __future__ import annotations
  2. import json
  3. import re
  4. from pathlib import Path
  5. import polars as pl
  6. import pytest
  7. from minio import Minio
  8. from sqlalchemy import create_engine, text
  9. from app.core.common.identifiers import new_governance_uid
  10. from app.core.data_rules.contracts import rule_spec_hash, validate_rule_spec
  11. from app.core.data_rules.execution_contracts import canonical_schema_hash
  12. COMPOSE = (
  13. Path(__file__).resolve().parents[2]
  14. / "deploy"
  15. / "docker"
  16. / "docker-compose.yml"
  17. )
  18. def _compose_value(pattern):
  19. source = COMPOSE.read_text(encoding="utf-8")
  20. match = re.search(pattern, source, flags=re.DOTALL)
  21. assert match is not None
  22. return match.group(1)
  23. def _schema(schema_ref, fields):
  24. normalized = [
  25. {"name": name, "type": field_type, "nullable": nullable}
  26. for name, field_type, nullable in fields
  27. ]
  28. return {
  29. "id": new_governance_uid(),
  30. "schema_ref": schema_ref,
  31. "schema_hash": canonical_schema_hash(normalized),
  32. "fields": normalized,
  33. "source_revision": "task5:real-cross-source",
  34. }
  35. def _binding(schema, *, source_uid, access_mode, object_ref):
  36. return {
  37. "id": new_governance_uid(),
  38. "data_source_uid": source_uid,
  39. "object_kind": "parquet_artifact",
  40. "object_ref": object_ref,
  41. "schema_snapshot_id": schema["id"],
  42. "access_mode": access_mode,
  43. "dialect": "parquet",
  44. "write_mode": "append",
  45. }
  46. def test_real_postgres_mysql_minio_polars_cross_source_execution(tmp_path):
  47. from app.core.data_rules.compilers.polars import PolarsRuleCompiler
  48. from app.runner.artifacts import ArtifactStore, PostgresArtifactResolver
  49. from app.runner.rule_polars import PolarsRulePlanAdapter
  50. from app.runner.rules import PostgresRulePlanRepository, RulePlanExecutor
  51. source_user = _compose_value(
  52. r"source-postgres:.*?POSTGRES_USER:\s*([^\s]+)"
  53. )
  54. source_password = _compose_value(
  55. r"source-postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
  56. )
  57. platform_user = _compose_value(
  58. r"\n postgres:.*?POSTGRES_USER:\s*([^\s]+)"
  59. )
  60. platform_password = _compose_value(
  61. r"\n postgres:.*?POSTGRES_PASSWORD:\s*([^\s]+)"
  62. )
  63. postgres_port = _compose_value(r'"(25432):5432"')
  64. mysql_port = _compose_value(r'"(23306):3306"')
  65. minio_user = _compose_value(r"MINIO_ROOT_USER:\s*([^\s]+)")
  66. minio_password = _compose_value(r"MINIO_ROOT_PASSWORD:\s*([^\s]+)")
  67. minio_port = _compose_value(r'"(19000):9000"')
  68. platform_port = _compose_value(r'"(15432):5432"')
  69. bucket = _compose_value(r"mc mb --ignore-existing local/([^\s]+)")
  70. postgres = create_engine(
  71. f"postgresql+psycopg2://{source_user}:{source_password}"
  72. f"@127.0.0.1:{postgres_port}/acceptance",
  73. pool_pre_ping=True,
  74. )
  75. mysql = create_engine(
  76. f"mysql+pymysql://{source_user}:{source_password}"
  77. f"@127.0.0.1:{mysql_port}/acceptance",
  78. pool_pre_ping=True,
  79. )
  80. platform = create_engine(
  81. f"postgresql+psycopg2://{platform_user}:{platform_password}"
  82. f"@127.0.0.1:{platform_port}/dataops",
  83. pool_pre_ping=True,
  84. )
  85. minio = Minio(
  86. f"127.0.0.1:{minio_port}",
  87. access_key=minio_user,
  88. secret_key=minio_password,
  89. secure=False,
  90. )
  91. store = ArtifactStore(
  92. minio,
  93. bucket=bucket,
  94. max_artifact_bytes=4 * 1024 * 1024,
  95. max_rows=1_000,
  96. memory_limit_bytes=256 * 1024 * 1024,
  97. max_ttl_seconds=3600,
  98. )
  99. correlation_id = new_governance_uid()
  100. failure_correlation_id = new_governance_uid()
  101. unknown_correlation_id = new_governance_uid()
  102. prefix = f"rules/{correlation_id}/"
  103. customer_table = "task5_polars_customers"
  104. segment_table = "task5_polars_segments"
  105. rule_uid = new_governance_uid()
  106. rule_id = new_governance_uid()
  107. downstream_rule_uid = new_governance_uid()
  108. downstream_rule_id = new_governance_uid()
  109. dataflow_uid = new_governance_uid()
  110. dataflow_version_id = new_governance_uid()
  111. deployment_id = new_governance_uid()
  112. component_binding_id = new_governance_uid()
  113. plan_id = new_governance_uid()
  114. downstream_component_binding_id = new_governance_uid()
  115. downstream_plan_id = new_governance_uid()
  116. ledger_jti = None
  117. retry_ledger_jti = None
  118. try:
  119. with postgres.begin() as connection:
  120. connection.execute(text(f"DROP TABLE IF EXISTS {customer_table}"))
  121. connection.execute(
  122. text(
  123. f"CREATE TABLE {customer_table} ("
  124. "customer_id BIGINT NOT NULL, "
  125. "name VARCHAR(100), mobile VARCHAR(30), "
  126. "segment_code VARCHAR(20), version_no BIGINT NOT NULL)"
  127. )
  128. )
  129. connection.execute(
  130. text(
  131. f"INSERT INTO {customer_table} "
  132. "(customer_id, name, mobile, segment_code, version_no) "
  133. "VALUES "
  134. "(1, ' Alice ', '13800138000', 'A', 1), "
  135. "(1, ' Alice Updated ', '13800138000', 'A', 2), "
  136. "(2, ' Bad ', 'invalid', 'B', 1), "
  137. "(3, ' Carol ', '13900139000', 'C', 1)"
  138. )
  139. )
  140. with mysql.begin() as connection:
  141. connection.execute(text(f"DROP TABLE IF EXISTS {segment_table}"))
  142. connection.execute(
  143. text(
  144. f"CREATE TABLE {segment_table} ("
  145. "code VARCHAR(20) PRIMARY KEY, "
  146. "segment_name VARCHAR(100) NOT NULL)"
  147. )
  148. )
  149. connection.execute(
  150. text(
  151. f"INSERT INTO {segment_table} (code, segment_name) "
  152. "VALUES ('A', 'Gold'), ('B', 'Basic'), ('C', 'Silver')"
  153. )
  154. )
  155. with postgres.connect() as connection:
  156. customer_rows = [
  157. dict(row)
  158. for row in connection.execute(
  159. text(
  160. f"SELECT customer_id, name, mobile, "
  161. f"segment_code, version_no FROM {customer_table}"
  162. )
  163. ).mappings()
  164. ]
  165. with mysql.connect() as connection:
  166. segment_rows = [
  167. dict(row)
  168. for row in connection.execute(
  169. text(
  170. f"SELECT code, segment_name FROM {segment_table}"
  171. )
  172. ).mappings()
  173. ]
  174. input_schema = _schema(
  175. "bd:task5:customer:raw",
  176. [
  177. ("customer_id", "integer", False),
  178. ("name", "string", True),
  179. ("mobile", "string", True),
  180. ("segment_code", "string", True),
  181. ("version_no", "integer", False),
  182. ],
  183. )
  184. lookup_schema = _schema(
  185. "bd:task5:segment:lookup",
  186. [
  187. ("code", "string", False),
  188. ("segment_name", "string", False),
  189. ],
  190. )
  191. output_schema = _schema(
  192. "bd:task5:customer:enriched",
  193. [
  194. ("customer_id", "integer", False),
  195. ("name", "string", True),
  196. ("mobile", "string", True),
  197. ("segment_code", "string", True),
  198. ("version_no", "integer", False),
  199. ("segment_name", "string", True),
  200. ],
  201. )
  202. input_binding = _binding(
  203. input_schema,
  204. source_uid=new_governance_uid(),
  205. access_mode="read",
  206. object_ref="postgres-customer-artifact",
  207. )
  208. lookup_binding = _binding(
  209. lookup_schema,
  210. source_uid=new_governance_uid(),
  211. access_mode="read",
  212. object_ref="mysql-segment-artifact",
  213. )
  214. output_binding = _binding(
  215. output_schema,
  216. source_uid=new_governance_uid(),
  217. access_mode="read_write",
  218. object_ref="polars-output-artifact",
  219. )
  220. downstream_output_binding = _binding(
  221. output_schema,
  222. source_uid=new_governance_uid(),
  223. access_mode="write",
  224. object_ref="polars-downstream-output-artifact",
  225. )
  226. spec = validate_rule_spec(
  227. {
  228. "schema_version": "2.0",
  229. "rule_uid": new_governance_uid(),
  230. "name": "task5_real_cross_source",
  231. "input_schema_ref": input_schema["schema_ref"],
  232. "output_schema_ref": output_schema["schema_ref"],
  233. "steps": [
  234. {
  235. "id": "normalize_name",
  236. "op": "normalize_text",
  237. "column": "name",
  238. "trim": True,
  239. },
  240. {
  241. "id": "join_segment",
  242. "op": "lookup_join",
  243. "lookup": {
  244. "binding_id": lookup_binding["id"],
  245. "left_on": ["segment_code"],
  246. "right_on": ["code"],
  247. "select": {
  248. "segment_name": "segment_name"
  249. },
  250. "how": "left",
  251. },
  252. },
  253. {
  254. "id": "valid_mobile",
  255. "op": "assert",
  256. "expression": "matches(mobile, '^[0-9]{11}$')",
  257. "on_failure": "reject",
  258. "severity": "error",
  259. },
  260. {
  261. "id": "latest_customer",
  262. "op": "deduplicate",
  263. "keys": ["customer_id"],
  264. "order_by": ["version_no"],
  265. "keep": "last",
  266. },
  267. ],
  268. "null_policy": "explicit",
  269. "timezone": "Asia/Shanghai",
  270. }
  271. )
  272. rule = {
  273. "id": rule_id,
  274. "status": "published",
  275. "rule_spec": spec,
  276. "spec_hash": rule_spec_hash(spec),
  277. }
  278. compiled = PolarsRuleCompiler().compile(
  279. rule_version=rule,
  280. input_schema=input_schema,
  281. output_schema=output_schema,
  282. input_binding=input_binding,
  283. output_binding=output_binding,
  284. backend={
  285. "max_rows": 1_000,
  286. "max_artifact_bytes": 4 * 1024 * 1024,
  287. "memory_limit_bytes": 256 * 1024 * 1024,
  288. "masking_policies": {},
  289. "lookup_bindings": {
  290. lookup_binding["id"]: {
  291. "binding": lookup_binding,
  292. "schema": lookup_schema,
  293. }
  294. },
  295. },
  296. )
  297. downstream_spec = validate_rule_spec(
  298. {
  299. "schema_version": "2.0",
  300. "rule_uid": new_governance_uid(),
  301. "name": "task6_real_artifact_handoff",
  302. "input_schema_ref": output_schema["schema_ref"],
  303. "output_schema_ref": output_schema["schema_ref"],
  304. "steps": [
  305. {
  306. "id": "normalize_downstream_name",
  307. "op": "normalize_text",
  308. "column": "name",
  309. "trim": True,
  310. }
  311. ],
  312. "null_policy": "explicit",
  313. "timezone": "Asia/Shanghai",
  314. }
  315. )
  316. downstream_rule = {
  317. "id": downstream_rule_id,
  318. "status": "published",
  319. "rule_spec": downstream_spec,
  320. "spec_hash": rule_spec_hash(downstream_spec),
  321. }
  322. downstream_compiled = PolarsRuleCompiler().compile(
  323. rule_version=downstream_rule,
  324. input_schema=output_schema,
  325. output_schema=output_schema,
  326. input_binding=output_binding,
  327. output_binding=downstream_output_binding,
  328. backend={
  329. "max_rows": 1_000,
  330. "max_artifact_bytes": 4 * 1024 * 1024,
  331. "memory_limit_bytes": 256 * 1024 * 1024,
  332. "masking_policies": {},
  333. "lookup_bindings": {},
  334. },
  335. )
  336. lookup_operation = compiled["plan"]["operations"][1]
  337. schema_hashes = {
  338. "rule_spec_hash": compiled["plan"]["rule_spec_hash"],
  339. "input_schema_snapshot_id": input_schema["id"],
  340. "input_schema_hash": input_schema["schema_hash"],
  341. "output_schema_snapshot_id": output_schema["id"],
  342. "output_schema_hash": output_schema["schema_hash"],
  343. }
  344. with platform.begin() as connection:
  345. for schema in (input_schema, lookup_schema, output_schema):
  346. connection.execute(
  347. text(
  348. """
  349. INSERT INTO public.data_schema_snapshots
  350. (id, schema_ref, schema_hash, fields, source_revision)
  351. VALUES (CAST(:id AS uuid), :schema_ref, :schema_hash,
  352. CAST(:fields AS jsonb), :source_revision)
  353. """
  354. ),
  355. {**schema, "fields": json.dumps(schema["fields"])},
  356. )
  357. connection.execute(
  358. text(
  359. """
  360. INSERT INTO public.data_rules
  361. (id, rule_uid, name, category, status)
  362. VALUES (CAST(:id AS uuid), CAST(:rule_uid AS uuid),
  363. :name, 'general', 'active')
  364. """
  365. ),
  366. {
  367. "id": new_governance_uid(),
  368. "rule_uid": rule_uid,
  369. "name": spec["name"],
  370. },
  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": downstream_rule_uid,
  384. "name": downstream_spec["name"],
  385. },
  386. )
  387. connection.execute(
  388. text(
  389. """
  390. INSERT INTO public.data_rule_versions
  391. (id, rule_uid, version_no, source_text, source_language,
  392. rule_spec, spec_hash, generated_kind, status, published_at)
  393. VALUES (CAST(:id AS uuid), CAST(:rule_uid AS uuid), 1,
  394. :source_text, 'en', CAST(:rule_spec AS jsonb),
  395. :spec_hash, 'polars', 'published',
  396. CURRENT_TIMESTAMP)
  397. """
  398. ),
  399. {
  400. "id": rule_id,
  401. "rule_uid": rule_uid,
  402. "source_text": "Task 5 real cross-source integration",
  403. "rule_spec": json.dumps(spec),
  404. "spec_hash": rule["spec_hash"],
  405. },
  406. )
  407. connection.execute(
  408. text(
  409. """
  410. INSERT INTO public.data_rule_versions
  411. (id, rule_uid, version_no, source_text,
  412. source_language, rule_spec, spec_hash,
  413. generated_kind, status, published_at)
  414. VALUES (
  415. CAST(:id AS uuid), CAST(:rule_uid AS uuid), 1,
  416. :source_text, 'en', CAST(:rule_spec AS jsonb),
  417. :spec_hash, 'polars', 'published',
  418. CURRENT_TIMESTAMP
  419. )
  420. """
  421. ),
  422. {
  423. "id": downstream_rule_id,
  424. "rule_uid": downstream_rule_uid,
  425. "source_text": (
  426. "Task 6 real two-node artifact handoff"
  427. ),
  428. "rule_spec": json.dumps(downstream_spec),
  429. "spec_hash": downstream_rule["spec_hash"],
  430. },
  431. )
  432. connection.execute(
  433. text(
  434. """
  435. INSERT INTO public.dataflow_versions
  436. (id, dataflow_uid, version_no, name, dataflow_spec,
  437. input_schema_hashes, output_schema_hash, status,
  438. released_at)
  439. VALUES (CAST(:id AS uuid), CAST(:dataflow_uid AS uuid), 1,
  440. :name, '{}'::jsonb,
  441. CAST(:input_schema_hashes AS jsonb),
  442. :output_schema_hash, 'released', CURRENT_TIMESTAMP)
  443. """
  444. ),
  445. {
  446. "id": dataflow_version_id,
  447. "dataflow_uid": dataflow_uid,
  448. "name": "Task 5 real cross-source integration",
  449. "input_schema_hashes": json.dumps(
  450. [
  451. input_schema["schema_hash"],
  452. lookup_schema["schema_hash"],
  453. ]
  454. ),
  455. "output_schema_hash": output_schema["schema_hash"],
  456. },
  457. )
  458. connection.execute(
  459. text(
  460. """
  461. INSERT INTO public.dataflow_deployments
  462. (id, dataflow_version_id, environment, deployment_config,
  463. status, activated_at)
  464. VALUES (CAST(:id AS uuid),
  465. CAST(:dataflow_version_id AS uuid), 'test',
  466. '{}'::jsonb, 'active', CURRENT_TIMESTAMP)
  467. """
  468. ),
  469. {
  470. "id": deployment_id,
  471. "dataflow_version_id": dataflow_version_id,
  472. },
  473. )
  474. for logical_ref, binding, binding_hash in (
  475. (
  476. "customers",
  477. input_binding,
  478. compiled["plan"]["input_binding_hash"],
  479. ),
  480. (
  481. "segments",
  482. lookup_binding,
  483. lookup_operation["lookup_binding_hash"],
  484. ),
  485. (
  486. "enriched",
  487. output_binding,
  488. compiled["plan"]["output_binding_hash"],
  489. ),
  490. (
  491. "downstream",
  492. downstream_output_binding,
  493. downstream_compiled["plan"][
  494. "output_binding_hash"
  495. ],
  496. ),
  497. ):
  498. connection.execute(
  499. text(
  500. """
  501. INSERT INTO public.dataflow_dataset_bindings
  502. (id, dataflow_deployment_id, logical_ref,
  503. data_source_uid, object_kind, object_ref,
  504. schema_snapshot_id, dialect, access_mode, write_mode,
  505. binding_hash)
  506. VALUES (CAST(:id AS uuid), CAST(:deployment_id AS uuid),
  507. :logical_ref, CAST(:source_uid AS uuid),
  508. 'parquet_artifact', :object_ref,
  509. CAST(:schema_snapshot_id AS uuid), 'parquet',
  510. :access_mode, 'append', :binding_hash)
  511. """
  512. ),
  513. {
  514. "id": binding["id"],
  515. "deployment_id": deployment_id,
  516. "logical_ref": logical_ref,
  517. "source_uid": binding["data_source_uid"],
  518. "object_ref": binding["object_ref"],
  519. "schema_snapshot_id": binding["schema_snapshot_id"],
  520. "access_mode": binding["access_mode"],
  521. "binding_hash": binding_hash,
  522. },
  523. )
  524. connection.execute(
  525. text(
  526. """
  527. INSERT INTO public.dataflow_component_bindings
  528. (id, dataflow_version_id, component_id, component_kind,
  529. rule_version_id, stage, order_no, idempotency, provenance)
  530. VALUES (CAST(:id AS uuid),
  531. CAST(:dataflow_version_id AS uuid),
  532. 'task5_real_polars', 'rule.apply',
  533. CAST(:rule_version_id AS uuid), 'transform', 0,
  534. CAST(:idempotency AS jsonb), '{}'::jsonb)
  535. """
  536. ),
  537. {
  538. "id": component_binding_id,
  539. "dataflow_version_id": dataflow_version_id,
  540. "rule_version_id": rule_id,
  541. "idempotency": json.dumps(
  542. {
  543. "strategy": "deduplication_key",
  544. "key": "customer_id",
  545. }
  546. ),
  547. },
  548. )
  549. connection.execute(
  550. text(
  551. """
  552. INSERT INTO public.dataflow_component_bindings
  553. (id, dataflow_version_id, component_id,
  554. component_kind, rule_version_id, stage, order_no,
  555. idempotency, provenance)
  556. VALUES (
  557. CAST(:id AS uuid),
  558. CAST(:dataflow_version_id AS uuid),
  559. 'task6_real_handoff', 'rule.apply',
  560. CAST(:rule_version_id AS uuid), 'transform', 1,
  561. CAST(:idempotency AS jsonb), '{}'::jsonb
  562. )
  563. """
  564. ),
  565. {
  566. "id": downstream_component_binding_id,
  567. "dataflow_version_id": dataflow_version_id,
  568. "rule_version_id": downstream_rule_id,
  569. "idempotency": json.dumps(
  570. {
  571. "strategy": "deduplication_key",
  572. "key": "customer_id",
  573. }
  574. ),
  575. },
  576. )
  577. connection.execute(
  578. text(
  579. """
  580. INSERT INTO public.rule_execution_plans
  581. (id, component_binding_id, backend, compiler_version, plan,
  582. plan_hash, schema_hashes, status)
  583. VALUES (CAST(:id AS uuid),
  584. CAST(:component_binding_id AS uuid),
  585. 'polars_batch', :compiler_version,
  586. CAST(:plan AS jsonb), :plan_hash,
  587. CAST(:schema_hashes AS jsonb), 'published')
  588. """
  589. ),
  590. {
  591. "id": plan_id,
  592. "component_binding_id": component_binding_id,
  593. "compiler_version": compiled["compiler_version"],
  594. "plan": json.dumps(compiled["plan"]),
  595. "plan_hash": compiled["plan_hash"],
  596. "schema_hashes": json.dumps(schema_hashes),
  597. },
  598. )
  599. connection.execute(
  600. text(
  601. """
  602. INSERT INTO public.rule_execution_plans
  603. (id, component_binding_id, backend,
  604. compiler_version, plan, plan_hash,
  605. schema_hashes, status)
  606. VALUES (
  607. CAST(:id AS uuid),
  608. CAST(:component_binding_id AS uuid),
  609. 'polars_batch', :compiler_version,
  610. CAST(:plan AS jsonb), :plan_hash,
  611. CAST(:schema_hashes AS jsonb), 'published'
  612. )
  613. """
  614. ),
  615. {
  616. "id": downstream_plan_id,
  617. "component_binding_id": (
  618. downstream_component_binding_id
  619. ),
  620. "compiler_version": downstream_compiled[
  621. "compiler_version"
  622. ],
  623. "plan": json.dumps(downstream_compiled["plan"]),
  624. "plan_hash": downstream_compiled["plan_hash"],
  625. "schema_hashes": json.dumps(
  626. {
  627. "rule_spec_hash": downstream_compiled[
  628. "plan"
  629. ]["rule_spec_hash"],
  630. "input_schema_snapshot_id": output_schema[
  631. "id"
  632. ],
  633. "input_schema_hash": output_schema[
  634. "schema_hash"
  635. ],
  636. "output_schema_snapshot_id": output_schema[
  637. "id"
  638. ],
  639. "output_schema_hash": output_schema[
  640. "schema_hash"
  641. ],
  642. }
  643. ),
  644. },
  645. )
  646. customer_path = tmp_path / "customers.parquet"
  647. segment_path = tmp_path / "segments.parquet"
  648. pl.DataFrame(customer_rows).write_parquet(customer_path)
  649. pl.DataFrame(segment_rows).write_parquet(segment_path)
  650. resolver = PostgresArtifactResolver(platform, store)
  651. customer_artifact = resolver.publish_path(
  652. str(customer_path),
  653. binding_id=input_binding["id"],
  654. binding_hash=compiled["plan"]["input_binding_hash"],
  655. correlation_id=correlation_id,
  656. kind="input",
  657. ttl_seconds=900,
  658. schema_fields=input_schema["fields"],
  659. limits=compiled["plan"]["resource_limits"],
  660. )
  661. segment_artifact = resolver.publish_path(
  662. str(segment_path),
  663. binding_id=lookup_binding["id"],
  664. binding_hash=lookup_operation["lookup_binding_hash"],
  665. correlation_id=correlation_id,
  666. kind="lookup",
  667. ttl_seconds=900,
  668. schema_fields=lookup_schema["fields"],
  669. limits=compiled["plan"]["resource_limits"],
  670. )
  671. node = {
  672. "id": "task5_real_polars",
  673. "type": "rule.apply",
  674. "purpose": "write",
  675. "idempotency": {
  676. "strategy": "deduplication_key",
  677. "key": "customer_id",
  678. },
  679. "config": {
  680. "component_binding_id": component_binding_id,
  681. "rule_version_id": rule["id"],
  682. "execution_plan_hash": compiled["plan_hash"],
  683. },
  684. }
  685. executor = RulePlanExecutor(
  686. PostgresRulePlanRepository(platform),
  687. adapters={
  688. "polars_batch": PolarsRulePlanAdapter(
  689. artifact_store=store,
  690. artifact_resolver=resolver,
  691. artifact_ttl_seconds=900,
  692. )
  693. },
  694. )
  695. result = executor.execute(
  696. node,
  697. {},
  698. write_authorized=True,
  699. correlation_id=correlation_id,
  700. )
  701. repeated = executor.execute(
  702. node,
  703. {},
  704. write_authorized=True,
  705. correlation_id=correlation_id,
  706. )
  707. assert result["rows_in"] == 4
  708. assert result["rows_out"] == 2
  709. assert result["rows_rejected"] == 1
  710. assert result["rows_deduplicated"] == 1
  711. assert result["rows_filtered"] == 0
  712. assert result["rows_join_dropped"] == 0
  713. assert result["rows_aggregated"] == 0
  714. assert result["violation_count"] == 1
  715. assert result["violations"] == [
  716. {"step_id": "valid_mobile", "count": 1}
  717. ]
  718. output = store.read(
  719. result["artifact_ref"],
  720. result["digest"],
  721. expected_schema_fields=output_schema["fields"],
  722. limits=compiled["plan"]["resource_limits"],
  723. ).collect()
  724. assert output.sort("customer_id").to_dicts() == [
  725. {
  726. "customer_id": 1,
  727. "mobile": "13800138000",
  728. "name": "Alice Updated",
  729. "segment_code": "A",
  730. "segment_name": "Gold",
  731. "version_no": 2,
  732. },
  733. {
  734. "customer_id": 3,
  735. "mobile": "13900139000",
  736. "name": "Carol",
  737. "segment_code": "C",
  738. "segment_name": "Silver",
  739. "version_no": 1,
  740. },
  741. ]
  742. assert all(
  743. item.object_name.startswith(prefix)
  744. for item in minio.list_objects(
  745. bucket, prefix=prefix, recursive=True
  746. )
  747. )
  748. assert repeated["rows_out"] == 2
  749. assert repeated["artifact_ref"] == result["artifact_ref"]
  750. assert "schema_fields" not in result
  751. from app.runner.api import create_runner_app
  752. from app.runner.auth import TaskTokenIssuer, TaskTokenVerifier
  753. from app.runner.ledger import PostgresTaskLedger
  754. from app.runner.nodes import NodeRegistry
  755. from app.runner.rule_evidence import PostgresRuleEvidenceWriter
  756. task_secret = "task5-real-http-secret-value-32-bytes"
  757. verifier = TaskTokenVerifier(task_secret)
  758. task_token = TaskTokenIssuer(task_secret).issue(
  759. task_uid=new_governance_uid(),
  760. dataflow_uid=dataflow_uid,
  761. workflow_version=1,
  762. correlation_id=correlation_id,
  763. node=node,
  764. write_authorized=True,
  765. )
  766. ledger_jti = verifier.verify(task_token, node=node).jti
  767. retry_token = TaskTokenIssuer(task_secret).issue(
  768. task_uid=new_governance_uid(),
  769. dataflow_uid=dataflow_uid,
  770. workflow_version=1,
  771. correlation_id=correlation_id,
  772. node=node,
  773. write_authorized=True,
  774. )
  775. retry_ledger_jti = verifier.verify(
  776. retry_token, node=node
  777. ).jti
  778. evidenced_executor = RulePlanExecutor(
  779. PostgresRulePlanRepository(platform),
  780. adapters={
  781. "polars_batch": PolarsRulePlanAdapter(
  782. artifact_store=store,
  783. artifact_resolver=resolver,
  784. artifact_ttl_seconds=900,
  785. )
  786. },
  787. evidence_writer=PostgresRuleEvidenceWriter(
  788. platform,
  789. store,
  790. sample_ttl_seconds=900,
  791. ),
  792. )
  793. runner_app = create_runner_app(
  794. verifier=verifier,
  795. ledger=PostgresTaskLedger(platform),
  796. registry=NodeRegistry({"rule.apply": evidenced_executor}),
  797. )
  798. with runner_app.test_client() as client:
  799. http_result = client.post(
  800. "/v1/tasks/execute",
  801. json={
  802. "task_token": task_token,
  803. "node": node,
  804. "parameters": {},
  805. },
  806. )
  807. replay = client.post(
  808. "/v1/tasks/execute",
  809. json={
  810. "task_token": task_token,
  811. "node": node,
  812. "parameters": {},
  813. },
  814. )
  815. retried = client.post(
  816. "/v1/tasks/execute",
  817. json={
  818. "task_token": retry_token,
  819. "node": node,
  820. "parameters": {},
  821. },
  822. )
  823. assert http_result.status_code == 200
  824. assert http_result.get_json()["output_artifact"] == result[
  825. "artifact_ref"
  826. ]
  827. assert http_result.get_json()["result"]["artifact_ref"] == result[
  828. "artifact_ref"
  829. ]
  830. assert replay.status_code == 409
  831. assert retried.status_code == 200
  832. assert retried.get_json()["output_artifact"] == result[
  833. "artifact_ref"
  834. ]
  835. ledger_record = PostgresTaskLedger(platform).get(ledger_jti)
  836. assert ledger_record is not None
  837. assert ledger_record.status == "success"
  838. assert ledger_record.commit_outcome == "committed"
  839. with platform.connect() as connection:
  840. run_evidence = connection.execute(
  841. text(
  842. """
  843. SELECT status, commit_outcome, rows_in, rows_out,
  844. rows_rejected, rows_quarantined, public_result
  845. FROM public.rule_runs
  846. WHERE correlation_id = CAST(:correlation_id AS uuid)
  847. AND component_binding_id =
  848. CAST(:component_binding_id AS uuid)
  849. """
  850. ),
  851. {
  852. "correlation_id": correlation_id,
  853. "component_binding_id": component_binding_id,
  854. },
  855. ).mappings().one()
  856. sample_evidence = connection.execute(
  857. text(
  858. """
  859. SELECT s.artifact_ref, s.artifact_digest,
  860. s.sample_count, s.redaction_policy,
  861. s.expires_at, s.handoff_status
  862. FROM public.rule_violation_samples s
  863. JOIN public.rule_runs r ON r.id = s.rule_run_id
  864. WHERE r.correlation_id =
  865. CAST(:correlation_id AS uuid)
  866. """
  867. ),
  868. {"correlation_id": correlation_id},
  869. ).mappings().one()
  870. assert run_evidence["status"] == "success"
  871. assert run_evidence["commit_outcome"] == "committed"
  872. assert run_evidence["rows_in"] == 4
  873. assert run_evidence["rows_out"] == 2
  874. assert run_evidence["rows_rejected"] == 1
  875. assert run_evidence["rows_quarantined"] == 0
  876. assert run_evidence["public_result"]["output_artifact"] == result[
  877. "artifact_ref"
  878. ]
  879. assert sample_evidence["artifact_digest"]
  880. assert sample_evidence["sample_count"] == 1
  881. assert (
  882. sample_evidence["redaction_policy"]
  883. == "rule-violation-default-v1"
  884. )
  885. assert sample_evidence["handoff_status"] == "ready"
  886. assert store.read(
  887. sample_evidence["artifact_ref"],
  888. sample_evidence["artifact_digest"],
  889. expected_schema_fields=[
  890. {
  891. "name": name,
  892. "type": "string",
  893. "nullable": True,
  894. }
  895. for name in (
  896. "customer_id",
  897. "mobile",
  898. "name",
  899. "segment_code",
  900. "segment_name",
  901. "version_no",
  902. )
  903. ],
  904. ).collect().to_dicts() == [
  905. {
  906. "customer_id": "[REDACTED]",
  907. "mobile": "[REDACTED]",
  908. "name": "[REDACTED]",
  909. "segment_code": "[REDACTED]",
  910. "segment_name": "[REDACTED]",
  911. "version_no": "[REDACTED]",
  912. }
  913. ]
  914. downstream_node = {
  915. "id": "task6_real_handoff",
  916. "type": "rule.apply",
  917. "purpose": "write",
  918. "idempotency": {
  919. "strategy": "deduplication_key",
  920. "key": "customer_id",
  921. },
  922. "config": {
  923. "component_binding_id": (
  924. downstream_component_binding_id
  925. ),
  926. "rule_version_id": downstream_rule_id,
  927. "execution_plan_hash": downstream_compiled[
  928. "plan_hash"
  929. ],
  930. },
  931. }
  932. downstream_result = evidenced_executor.execute(
  933. downstream_node,
  934. {"input_artifact": result["artifact_ref"]},
  935. write_authorized=True,
  936. correlation_id=correlation_id,
  937. dataflow_uid=dataflow_uid,
  938. workflow_version=1,
  939. node_id="task6_real_handoff",
  940. )
  941. assert downstream_result["rows_in"] == 2
  942. assert downstream_result["rows_out"] == 2
  943. assert downstream_result["output_artifact"] != result[
  944. "artifact_ref"
  945. ]
  946. assert store.read(
  947. downstream_result["artifact_ref"],
  948. downstream_result["digest"],
  949. expected_schema_fields=output_schema["fields"],
  950. limits=downstream_compiled["plan"]["resource_limits"],
  951. ).collect().sort("customer_id").to_dicts() == (
  952. output.sort("customer_id").to_dicts()
  953. )
  954. replayed_downstream = evidenced_executor.execute(
  955. downstream_node,
  956. {"input_artifact": result["artifact_ref"]},
  957. write_authorized=True,
  958. correlation_id=correlation_id,
  959. dataflow_uid=dataflow_uid,
  960. workflow_version=1,
  961. node_id="task6_real_handoff",
  962. )
  963. assert replayed_downstream["artifact_ref"] == downstream_result[
  964. "artifact_ref"
  965. ]
  966. with platform.connect() as connection:
  967. assert connection.execute(
  968. text(
  969. """
  970. SELECT COUNT(*)
  971. FROM public.rule_runs
  972. WHERE correlation_id =
  973. CAST(:correlation_id AS uuid)
  974. """
  975. ),
  976. {"correlation_id": correlation_id},
  977. ).scalar_one() == 2
  978. from app.runner.nodes import NodeExecutionError
  979. class FailingAdapter:
  980. def execute(self, **_kwargs):
  981. raise NodeExecutionError(
  982. "safe downstream failure",
  983. commit_outcome="not_committed",
  984. )
  985. failed_executor = RulePlanExecutor(
  986. PostgresRulePlanRepository(platform),
  987. adapters={"polars_batch": FailingAdapter()},
  988. evidence_writer=PostgresRuleEvidenceWriter(
  989. platform,
  990. store,
  991. sample_ttl_seconds=900,
  992. ),
  993. )
  994. with pytest.raises(NodeExecutionError, match="safe downstream"):
  995. failed_executor.execute(
  996. node,
  997. {},
  998. write_authorized=True,
  999. correlation_id=failure_correlation_id,
  1000. dataflow_uid=dataflow_uid,
  1001. workflow_version=1,
  1002. node_id="task5_real_polars",
  1003. )
  1004. with platform.connect() as connection:
  1005. failed_evidence = connection.execute(
  1006. text(
  1007. """
  1008. SELECT status, commit_outcome
  1009. FROM public.rule_runs
  1010. WHERE correlation_id =
  1011. CAST(:correlation_id AS uuid)
  1012. """
  1013. ),
  1014. {"correlation_id": failure_correlation_id},
  1015. ).mappings().one()
  1016. assert dict(failed_evidence) == {
  1017. "status": "failed",
  1018. "commit_outcome": "not_committed",
  1019. }
  1020. class UnknownAdapter:
  1021. def execute(self, **_kwargs):
  1022. raise NodeExecutionError(
  1023. "safe uncertain commit",
  1024. commit_outcome="unknown",
  1025. )
  1026. unknown_executor = RulePlanExecutor(
  1027. PostgresRulePlanRepository(platform),
  1028. adapters={"polars_batch": UnknownAdapter()},
  1029. evidence_writer=PostgresRuleEvidenceWriter(
  1030. platform,
  1031. store,
  1032. sample_ttl_seconds=900,
  1033. ),
  1034. )
  1035. with pytest.raises(NodeExecutionError, match="uncertain commit"):
  1036. unknown_executor.execute(
  1037. node,
  1038. {},
  1039. write_authorized=True,
  1040. correlation_id=unknown_correlation_id,
  1041. dataflow_uid=dataflow_uid,
  1042. workflow_version=1,
  1043. node_id="task5_real_polars",
  1044. )
  1045. with platform.connect() as connection:
  1046. unknown_evidence = connection.execute(
  1047. text(
  1048. """
  1049. SELECT status, commit_outcome
  1050. FROM public.rule_runs
  1051. WHERE correlation_id =
  1052. CAST(:correlation_id AS uuid)
  1053. """
  1054. ),
  1055. {"correlation_id": unknown_correlation_id},
  1056. ).mappings().one()
  1057. assert dict(unknown_evidence) == {
  1058. "status": "unknown",
  1059. "commit_outcome": "unknown",
  1060. }
  1061. conflict_path = tmp_path / "conflict.parquet"
  1062. pl.DataFrame(
  1063. {
  1064. "customer_id": [99],
  1065. "mobile": ["13800138000"],
  1066. "name": ["Conflict"],
  1067. "segment_code": ["A"],
  1068. "segment_name": ["Gold"],
  1069. "version_no": [1],
  1070. }
  1071. ).write_parquet(conflict_path)
  1072. object_count_before_conflict = len(
  1073. list(minio.list_objects(bucket, prefix=prefix, recursive=True))
  1074. )
  1075. with pytest.raises(ValueError, match="immutable|digest"):
  1076. resolver.publish_path(
  1077. str(conflict_path),
  1078. binding_id=output_binding["id"],
  1079. binding_hash=compiled["plan"]["output_binding_hash"],
  1080. correlation_id=correlation_id,
  1081. kind="output",
  1082. ttl_seconds=900,
  1083. schema_fields=output_schema["fields"],
  1084. limits=compiled["plan"]["resource_limits"],
  1085. )
  1086. assert len(
  1087. list(minio.list_objects(bucket, prefix=prefix, recursive=True))
  1088. ) == object_count_before_conflict
  1089. assert customer_artifact["digest"]
  1090. assert segment_artifact["digest"]
  1091. with platform.connect() as connection:
  1092. catalog_rows = connection.execute(
  1093. text(
  1094. """
  1095. SELECT artifact_kind, artifact_ref, handoff_status,
  1096. binding_hash
  1097. FROM public.rule_run_artifacts
  1098. WHERE correlation_id = CAST(:correlation_id AS uuid)
  1099. ORDER BY artifact_kind
  1100. """
  1101. ),
  1102. {"correlation_id": correlation_id},
  1103. ).mappings().all()
  1104. # The repeated deterministic output has the same digest and is
  1105. # idempotently retained as one stable catalog handoff.
  1106. assert len(catalog_rows) == 4
  1107. assert all(row["handoff_status"] == "ready" for row in catalog_rows)
  1108. assert all(len(row["binding_hash"]) == 64 for row in catalog_rows)
  1109. assert {
  1110. row["artifact_ref"]
  1111. for row in catalog_rows
  1112. if row["artifact_kind"] == "output"
  1113. } == {
  1114. result["artifact_ref"],
  1115. downstream_result["artifact_ref"],
  1116. }
  1117. assert len(
  1118. list(minio.list_objects(bucket, prefix=prefix, recursive=True))
  1119. ) == 5
  1120. finally:
  1121. for item in list(
  1122. minio.list_objects(bucket, prefix=prefix, recursive=True)
  1123. ):
  1124. minio.remove_object(bucket, item.object_name)
  1125. assert list(
  1126. minio.list_objects(bucket, prefix=prefix, recursive=True)
  1127. ) == []
  1128. with postgres.begin() as connection:
  1129. connection.execute(text(f"DROP TABLE IF EXISTS {customer_table}"))
  1130. with mysql.begin() as connection:
  1131. connection.execute(text(f"DROP TABLE IF EXISTS {segment_table}"))
  1132. with platform.begin() as connection:
  1133. if ledger_jti is not None:
  1134. connection.execute(
  1135. text(
  1136. "DELETE FROM public.runner_task_executions "
  1137. "WHERE token_jti = CAST(:jti AS uuid)"
  1138. ),
  1139. {"jti": ledger_jti},
  1140. )
  1141. if retry_ledger_jti is not None:
  1142. connection.execute(
  1143. text(
  1144. "DELETE FROM public.runner_task_executions "
  1145. "WHERE token_jti = CAST(:jti AS uuid)"
  1146. ),
  1147. {"jti": retry_ledger_jti},
  1148. )
  1149. connection.execute(
  1150. text(
  1151. """
  1152. DELETE FROM public.rule_violation_samples s
  1153. USING public.rule_runs r
  1154. WHERE s.rule_run_id = r.id
  1155. AND r.correlation_id =
  1156. CAST(:correlation_id AS uuid)
  1157. """
  1158. ),
  1159. {"correlation_id": correlation_id},
  1160. )
  1161. connection.execute(
  1162. text(
  1163. "DELETE FROM public.rule_runs "
  1164. "WHERE correlation_id IN ("
  1165. "CAST(:correlation_id AS uuid), "
  1166. "CAST(:failure_correlation_id AS uuid), "
  1167. "CAST(:unknown_correlation_id AS uuid))"
  1168. ),
  1169. {
  1170. "correlation_id": correlation_id,
  1171. "failure_correlation_id": failure_correlation_id,
  1172. "unknown_correlation_id": unknown_correlation_id,
  1173. },
  1174. )
  1175. connection.execute(
  1176. text(
  1177. "DELETE FROM public.rule_run_artifacts "
  1178. "WHERE correlation_id = CAST(:correlation_id AS uuid)"
  1179. ),
  1180. {"correlation_id": correlation_id},
  1181. )
  1182. connection.execute(
  1183. text(
  1184. "DELETE FROM public.rule_execution_plans "
  1185. "WHERE id IN (CAST(:id AS uuid), "
  1186. "CAST(:downstream_id AS uuid))"
  1187. ),
  1188. {
  1189. "id": plan_id,
  1190. "downstream_id": downstream_plan_id,
  1191. },
  1192. )
  1193. connection.execute(
  1194. text(
  1195. "DELETE FROM public.dataflow_dataset_bindings "
  1196. "WHERE dataflow_deployment_id = CAST(:id AS uuid)"
  1197. ),
  1198. {"id": deployment_id},
  1199. )
  1200. connection.execute(
  1201. text(
  1202. "DELETE FROM public.dataflow_component_bindings "
  1203. "WHERE id IN (CAST(:id AS uuid), "
  1204. "CAST(:downstream_id AS uuid))"
  1205. ),
  1206. {
  1207. "id": component_binding_id,
  1208. "downstream_id": downstream_component_binding_id,
  1209. },
  1210. )
  1211. connection.execute(
  1212. text(
  1213. "DELETE FROM public.dataflow_deployments "
  1214. "WHERE id = CAST(:id AS uuid)"
  1215. ),
  1216. {"id": deployment_id},
  1217. )
  1218. connection.execute(
  1219. text(
  1220. "DELETE FROM public.dataflow_versions "
  1221. "WHERE id = CAST(:id AS uuid)"
  1222. ),
  1223. {"id": dataflow_version_id},
  1224. )
  1225. connection.execute(
  1226. text(
  1227. "DELETE FROM public.data_rule_versions "
  1228. "WHERE id IN (CAST(:id AS uuid), "
  1229. "CAST(:downstream_id AS uuid))"
  1230. ),
  1231. {
  1232. "id": rule_id,
  1233. "downstream_id": downstream_rule_id,
  1234. },
  1235. )
  1236. connection.execute(
  1237. text(
  1238. "DELETE FROM public.data_rules "
  1239. "WHERE rule_uid IN (CAST(:rule_uid AS uuid), "
  1240. "CAST(:downstream_rule_uid AS uuid))"
  1241. ),
  1242. {
  1243. "rule_uid": rule_uid,
  1244. "downstream_rule_uid": downstream_rule_uid,
  1245. },
  1246. )
  1247. for schema in (input_schema, lookup_schema, output_schema):
  1248. connection.execute(
  1249. text(
  1250. "DELETE FROM public.data_schema_snapshots "
  1251. "WHERE id = CAST(:id AS uuid)"
  1252. ),
  1253. {"id": schema["id"]},
  1254. )
  1255. postgres.dispose()
  1256. mysql.dispose()
  1257. platform.dispose()