relationalai 0.13.0.dev0__py3-none-any.whl → 0.13.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (838) hide show
  1. frontend/debugger/dist/.gitignore +2 -0
  2. frontend/debugger/dist/assets/favicon-Dy0ZgA6N.png +0 -0
  3. frontend/debugger/dist/assets/index-Cssla-O7.js +208 -0
  4. frontend/debugger/dist/assets/index-DlHsYx1V.css +9 -0
  5. frontend/debugger/dist/index.html +17 -0
  6. relationalai/__init__.py +256 -1
  7. relationalai/clients/__init__.py +18 -0
  8. relationalai/clients/client.py +947 -0
  9. relationalai/clients/config.py +673 -0
  10. relationalai/clients/direct_access_client.py +118 -0
  11. relationalai/clients/exec_txn_poller.py +91 -0
  12. relationalai/clients/hash_util.py +31 -0
  13. relationalai/clients/local.py +586 -0
  14. relationalai/clients/profile_polling.py +73 -0
  15. relationalai/clients/resources/__init__.py +8 -0
  16. relationalai/clients/resources/azure/azure.py +502 -0
  17. relationalai/clients/resources/snowflake/__init__.py +20 -0
  18. relationalai/clients/resources/snowflake/cli_resources.py +98 -0
  19. relationalai/clients/resources/snowflake/direct_access_resources.py +734 -0
  20. relationalai/clients/resources/snowflake/engine_service.py +381 -0
  21. relationalai/clients/resources/snowflake/engine_state_handlers.py +315 -0
  22. relationalai/clients/resources/snowflake/error_handlers.py +240 -0
  23. relationalai/clients/resources/snowflake/export_procedure.py.jinja +249 -0
  24. relationalai/clients/resources/snowflake/resources_factory.py +99 -0
  25. relationalai/clients/resources/snowflake/snowflake.py +3185 -0
  26. relationalai/clients/resources/snowflake/use_index_poller.py +1019 -0
  27. relationalai/clients/resources/snowflake/use_index_resources.py +188 -0
  28. relationalai/clients/resources/snowflake/util.py +387 -0
  29. relationalai/clients/result_helpers.py +420 -0
  30. relationalai/clients/types.py +118 -0
  31. relationalai/clients/util.py +356 -0
  32. relationalai/debugging.py +389 -0
  33. relationalai/dsl.py +1749 -0
  34. relationalai/early_access/builder/__init__.py +30 -0
  35. relationalai/early_access/builder/builder/__init__.py +35 -0
  36. relationalai/early_access/builder/snowflake/__init__.py +12 -0
  37. relationalai/early_access/builder/std/__init__.py +25 -0
  38. relationalai/early_access/builder/std/decimals/__init__.py +12 -0
  39. relationalai/early_access/builder/std/integers/__init__.py +12 -0
  40. relationalai/early_access/builder/std/math/__init__.py +12 -0
  41. relationalai/early_access/builder/std/strings/__init__.py +14 -0
  42. relationalai/early_access/devtools/__init__.py +12 -0
  43. relationalai/early_access/devtools/benchmark_lqp/__init__.py +12 -0
  44. relationalai/early_access/devtools/extract_lqp/__init__.py +12 -0
  45. relationalai/early_access/dsl/adapters/orm/adapter_qb.py +427 -0
  46. relationalai/early_access/dsl/adapters/orm/parser.py +636 -0
  47. relationalai/early_access/dsl/adapters/owl/adapter.py +176 -0
  48. relationalai/early_access/dsl/adapters/owl/parser.py +160 -0
  49. relationalai/early_access/dsl/bindings/common.py +402 -0
  50. relationalai/early_access/dsl/bindings/csv.py +170 -0
  51. relationalai/early_access/dsl/bindings/legacy/binding_models.py +143 -0
  52. relationalai/early_access/dsl/bindings/snowflake.py +64 -0
  53. relationalai/early_access/dsl/codegen/binder.py +411 -0
  54. relationalai/early_access/dsl/codegen/common.py +79 -0
  55. relationalai/early_access/dsl/codegen/helpers.py +23 -0
  56. relationalai/early_access/dsl/codegen/relations.py +700 -0
  57. relationalai/early_access/dsl/codegen/weaver.py +417 -0
  58. relationalai/early_access/dsl/core/builders/__init__.py +47 -0
  59. relationalai/early_access/dsl/core/builders/logic.py +19 -0
  60. relationalai/early_access/dsl/core/builders/scalar_constraint.py +11 -0
  61. relationalai/early_access/dsl/core/constraints/predicate/atomic.py +455 -0
  62. relationalai/early_access/dsl/core/constraints/predicate/universal.py +73 -0
  63. relationalai/early_access/dsl/core/constraints/scalar.py +310 -0
  64. relationalai/early_access/dsl/core/context.py +13 -0
  65. relationalai/early_access/dsl/core/cset.py +132 -0
  66. relationalai/early_access/dsl/core/exprs/__init__.py +116 -0
  67. relationalai/early_access/dsl/core/exprs/relational.py +18 -0
  68. relationalai/early_access/dsl/core/exprs/scalar.py +412 -0
  69. relationalai/early_access/dsl/core/instances.py +44 -0
  70. relationalai/early_access/dsl/core/logic/__init__.py +193 -0
  71. relationalai/early_access/dsl/core/logic/aggregation.py +98 -0
  72. relationalai/early_access/dsl/core/logic/exists.py +223 -0
  73. relationalai/early_access/dsl/core/logic/helper.py +163 -0
  74. relationalai/early_access/dsl/core/namespaces.py +32 -0
  75. relationalai/early_access/dsl/core/relations.py +276 -0
  76. relationalai/early_access/dsl/core/rules.py +112 -0
  77. relationalai/early_access/dsl/core/std/__init__.py +45 -0
  78. relationalai/early_access/dsl/core/temporal/recall.py +6 -0
  79. relationalai/early_access/dsl/core/types/__init__.py +270 -0
  80. relationalai/early_access/dsl/core/types/concepts.py +128 -0
  81. relationalai/early_access/dsl/core/types/constrained/__init__.py +267 -0
  82. relationalai/early_access/dsl/core/types/constrained/nominal.py +143 -0
  83. relationalai/early_access/dsl/core/types/constrained/subtype.py +124 -0
  84. relationalai/early_access/dsl/core/types/standard.py +92 -0
  85. relationalai/early_access/dsl/core/types/unconstrained.py +50 -0
  86. relationalai/early_access/dsl/core/types/variables.py +203 -0
  87. relationalai/early_access/dsl/ir/compiler.py +318 -0
  88. relationalai/early_access/dsl/ir/executor.py +260 -0
  89. relationalai/early_access/dsl/ontologies/constraints.py +88 -0
  90. relationalai/early_access/dsl/ontologies/export.py +30 -0
  91. relationalai/early_access/dsl/ontologies/models.py +453 -0
  92. relationalai/early_access/dsl/ontologies/python_printer.py +303 -0
  93. relationalai/early_access/dsl/ontologies/readings.py +60 -0
  94. relationalai/early_access/dsl/ontologies/relationships.py +322 -0
  95. relationalai/early_access/dsl/ontologies/roles.py +87 -0
  96. relationalai/early_access/dsl/ontologies/subtyping.py +55 -0
  97. relationalai/early_access/dsl/orm/constraints.py +438 -0
  98. relationalai/early_access/dsl/orm/measures/dimensions.py +200 -0
  99. relationalai/early_access/dsl/orm/measures/initializer.py +16 -0
  100. relationalai/early_access/dsl/orm/measures/measure_rules.py +275 -0
  101. relationalai/early_access/dsl/orm/measures/measures.py +299 -0
  102. relationalai/early_access/dsl/orm/measures/role_exprs.py +268 -0
  103. relationalai/early_access/dsl/orm/models.py +256 -0
  104. relationalai/early_access/dsl/orm/object_oriented_printer.py +344 -0
  105. relationalai/early_access/dsl/orm/printer.py +469 -0
  106. relationalai/early_access/dsl/orm/reasoners.py +480 -0
  107. relationalai/early_access/dsl/orm/relations.py +19 -0
  108. relationalai/early_access/dsl/orm/relationships.py +251 -0
  109. relationalai/early_access/dsl/orm/types.py +42 -0
  110. relationalai/early_access/dsl/orm/utils.py +79 -0
  111. relationalai/early_access/dsl/orm/verb.py +204 -0
  112. relationalai/early_access/dsl/physical_metadata/tables.py +133 -0
  113. relationalai/early_access/dsl/relations.py +170 -0
  114. relationalai/early_access/dsl/rulesets.py +69 -0
  115. relationalai/early_access/dsl/schemas/__init__.py +450 -0
  116. relationalai/early_access/dsl/schemas/builder.py +48 -0
  117. relationalai/early_access/dsl/schemas/comp_names.py +51 -0
  118. relationalai/early_access/dsl/schemas/components.py +203 -0
  119. relationalai/early_access/dsl/schemas/contexts.py +156 -0
  120. relationalai/early_access/dsl/schemas/exprs.py +89 -0
  121. relationalai/early_access/dsl/schemas/fragments.py +464 -0
  122. relationalai/early_access/dsl/serialization.py +79 -0
  123. relationalai/early_access/dsl/serialize/exporter.py +163 -0
  124. relationalai/early_access/dsl/snow/api.py +105 -0
  125. relationalai/early_access/dsl/snow/common.py +76 -0
  126. relationalai/early_access/dsl/state_mgmt/__init__.py +129 -0
  127. relationalai/early_access/dsl/state_mgmt/state_charts.py +125 -0
  128. relationalai/early_access/dsl/state_mgmt/transitions.py +130 -0
  129. relationalai/early_access/dsl/types/__init__.py +40 -0
  130. relationalai/early_access/dsl/types/concepts.py +12 -0
  131. relationalai/early_access/dsl/types/entities.py +135 -0
  132. relationalai/early_access/dsl/types/values.py +17 -0
  133. relationalai/early_access/dsl/utils.py +102 -0
  134. relationalai/early_access/graphs/__init__.py +13 -0
  135. relationalai/early_access/lqp/__init__.py +12 -0
  136. relationalai/early_access/lqp/compiler/__init__.py +12 -0
  137. relationalai/early_access/lqp/constructors/__init__.py +18 -0
  138. relationalai/early_access/lqp/executor/__init__.py +12 -0
  139. relationalai/early_access/lqp/ir/__init__.py +12 -0
  140. relationalai/early_access/lqp/passes/__init__.py +12 -0
  141. relationalai/early_access/lqp/pragmas/__init__.py +12 -0
  142. relationalai/early_access/lqp/primitives/__init__.py +12 -0
  143. relationalai/early_access/lqp/types/__init__.py +12 -0
  144. relationalai/early_access/lqp/utils/__init__.py +12 -0
  145. relationalai/early_access/lqp/validators/__init__.py +12 -0
  146. relationalai/early_access/metamodel/__init__.py +58 -0
  147. relationalai/early_access/metamodel/builtins/__init__.py +12 -0
  148. relationalai/early_access/metamodel/compiler/__init__.py +12 -0
  149. relationalai/early_access/metamodel/dependency/__init__.py +12 -0
  150. relationalai/early_access/metamodel/factory/__init__.py +17 -0
  151. relationalai/early_access/metamodel/helpers/__init__.py +12 -0
  152. relationalai/early_access/metamodel/ir/__init__.py +14 -0
  153. relationalai/early_access/metamodel/rewrite/__init__.py +7 -0
  154. relationalai/early_access/metamodel/typer/__init__.py +3 -0
  155. relationalai/early_access/metamodel/typer/typer/__init__.py +12 -0
  156. relationalai/early_access/metamodel/types/__init__.py +15 -0
  157. relationalai/early_access/metamodel/util/__init__.py +15 -0
  158. relationalai/early_access/metamodel/visitor/__init__.py +12 -0
  159. relationalai/early_access/rel/__init__.py +12 -0
  160. relationalai/early_access/rel/executor/__init__.py +12 -0
  161. relationalai/early_access/rel/rel_utils/__init__.py +12 -0
  162. relationalai/early_access/rel/rewrite/__init__.py +7 -0
  163. relationalai/early_access/solvers/__init__.py +19 -0
  164. relationalai/early_access/sql/__init__.py +11 -0
  165. relationalai/early_access/sql/executor/__init__.py +3 -0
  166. relationalai/early_access/sql/rewrite/__init__.py +3 -0
  167. relationalai/early_access/tests/logging/__init__.py +12 -0
  168. relationalai/early_access/tests/test_snapshot_base/__init__.py +12 -0
  169. relationalai/early_access/tests/utils/__init__.py +12 -0
  170. relationalai/environments/__init__.py +35 -0
  171. relationalai/environments/base.py +381 -0
  172. relationalai/environments/colab.py +14 -0
  173. relationalai/environments/generic.py +71 -0
  174. relationalai/environments/ipython.py +68 -0
  175. relationalai/environments/jupyter.py +9 -0
  176. relationalai/environments/snowbook.py +169 -0
  177. relationalai/errors.py +2496 -0
  178. relationalai/experimental/SF.py +38 -0
  179. relationalai/experimental/inspect.py +47 -0
  180. relationalai/experimental/pathfinder/__init__.py +158 -0
  181. relationalai/experimental/pathfinder/api.py +160 -0
  182. relationalai/experimental/pathfinder/automaton.py +584 -0
  183. relationalai/experimental/pathfinder/bridge.py +226 -0
  184. relationalai/experimental/pathfinder/compiler.py +416 -0
  185. relationalai/experimental/pathfinder/datalog.py +214 -0
  186. relationalai/experimental/pathfinder/diagnostics.py +56 -0
  187. relationalai/experimental/pathfinder/filter.py +236 -0
  188. relationalai/experimental/pathfinder/glushkov.py +439 -0
  189. relationalai/experimental/pathfinder/options.py +265 -0
  190. relationalai/experimental/pathfinder/pathfinder-v0.7.0.rel +1951 -0
  191. relationalai/experimental/pathfinder/rpq.py +344 -0
  192. relationalai/experimental/pathfinder/transition.py +200 -0
  193. relationalai/experimental/pathfinder/utils.py +26 -0
  194. relationalai/experimental/paths/README.md +107 -0
  195. relationalai/experimental/paths/api.py +143 -0
  196. relationalai/experimental/paths/benchmarks/grid_graph.py +37 -0
  197. relationalai/experimental/paths/code_organization.md +2 -0
  198. relationalai/experimental/paths/examples/Movies.ipynb +16328 -0
  199. relationalai/experimental/paths/examples/basic_example.py +40 -0
  200. relationalai/experimental/paths/examples/minimal_engine_warmup.py +3 -0
  201. relationalai/experimental/paths/examples/movie_example.py +77 -0
  202. relationalai/experimental/paths/examples/movies_data/actedin.csv +193 -0
  203. relationalai/experimental/paths/examples/movies_data/directed.csv +45 -0
  204. relationalai/experimental/paths/examples/movies_data/follows.csv +7 -0
  205. relationalai/experimental/paths/examples/movies_data/movies.csv +39 -0
  206. relationalai/experimental/paths/examples/movies_data/person.csv +134 -0
  207. relationalai/experimental/paths/examples/movies_data/produced.csv +16 -0
  208. relationalai/experimental/paths/examples/movies_data/ratings.csv +10 -0
  209. relationalai/experimental/paths/examples/movies_data/wrote.csv +11 -0
  210. relationalai/experimental/paths/examples/paths_benchmark.py +115 -0
  211. relationalai/experimental/paths/examples/paths_example.py +116 -0
  212. relationalai/experimental/paths/examples/pattern_to_automaton.py +28 -0
  213. relationalai/experimental/paths/find_paths_via_automaton.py +85 -0
  214. relationalai/experimental/paths/graph.py +185 -0
  215. relationalai/experimental/paths/path_algorithms/find_paths.py +280 -0
  216. relationalai/experimental/paths/path_algorithms/one_sided_ball_repetition.py +26 -0
  217. relationalai/experimental/paths/path_algorithms/one_sided_ball_upto.py +111 -0
  218. relationalai/experimental/paths/path_algorithms/single.py +59 -0
  219. relationalai/experimental/paths/path_algorithms/two_sided_balls_repetition.py +39 -0
  220. relationalai/experimental/paths/path_algorithms/two_sided_balls_upto.py +103 -0
  221. relationalai/experimental/paths/path_algorithms/usp-old.py +130 -0
  222. relationalai/experimental/paths/path_algorithms/usp-tuple.py +183 -0
  223. relationalai/experimental/paths/path_algorithms/usp.py +150 -0
  224. relationalai/experimental/paths/product_graph.py +93 -0
  225. relationalai/experimental/paths/rpq/automaton.py +584 -0
  226. relationalai/experimental/paths/rpq/diagnostics.py +56 -0
  227. relationalai/experimental/paths/rpq/rpq.py +378 -0
  228. relationalai/experimental/paths/tests/tests_limit_sp_max_length.py +90 -0
  229. relationalai/experimental/paths/tests/tests_limit_sp_multiple.py +119 -0
  230. relationalai/experimental/paths/tests/tests_limit_sp_single.py +104 -0
  231. relationalai/experimental/paths/tests/tests_limit_walks_multiple.py +113 -0
  232. relationalai/experimental/paths/tests/tests_limit_walks_single.py +149 -0
  233. relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_multiple.py +70 -0
  234. relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_single.py +64 -0
  235. relationalai/experimental/paths/tests/tests_one_sided_ball_upto_multiple.py +115 -0
  236. relationalai/experimental/paths/tests/tests_one_sided_ball_upto_single.py +75 -0
  237. relationalai/experimental/paths/tests/tests_single_paths.py +152 -0
  238. relationalai/experimental/paths/tests/tests_single_walks.py +208 -0
  239. relationalai/experimental/paths/tests/tests_single_walks_undirected.py +297 -0
  240. relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_multiple.py +107 -0
  241. relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_single.py +76 -0
  242. relationalai/experimental/paths/tests/tests_two_sided_balls_upto_multiple.py +76 -0
  243. relationalai/experimental/paths/tests/tests_two_sided_balls_upto_single.py +110 -0
  244. relationalai/experimental/paths/tests/tests_usp_nsp_multiple.py +229 -0
  245. relationalai/experimental/paths/tests/tests_usp_nsp_single.py +108 -0
  246. relationalai/experimental/paths/tree_agg.py +168 -0
  247. relationalai/experimental/paths/utilities/iterators.py +27 -0
  248. relationalai/experimental/paths/utilities/prefix_sum.py +91 -0
  249. relationalai/experimental/solvers.py +1087 -0
  250. relationalai/loaders/csv.py +195 -0
  251. relationalai/loaders/loader.py +177 -0
  252. relationalai/loaders/types.py +23 -0
  253. relationalai/rel_emitter.py +373 -0
  254. relationalai/rel_utils.py +185 -0
  255. relationalai/semantics/__init__.py +22 -146
  256. relationalai/semantics/designs/query_builder/identify_by.md +106 -0
  257. relationalai/semantics/devtools/benchmark_lqp.py +535 -0
  258. relationalai/semantics/devtools/compilation_manager.py +294 -0
  259. relationalai/semantics/devtools/extract_lqp.py +110 -0
  260. relationalai/semantics/internal/internal.py +3785 -0
  261. relationalai/semantics/internal/snowflake.py +325 -0
  262. relationalai/semantics/lqp/README.md +34 -0
  263. relationalai/semantics/lqp/builtins.py +16 -0
  264. relationalai/semantics/lqp/compiler.py +22 -0
  265. relationalai/semantics/lqp/constructors.py +68 -0
  266. relationalai/semantics/lqp/executor.py +469 -0
  267. relationalai/semantics/lqp/intrinsics.py +24 -0
  268. relationalai/semantics/lqp/model2lqp.py +877 -0
  269. relationalai/semantics/lqp/passes.py +680 -0
  270. relationalai/semantics/lqp/primitives.py +252 -0
  271. relationalai/semantics/lqp/result_helpers.py +202 -0
  272. relationalai/semantics/lqp/rewrite/annotate_constraints.py +57 -0
  273. relationalai/semantics/lqp/rewrite/cdc.py +216 -0
  274. relationalai/semantics/lqp/rewrite/extract_common.py +338 -0
  275. relationalai/semantics/lqp/rewrite/extract_keys.py +512 -0
  276. relationalai/semantics/lqp/rewrite/function_annotations.py +114 -0
  277. relationalai/semantics/lqp/rewrite/functional_dependencies.py +314 -0
  278. relationalai/semantics/lqp/rewrite/quantify_vars.py +296 -0
  279. relationalai/semantics/lqp/rewrite/splinter.py +76 -0
  280. relationalai/semantics/lqp/types.py +101 -0
  281. relationalai/semantics/lqp/utils.py +160 -0
  282. relationalai/semantics/lqp/validators.py +57 -0
  283. relationalai/semantics/metamodel/__init__.py +40 -6
  284. relationalai/semantics/metamodel/builtins.py +771 -205
  285. relationalai/semantics/metamodel/compiler.py +133 -0
  286. relationalai/semantics/metamodel/dependency.py +862 -0
  287. relationalai/semantics/metamodel/executor.py +61 -0
  288. relationalai/semantics/metamodel/factory.py +287 -0
  289. relationalai/semantics/metamodel/helpers.py +361 -0
  290. relationalai/semantics/metamodel/rewrite/discharge_constraints.py +39 -0
  291. relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py +210 -0
  292. relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py +78 -0
  293. relationalai/semantics/metamodel/rewrite/flatten.py +554 -0
  294. relationalai/semantics/metamodel/rewrite/format_outputs.py +165 -0
  295. relationalai/semantics/metamodel/typer/checker.py +353 -0
  296. relationalai/semantics/metamodel/typer/typer.py +1399 -0
  297. relationalai/semantics/metamodel/util.py +506 -0
  298. relationalai/semantics/reasoners/__init__.py +10 -0
  299. relationalai/semantics/reasoners/graph/README.md +620 -0
  300. relationalai/semantics/reasoners/graph/__init__.py +37 -0
  301. relationalai/semantics/reasoners/graph/core.py +9019 -0
  302. relationalai/semantics/reasoners/graph/design/beyond_demand_transform.md +797 -0
  303. relationalai/semantics/reasoners/graph/tests/README.md +21 -0
  304. relationalai/semantics/reasoners/optimization/__init__.py +68 -0
  305. relationalai/semantics/reasoners/optimization/common.py +88 -0
  306. relationalai/semantics/reasoners/optimization/solvers_dev.py +568 -0
  307. relationalai/semantics/reasoners/optimization/solvers_pb.py +1414 -0
  308. relationalai/semantics/rel/builtins.py +40 -0
  309. relationalai/semantics/rel/compiler.py +989 -0
  310. relationalai/semantics/rel/executor.py +362 -0
  311. relationalai/semantics/rel/rel.py +482 -0
  312. relationalai/semantics/rel/rel_utils.py +276 -0
  313. relationalai/semantics/snowflake/__init__.py +3 -0
  314. relationalai/semantics/sql/compiler.py +2503 -0
  315. relationalai/semantics/sql/executor/duck_db.py +52 -0
  316. relationalai/semantics/sql/executor/result_helpers.py +64 -0
  317. relationalai/semantics/sql/executor/snowflake.py +149 -0
  318. relationalai/semantics/sql/rewrite/denormalize.py +222 -0
  319. relationalai/semantics/sql/rewrite/double_negation.py +49 -0
  320. relationalai/semantics/sql/rewrite/recursive_union.py +127 -0
  321. relationalai/semantics/sql/rewrite/sort_output_query.py +246 -0
  322. relationalai/semantics/sql/sql.py +504 -0
  323. relationalai/semantics/std/__init__.py +40 -60
  324. relationalai/semantics/std/constraints.py +43 -37
  325. relationalai/semantics/std/datetime.py +135 -246
  326. relationalai/semantics/std/decimals.py +52 -45
  327. relationalai/semantics/std/floats.py +5 -13
  328. relationalai/semantics/std/integers.py +11 -26
  329. relationalai/semantics/std/math.py +112 -183
  330. relationalai/semantics/std/pragmas.py +11 -0
  331. relationalai/semantics/std/re.py +62 -80
  332. relationalai/semantics/std/std.py +14 -0
  333. relationalai/semantics/std/strings.py +60 -117
  334. relationalai/semantics/tests/test_snapshot_abstract.py +143 -0
  335. relationalai/semantics/tests/test_snapshot_base.py +9 -0
  336. relationalai/semantics/tests/utils.py +46 -0
  337. relationalai/std/__init__.py +70 -0
  338. relationalai/tools/cli.py +2089 -0
  339. relationalai/tools/cli_controls.py +1826 -0
  340. relationalai/tools/cli_helpers.py +802 -0
  341. relationalai/tools/debugger.py +183 -289
  342. relationalai/tools/debugger_client.py +109 -0
  343. relationalai/tools/debugger_server.py +302 -0
  344. relationalai/tools/dev.py +685 -0
  345. relationalai/tools/notes +7 -0
  346. relationalai/tools/qb_debugger.py +425 -0
  347. relationalai/util/clean_up_databases.py +95 -0
  348. relationalai/util/format.py +106 -48
  349. relationalai/util/list_databases.py +9 -0
  350. relationalai/util/otel_configuration.py +26 -0
  351. relationalai/util/otel_handler.py +484 -0
  352. relationalai/util/snowflake_handler.py +88 -0
  353. relationalai/util/span_format_test.py +43 -0
  354. relationalai/util/span_tracker.py +207 -0
  355. relationalai/util/spans_file_handler.py +72 -0
  356. relationalai/util/tracing_handler.py +34 -0
  357. relationalai-0.13.2.dist-info/METADATA +74 -0
  358. relationalai-0.13.2.dist-info/RECORD +460 -0
  359. relationalai-0.13.2.dist-info/WHEEL +4 -0
  360. relationalai-0.13.2.dist-info/entry_points.txt +3 -0
  361. relationalai-0.13.2.dist-info/licenses/LICENSE +202 -0
  362. relationalai_test_util/__init__.py +4 -0
  363. relationalai_test_util/fixtures.py +233 -0
  364. relationalai_test_util/snapshot.py +252 -0
  365. relationalai_test_util/traceback.py +118 -0
  366. relationalai/config/__init__.py +0 -56
  367. relationalai/config/config.py +0 -289
  368. relationalai/config/config_fields.py +0 -86
  369. relationalai/config/connections/__init__.py +0 -46
  370. relationalai/config/connections/base.py +0 -23
  371. relationalai/config/connections/duckdb.py +0 -29
  372. relationalai/config/connections/snowflake.py +0 -243
  373. relationalai/config/external/__init__.py +0 -17
  374. relationalai/config/external/dbt_converter.py +0 -101
  375. relationalai/config/external/dbt_models.py +0 -93
  376. relationalai/config/external/snowflake_converter.py +0 -41
  377. relationalai/config/external/snowflake_models.py +0 -85
  378. relationalai/config/external/utils.py +0 -19
  379. relationalai/semantics/backends/lqp/annotations.py +0 -11
  380. relationalai/semantics/backends/sql/sql_compiler.py +0 -327
  381. relationalai/semantics/frontend/base.py +0 -1707
  382. relationalai/semantics/frontend/core.py +0 -179
  383. relationalai/semantics/frontend/front_compiler.py +0 -1313
  384. relationalai/semantics/frontend/pprint.py +0 -408
  385. relationalai/semantics/metamodel/metamodel.py +0 -437
  386. relationalai/semantics/metamodel/metamodel_analyzer.py +0 -519
  387. relationalai/semantics/metamodel/metamodel_compiler.py +0 -0
  388. relationalai/semantics/metamodel/pprint.py +0 -412
  389. relationalai/semantics/metamodel/rewriter.py +0 -266
  390. relationalai/semantics/metamodel/typer.py +0 -1378
  391. relationalai/semantics/std/aggregates.py +0 -149
  392. relationalai/semantics/std/common.py +0 -44
  393. relationalai/semantics/std/numbers.py +0 -86
  394. relationalai/shims/executor.py +0 -147
  395. relationalai/shims/helpers.py +0 -126
  396. relationalai/shims/hoister.py +0 -221
  397. relationalai/shims/mm2v0.py +0 -1290
  398. relationalai/tools/cli/__init__.py +0 -6
  399. relationalai/tools/cli/cli.py +0 -90
  400. relationalai/tools/cli/components/__init__.py +0 -5
  401. relationalai/tools/cli/components/progress_reader.py +0 -1524
  402. relationalai/tools/cli/components/utils.py +0 -58
  403. relationalai/tools/cli/config_template.py +0 -45
  404. relationalai/tools/cli/dev.py +0 -19
  405. relationalai/tools/typer_debugger.py +0 -93
  406. relationalai/util/dataclasses.py +0 -43
  407. relationalai/util/docutils.py +0 -40
  408. relationalai/util/error.py +0 -199
  409. relationalai/util/naming.py +0 -145
  410. relationalai/util/python.py +0 -35
  411. relationalai/util/runtime.py +0 -156
  412. relationalai/util/schema.py +0 -197
  413. relationalai/util/source.py +0 -185
  414. relationalai/util/structures.py +0 -163
  415. relationalai/util/tracing.py +0 -261
  416. relationalai-0.13.0.dev0.dist-info/METADATA +0 -46
  417. relationalai-0.13.0.dev0.dist-info/RECORD +0 -488
  418. relationalai-0.13.0.dev0.dist-info/WHEEL +0 -5
  419. relationalai-0.13.0.dev0.dist-info/entry_points.txt +0 -3
  420. relationalai-0.13.0.dev0.dist-info/top_level.txt +0 -2
  421. v0/relationalai/__init__.py +0 -216
  422. v0/relationalai/clients/__init__.py +0 -5
  423. v0/relationalai/clients/azure.py +0 -477
  424. v0/relationalai/clients/client.py +0 -912
  425. v0/relationalai/clients/config.py +0 -673
  426. v0/relationalai/clients/direct_access_client.py +0 -118
  427. v0/relationalai/clients/hash_util.py +0 -31
  428. v0/relationalai/clients/local.py +0 -571
  429. v0/relationalai/clients/profile_polling.py +0 -73
  430. v0/relationalai/clients/result_helpers.py +0 -420
  431. v0/relationalai/clients/snowflake.py +0 -3869
  432. v0/relationalai/clients/types.py +0 -113
  433. v0/relationalai/clients/use_index_poller.py +0 -980
  434. v0/relationalai/clients/util.py +0 -356
  435. v0/relationalai/debugging.py +0 -389
  436. v0/relationalai/dsl.py +0 -1749
  437. v0/relationalai/early_access/builder/__init__.py +0 -30
  438. v0/relationalai/early_access/builder/builder/__init__.py +0 -35
  439. v0/relationalai/early_access/builder/snowflake/__init__.py +0 -12
  440. v0/relationalai/early_access/builder/std/__init__.py +0 -25
  441. v0/relationalai/early_access/builder/std/decimals/__init__.py +0 -12
  442. v0/relationalai/early_access/builder/std/integers/__init__.py +0 -12
  443. v0/relationalai/early_access/builder/std/math/__init__.py +0 -12
  444. v0/relationalai/early_access/builder/std/strings/__init__.py +0 -14
  445. v0/relationalai/early_access/devtools/__init__.py +0 -12
  446. v0/relationalai/early_access/devtools/benchmark_lqp/__init__.py +0 -12
  447. v0/relationalai/early_access/devtools/extract_lqp/__init__.py +0 -12
  448. v0/relationalai/early_access/dsl/adapters/orm/adapter_qb.py +0 -427
  449. v0/relationalai/early_access/dsl/adapters/orm/parser.py +0 -636
  450. v0/relationalai/early_access/dsl/adapters/owl/adapter.py +0 -176
  451. v0/relationalai/early_access/dsl/adapters/owl/parser.py +0 -160
  452. v0/relationalai/early_access/dsl/bindings/common.py +0 -402
  453. v0/relationalai/early_access/dsl/bindings/csv.py +0 -170
  454. v0/relationalai/early_access/dsl/bindings/legacy/binding_models.py +0 -143
  455. v0/relationalai/early_access/dsl/bindings/snowflake.py +0 -64
  456. v0/relationalai/early_access/dsl/codegen/binder.py +0 -411
  457. v0/relationalai/early_access/dsl/codegen/common.py +0 -79
  458. v0/relationalai/early_access/dsl/codegen/helpers.py +0 -23
  459. v0/relationalai/early_access/dsl/codegen/relations.py +0 -700
  460. v0/relationalai/early_access/dsl/codegen/weaver.py +0 -417
  461. v0/relationalai/early_access/dsl/core/builders/__init__.py +0 -47
  462. v0/relationalai/early_access/dsl/core/builders/logic.py +0 -19
  463. v0/relationalai/early_access/dsl/core/builders/scalar_constraint.py +0 -11
  464. v0/relationalai/early_access/dsl/core/constraints/predicate/atomic.py +0 -455
  465. v0/relationalai/early_access/dsl/core/constraints/predicate/universal.py +0 -73
  466. v0/relationalai/early_access/dsl/core/constraints/scalar.py +0 -310
  467. v0/relationalai/early_access/dsl/core/context.py +0 -13
  468. v0/relationalai/early_access/dsl/core/cset.py +0 -132
  469. v0/relationalai/early_access/dsl/core/exprs/__init__.py +0 -116
  470. v0/relationalai/early_access/dsl/core/exprs/relational.py +0 -18
  471. v0/relationalai/early_access/dsl/core/exprs/scalar.py +0 -412
  472. v0/relationalai/early_access/dsl/core/instances.py +0 -44
  473. v0/relationalai/early_access/dsl/core/logic/__init__.py +0 -193
  474. v0/relationalai/early_access/dsl/core/logic/aggregation.py +0 -98
  475. v0/relationalai/early_access/dsl/core/logic/exists.py +0 -223
  476. v0/relationalai/early_access/dsl/core/logic/helper.py +0 -163
  477. v0/relationalai/early_access/dsl/core/namespaces.py +0 -32
  478. v0/relationalai/early_access/dsl/core/relations.py +0 -276
  479. v0/relationalai/early_access/dsl/core/rules.py +0 -112
  480. v0/relationalai/early_access/dsl/core/std/__init__.py +0 -45
  481. v0/relationalai/early_access/dsl/core/temporal/recall.py +0 -6
  482. v0/relationalai/early_access/dsl/core/types/__init__.py +0 -270
  483. v0/relationalai/early_access/dsl/core/types/concepts.py +0 -128
  484. v0/relationalai/early_access/dsl/core/types/constrained/__init__.py +0 -267
  485. v0/relationalai/early_access/dsl/core/types/constrained/nominal.py +0 -143
  486. v0/relationalai/early_access/dsl/core/types/constrained/subtype.py +0 -124
  487. v0/relationalai/early_access/dsl/core/types/standard.py +0 -92
  488. v0/relationalai/early_access/dsl/core/types/unconstrained.py +0 -50
  489. v0/relationalai/early_access/dsl/core/types/variables.py +0 -203
  490. v0/relationalai/early_access/dsl/ir/compiler.py +0 -318
  491. v0/relationalai/early_access/dsl/ir/executor.py +0 -260
  492. v0/relationalai/early_access/dsl/ontologies/constraints.py +0 -88
  493. v0/relationalai/early_access/dsl/ontologies/export.py +0 -30
  494. v0/relationalai/early_access/dsl/ontologies/models.py +0 -453
  495. v0/relationalai/early_access/dsl/ontologies/python_printer.py +0 -303
  496. v0/relationalai/early_access/dsl/ontologies/readings.py +0 -60
  497. v0/relationalai/early_access/dsl/ontologies/relationships.py +0 -322
  498. v0/relationalai/early_access/dsl/ontologies/roles.py +0 -87
  499. v0/relationalai/early_access/dsl/ontologies/subtyping.py +0 -55
  500. v0/relationalai/early_access/dsl/orm/constraints.py +0 -438
  501. v0/relationalai/early_access/dsl/orm/measures/dimensions.py +0 -200
  502. v0/relationalai/early_access/dsl/orm/measures/initializer.py +0 -16
  503. v0/relationalai/early_access/dsl/orm/measures/measure_rules.py +0 -275
  504. v0/relationalai/early_access/dsl/orm/measures/measures.py +0 -299
  505. v0/relationalai/early_access/dsl/orm/measures/role_exprs.py +0 -268
  506. v0/relationalai/early_access/dsl/orm/models.py +0 -256
  507. v0/relationalai/early_access/dsl/orm/object_oriented_printer.py +0 -344
  508. v0/relationalai/early_access/dsl/orm/printer.py +0 -469
  509. v0/relationalai/early_access/dsl/orm/reasoners.py +0 -480
  510. v0/relationalai/early_access/dsl/orm/relations.py +0 -19
  511. v0/relationalai/early_access/dsl/orm/relationships.py +0 -251
  512. v0/relationalai/early_access/dsl/orm/types.py +0 -42
  513. v0/relationalai/early_access/dsl/orm/utils.py +0 -79
  514. v0/relationalai/early_access/dsl/orm/verb.py +0 -204
  515. v0/relationalai/early_access/dsl/physical_metadata/tables.py +0 -133
  516. v0/relationalai/early_access/dsl/relations.py +0 -170
  517. v0/relationalai/early_access/dsl/rulesets.py +0 -69
  518. v0/relationalai/early_access/dsl/schemas/__init__.py +0 -450
  519. v0/relationalai/early_access/dsl/schemas/builder.py +0 -48
  520. v0/relationalai/early_access/dsl/schemas/comp_names.py +0 -51
  521. v0/relationalai/early_access/dsl/schemas/components.py +0 -203
  522. v0/relationalai/early_access/dsl/schemas/contexts.py +0 -156
  523. v0/relationalai/early_access/dsl/schemas/exprs.py +0 -89
  524. v0/relationalai/early_access/dsl/schemas/fragments.py +0 -464
  525. v0/relationalai/early_access/dsl/serialization.py +0 -79
  526. v0/relationalai/early_access/dsl/serialize/exporter.py +0 -163
  527. v0/relationalai/early_access/dsl/snow/api.py +0 -104
  528. v0/relationalai/early_access/dsl/snow/common.py +0 -76
  529. v0/relationalai/early_access/dsl/state_mgmt/__init__.py +0 -129
  530. v0/relationalai/early_access/dsl/state_mgmt/state_charts.py +0 -125
  531. v0/relationalai/early_access/dsl/state_mgmt/transitions.py +0 -130
  532. v0/relationalai/early_access/dsl/types/__init__.py +0 -40
  533. v0/relationalai/early_access/dsl/types/concepts.py +0 -12
  534. v0/relationalai/early_access/dsl/types/entities.py +0 -135
  535. v0/relationalai/early_access/dsl/types/values.py +0 -17
  536. v0/relationalai/early_access/dsl/utils.py +0 -102
  537. v0/relationalai/early_access/graphs/__init__.py +0 -13
  538. v0/relationalai/early_access/lqp/__init__.py +0 -12
  539. v0/relationalai/early_access/lqp/compiler/__init__.py +0 -12
  540. v0/relationalai/early_access/lqp/constructors/__init__.py +0 -18
  541. v0/relationalai/early_access/lqp/executor/__init__.py +0 -12
  542. v0/relationalai/early_access/lqp/ir/__init__.py +0 -12
  543. v0/relationalai/early_access/lqp/passes/__init__.py +0 -12
  544. v0/relationalai/early_access/lqp/pragmas/__init__.py +0 -12
  545. v0/relationalai/early_access/lqp/primitives/__init__.py +0 -12
  546. v0/relationalai/early_access/lqp/types/__init__.py +0 -12
  547. v0/relationalai/early_access/lqp/utils/__init__.py +0 -12
  548. v0/relationalai/early_access/lqp/validators/__init__.py +0 -12
  549. v0/relationalai/early_access/metamodel/__init__.py +0 -58
  550. v0/relationalai/early_access/metamodel/builtins/__init__.py +0 -12
  551. v0/relationalai/early_access/metamodel/compiler/__init__.py +0 -12
  552. v0/relationalai/early_access/metamodel/dependency/__init__.py +0 -12
  553. v0/relationalai/early_access/metamodel/factory/__init__.py +0 -17
  554. v0/relationalai/early_access/metamodel/helpers/__init__.py +0 -12
  555. v0/relationalai/early_access/metamodel/ir/__init__.py +0 -14
  556. v0/relationalai/early_access/metamodel/rewrite/__init__.py +0 -7
  557. v0/relationalai/early_access/metamodel/typer/__init__.py +0 -3
  558. v0/relationalai/early_access/metamodel/typer/typer/__init__.py +0 -12
  559. v0/relationalai/early_access/metamodel/types/__init__.py +0 -15
  560. v0/relationalai/early_access/metamodel/util/__init__.py +0 -15
  561. v0/relationalai/early_access/metamodel/visitor/__init__.py +0 -12
  562. v0/relationalai/early_access/rel/__init__.py +0 -12
  563. v0/relationalai/early_access/rel/executor/__init__.py +0 -12
  564. v0/relationalai/early_access/rel/rel_utils/__init__.py +0 -12
  565. v0/relationalai/early_access/rel/rewrite/__init__.py +0 -7
  566. v0/relationalai/early_access/solvers/__init__.py +0 -19
  567. v0/relationalai/early_access/sql/__init__.py +0 -11
  568. v0/relationalai/early_access/sql/executor/__init__.py +0 -3
  569. v0/relationalai/early_access/sql/rewrite/__init__.py +0 -3
  570. v0/relationalai/early_access/tests/logging/__init__.py +0 -12
  571. v0/relationalai/early_access/tests/test_snapshot_base/__init__.py +0 -12
  572. v0/relationalai/early_access/tests/utils/__init__.py +0 -12
  573. v0/relationalai/environments/__init__.py +0 -35
  574. v0/relationalai/environments/base.py +0 -381
  575. v0/relationalai/environments/colab.py +0 -14
  576. v0/relationalai/environments/generic.py +0 -71
  577. v0/relationalai/environments/ipython.py +0 -68
  578. v0/relationalai/environments/jupyter.py +0 -9
  579. v0/relationalai/environments/snowbook.py +0 -169
  580. v0/relationalai/errors.py +0 -2455
  581. v0/relationalai/experimental/SF.py +0 -38
  582. v0/relationalai/experimental/inspect.py +0 -47
  583. v0/relationalai/experimental/pathfinder/__init__.py +0 -158
  584. v0/relationalai/experimental/pathfinder/api.py +0 -160
  585. v0/relationalai/experimental/pathfinder/automaton.py +0 -584
  586. v0/relationalai/experimental/pathfinder/bridge.py +0 -226
  587. v0/relationalai/experimental/pathfinder/compiler.py +0 -416
  588. v0/relationalai/experimental/pathfinder/datalog.py +0 -214
  589. v0/relationalai/experimental/pathfinder/diagnostics.py +0 -56
  590. v0/relationalai/experimental/pathfinder/filter.py +0 -236
  591. v0/relationalai/experimental/pathfinder/glushkov.py +0 -439
  592. v0/relationalai/experimental/pathfinder/options.py +0 -265
  593. v0/relationalai/experimental/pathfinder/rpq.py +0 -344
  594. v0/relationalai/experimental/pathfinder/transition.py +0 -200
  595. v0/relationalai/experimental/pathfinder/utils.py +0 -26
  596. v0/relationalai/experimental/paths/api.py +0 -143
  597. v0/relationalai/experimental/paths/benchmarks/grid_graph.py +0 -37
  598. v0/relationalai/experimental/paths/examples/basic_example.py +0 -40
  599. v0/relationalai/experimental/paths/examples/minimal_engine_warmup.py +0 -3
  600. v0/relationalai/experimental/paths/examples/movie_example.py +0 -77
  601. v0/relationalai/experimental/paths/examples/paths_benchmark.py +0 -115
  602. v0/relationalai/experimental/paths/examples/paths_example.py +0 -116
  603. v0/relationalai/experimental/paths/examples/pattern_to_automaton.py +0 -28
  604. v0/relationalai/experimental/paths/find_paths_via_automaton.py +0 -85
  605. v0/relationalai/experimental/paths/graph.py +0 -185
  606. v0/relationalai/experimental/paths/path_algorithms/find_paths.py +0 -280
  607. v0/relationalai/experimental/paths/path_algorithms/one_sided_ball_repetition.py +0 -26
  608. v0/relationalai/experimental/paths/path_algorithms/one_sided_ball_upto.py +0 -111
  609. v0/relationalai/experimental/paths/path_algorithms/single.py +0 -59
  610. v0/relationalai/experimental/paths/path_algorithms/two_sided_balls_repetition.py +0 -39
  611. v0/relationalai/experimental/paths/path_algorithms/two_sided_balls_upto.py +0 -103
  612. v0/relationalai/experimental/paths/path_algorithms/usp-old.py +0 -130
  613. v0/relationalai/experimental/paths/path_algorithms/usp-tuple.py +0 -183
  614. v0/relationalai/experimental/paths/path_algorithms/usp.py +0 -150
  615. v0/relationalai/experimental/paths/product_graph.py +0 -93
  616. v0/relationalai/experimental/paths/rpq/automaton.py +0 -584
  617. v0/relationalai/experimental/paths/rpq/diagnostics.py +0 -56
  618. v0/relationalai/experimental/paths/rpq/rpq.py +0 -378
  619. v0/relationalai/experimental/paths/tests/tests_limit_sp_max_length.py +0 -90
  620. v0/relationalai/experimental/paths/tests/tests_limit_sp_multiple.py +0 -119
  621. v0/relationalai/experimental/paths/tests/tests_limit_sp_single.py +0 -104
  622. v0/relationalai/experimental/paths/tests/tests_limit_walks_multiple.py +0 -113
  623. v0/relationalai/experimental/paths/tests/tests_limit_walks_single.py +0 -149
  624. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_multiple.py +0 -70
  625. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_repetition_single.py +0 -64
  626. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_upto_multiple.py +0 -115
  627. v0/relationalai/experimental/paths/tests/tests_one_sided_ball_upto_single.py +0 -75
  628. v0/relationalai/experimental/paths/tests/tests_single_paths.py +0 -152
  629. v0/relationalai/experimental/paths/tests/tests_single_walks.py +0 -208
  630. v0/relationalai/experimental/paths/tests/tests_single_walks_undirected.py +0 -297
  631. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_multiple.py +0 -107
  632. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_repetition_single.py +0 -76
  633. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_upto_multiple.py +0 -76
  634. v0/relationalai/experimental/paths/tests/tests_two_sided_balls_upto_single.py +0 -110
  635. v0/relationalai/experimental/paths/tests/tests_usp_nsp_multiple.py +0 -229
  636. v0/relationalai/experimental/paths/tests/tests_usp_nsp_single.py +0 -108
  637. v0/relationalai/experimental/paths/tree_agg.py +0 -168
  638. v0/relationalai/experimental/paths/utilities/iterators.py +0 -27
  639. v0/relationalai/experimental/paths/utilities/prefix_sum.py +0 -91
  640. v0/relationalai/experimental/solvers.py +0 -1087
  641. v0/relationalai/loaders/csv.py +0 -195
  642. v0/relationalai/loaders/loader.py +0 -177
  643. v0/relationalai/loaders/types.py +0 -23
  644. v0/relationalai/rel_emitter.py +0 -373
  645. v0/relationalai/rel_utils.py +0 -185
  646. v0/relationalai/semantics/__init__.py +0 -29
  647. v0/relationalai/semantics/devtools/benchmark_lqp.py +0 -536
  648. v0/relationalai/semantics/devtools/compilation_manager.py +0 -294
  649. v0/relationalai/semantics/devtools/extract_lqp.py +0 -110
  650. v0/relationalai/semantics/internal/internal.py +0 -3785
  651. v0/relationalai/semantics/internal/snowflake.py +0 -324
  652. v0/relationalai/semantics/lqp/builtins.py +0 -16
  653. v0/relationalai/semantics/lqp/compiler.py +0 -22
  654. v0/relationalai/semantics/lqp/constructors.py +0 -68
  655. v0/relationalai/semantics/lqp/executor.py +0 -469
  656. v0/relationalai/semantics/lqp/intrinsics.py +0 -24
  657. v0/relationalai/semantics/lqp/model2lqp.py +0 -839
  658. v0/relationalai/semantics/lqp/passes.py +0 -680
  659. v0/relationalai/semantics/lqp/primitives.py +0 -252
  660. v0/relationalai/semantics/lqp/result_helpers.py +0 -202
  661. v0/relationalai/semantics/lqp/rewrite/annotate_constraints.py +0 -57
  662. v0/relationalai/semantics/lqp/rewrite/cdc.py +0 -216
  663. v0/relationalai/semantics/lqp/rewrite/extract_common.py +0 -338
  664. v0/relationalai/semantics/lqp/rewrite/extract_keys.py +0 -449
  665. v0/relationalai/semantics/lqp/rewrite/function_annotations.py +0 -114
  666. v0/relationalai/semantics/lqp/rewrite/functional_dependencies.py +0 -314
  667. v0/relationalai/semantics/lqp/rewrite/quantify_vars.py +0 -296
  668. v0/relationalai/semantics/lqp/rewrite/splinter.py +0 -76
  669. v0/relationalai/semantics/lqp/types.py +0 -101
  670. v0/relationalai/semantics/lqp/utils.py +0 -160
  671. v0/relationalai/semantics/lqp/validators.py +0 -57
  672. v0/relationalai/semantics/metamodel/__init__.py +0 -40
  673. v0/relationalai/semantics/metamodel/builtins.py +0 -774
  674. v0/relationalai/semantics/metamodel/compiler.py +0 -133
  675. v0/relationalai/semantics/metamodel/dependency.py +0 -862
  676. v0/relationalai/semantics/metamodel/executor.py +0 -61
  677. v0/relationalai/semantics/metamodel/factory.py +0 -287
  678. v0/relationalai/semantics/metamodel/helpers.py +0 -361
  679. v0/relationalai/semantics/metamodel/rewrite/discharge_constraints.py +0 -39
  680. v0/relationalai/semantics/metamodel/rewrite/dnf_union_splitter.py +0 -210
  681. v0/relationalai/semantics/metamodel/rewrite/extract_nested_logicals.py +0 -78
  682. v0/relationalai/semantics/metamodel/rewrite/flatten.py +0 -549
  683. v0/relationalai/semantics/metamodel/rewrite/format_outputs.py +0 -165
  684. v0/relationalai/semantics/metamodel/typer/checker.py +0 -353
  685. v0/relationalai/semantics/metamodel/typer/typer.py +0 -1395
  686. v0/relationalai/semantics/metamodel/util.py +0 -505
  687. v0/relationalai/semantics/reasoners/__init__.py +0 -10
  688. v0/relationalai/semantics/reasoners/graph/__init__.py +0 -37
  689. v0/relationalai/semantics/reasoners/graph/core.py +0 -9020
  690. v0/relationalai/semantics/reasoners/optimization/__init__.py +0 -68
  691. v0/relationalai/semantics/reasoners/optimization/common.py +0 -88
  692. v0/relationalai/semantics/reasoners/optimization/solvers_dev.py +0 -568
  693. v0/relationalai/semantics/reasoners/optimization/solvers_pb.py +0 -1163
  694. v0/relationalai/semantics/rel/builtins.py +0 -40
  695. v0/relationalai/semantics/rel/compiler.py +0 -989
  696. v0/relationalai/semantics/rel/executor.py +0 -359
  697. v0/relationalai/semantics/rel/rel.py +0 -482
  698. v0/relationalai/semantics/rel/rel_utils.py +0 -276
  699. v0/relationalai/semantics/snowflake/__init__.py +0 -3
  700. v0/relationalai/semantics/sql/compiler.py +0 -2503
  701. v0/relationalai/semantics/sql/executor/duck_db.py +0 -52
  702. v0/relationalai/semantics/sql/executor/result_helpers.py +0 -64
  703. v0/relationalai/semantics/sql/executor/snowflake.py +0 -145
  704. v0/relationalai/semantics/sql/rewrite/denormalize.py +0 -222
  705. v0/relationalai/semantics/sql/rewrite/double_negation.py +0 -49
  706. v0/relationalai/semantics/sql/rewrite/recursive_union.py +0 -127
  707. v0/relationalai/semantics/sql/rewrite/sort_output_query.py +0 -246
  708. v0/relationalai/semantics/sql/sql.py +0 -504
  709. v0/relationalai/semantics/std/__init__.py +0 -54
  710. v0/relationalai/semantics/std/constraints.py +0 -43
  711. v0/relationalai/semantics/std/datetime.py +0 -363
  712. v0/relationalai/semantics/std/decimals.py +0 -62
  713. v0/relationalai/semantics/std/floats.py +0 -7
  714. v0/relationalai/semantics/std/integers.py +0 -22
  715. v0/relationalai/semantics/std/math.py +0 -141
  716. v0/relationalai/semantics/std/pragmas.py +0 -11
  717. v0/relationalai/semantics/std/re.py +0 -83
  718. v0/relationalai/semantics/std/std.py +0 -14
  719. v0/relationalai/semantics/std/strings.py +0 -63
  720. v0/relationalai/semantics/tests/__init__.py +0 -0
  721. v0/relationalai/semantics/tests/test_snapshot_abstract.py +0 -143
  722. v0/relationalai/semantics/tests/test_snapshot_base.py +0 -9
  723. v0/relationalai/semantics/tests/utils.py +0 -46
  724. v0/relationalai/std/__init__.py +0 -70
  725. v0/relationalai/tools/__init__.py +0 -0
  726. v0/relationalai/tools/cli.py +0 -1940
  727. v0/relationalai/tools/cli_controls.py +0 -1826
  728. v0/relationalai/tools/cli_helpers.py +0 -390
  729. v0/relationalai/tools/debugger.py +0 -183
  730. v0/relationalai/tools/debugger_client.py +0 -109
  731. v0/relationalai/tools/debugger_server.py +0 -302
  732. v0/relationalai/tools/dev.py +0 -685
  733. v0/relationalai/tools/qb_debugger.py +0 -425
  734. v0/relationalai/util/clean_up_databases.py +0 -95
  735. v0/relationalai/util/format.py +0 -123
  736. v0/relationalai/util/list_databases.py +0 -9
  737. v0/relationalai/util/otel_configuration.py +0 -25
  738. v0/relationalai/util/otel_handler.py +0 -484
  739. v0/relationalai/util/snowflake_handler.py +0 -88
  740. v0/relationalai/util/span_format_test.py +0 -43
  741. v0/relationalai/util/span_tracker.py +0 -207
  742. v0/relationalai/util/spans_file_handler.py +0 -72
  743. v0/relationalai/util/tracing_handler.py +0 -34
  744. /relationalai/{semantics/frontend → analysis}/__init__.py +0 -0
  745. {v0/relationalai → relationalai}/analysis/mechanistic.py +0 -0
  746. {v0/relationalai → relationalai}/analysis/whynot.py +0 -0
  747. /relationalai/{shims → auth}/__init__.py +0 -0
  748. {v0/relationalai → relationalai}/auth/jwt_generator.py +0 -0
  749. {v0/relationalai → relationalai}/auth/oauth_callback_server.py +0 -0
  750. {v0/relationalai → relationalai}/auth/token_handler.py +0 -0
  751. {v0/relationalai → relationalai}/auth/util.py +0 -0
  752. {v0/relationalai/clients → relationalai/clients/resources/snowflake}/cache_store.py +0 -0
  753. {v0/relationalai → relationalai}/compiler.py +0 -0
  754. {v0/relationalai → relationalai}/dependencies.py +0 -0
  755. {v0/relationalai → relationalai}/docutils.py +0 -0
  756. {v0/relationalai/analysis → relationalai/early_access}/__init__.py +0 -0
  757. {v0/relationalai → relationalai}/early_access/dsl/__init__.py +0 -0
  758. {v0/relationalai/auth → relationalai/early_access/dsl/adapters}/__init__.py +0 -0
  759. {v0/relationalai/early_access → relationalai/early_access/dsl/adapters/orm}/__init__.py +0 -0
  760. {v0/relationalai → relationalai}/early_access/dsl/adapters/orm/model.py +0 -0
  761. {v0/relationalai/early_access/dsl/adapters → relationalai/early_access/dsl/adapters/owl}/__init__.py +0 -0
  762. {v0/relationalai → relationalai}/early_access/dsl/adapters/owl/model.py +0 -0
  763. {v0/relationalai/early_access/dsl/adapters/orm → relationalai/early_access/dsl/bindings}/__init__.py +0 -0
  764. {v0/relationalai/early_access/dsl/adapters/owl → relationalai/early_access/dsl/bindings/legacy}/__init__.py +0 -0
  765. {v0/relationalai/early_access/dsl/bindings → relationalai/early_access/dsl/codegen}/__init__.py +0 -0
  766. {v0/relationalai → relationalai}/early_access/dsl/constants.py +0 -0
  767. {v0/relationalai → relationalai}/early_access/dsl/core/__init__.py +0 -0
  768. {v0/relationalai → relationalai}/early_access/dsl/core/constraints/__init__.py +0 -0
  769. {v0/relationalai → relationalai}/early_access/dsl/core/constraints/predicate/__init__.py +0 -0
  770. {v0/relationalai → relationalai}/early_access/dsl/core/stack.py +0 -0
  771. {v0/relationalai/early_access/dsl/bindings/legacy → relationalai/early_access/dsl/core/temporal}/__init__.py +0 -0
  772. {v0/relationalai → relationalai}/early_access/dsl/core/utils.py +0 -0
  773. {v0/relationalai/early_access/dsl/codegen → relationalai/early_access/dsl/ir}/__init__.py +0 -0
  774. {v0/relationalai/early_access/dsl/core/temporal → relationalai/early_access/dsl/ontologies}/__init__.py +0 -0
  775. {v0/relationalai → relationalai}/early_access/dsl/ontologies/raw_source.py +0 -0
  776. {v0/relationalai/early_access/dsl/ir → relationalai/early_access/dsl/orm}/__init__.py +0 -0
  777. {v0/relationalai/early_access/dsl/ontologies → relationalai/early_access/dsl/orm/measures}/__init__.py +0 -0
  778. {v0/relationalai → relationalai}/early_access/dsl/orm/reasoner_errors.py +0 -0
  779. {v0/relationalai/early_access/dsl/orm → relationalai/early_access/dsl/physical_metadata}/__init__.py +0 -0
  780. {v0/relationalai/early_access/dsl/orm/measures → relationalai/early_access/dsl/serialize}/__init__.py +0 -0
  781. {v0/relationalai → relationalai}/early_access/dsl/serialize/binding_model.py +0 -0
  782. {v0/relationalai → relationalai}/early_access/dsl/serialize/model.py +0 -0
  783. {v0/relationalai/early_access/dsl/physical_metadata → relationalai/early_access/dsl/snow}/__init__.py +0 -0
  784. {v0/relationalai → relationalai}/early_access/tests/__init__.py +0 -0
  785. {v0/relationalai → relationalai}/environments/ci.py +0 -0
  786. {v0/relationalai → relationalai}/environments/hex.py +0 -0
  787. {v0/relationalai → relationalai}/environments/terminal.py +0 -0
  788. {v0/relationalai → relationalai}/experimental/__init__.py +0 -0
  789. {v0/relationalai → relationalai}/experimental/graphs.py +0 -0
  790. {v0/relationalai → relationalai}/experimental/paths/__init__.py +0 -0
  791. {v0/relationalai → relationalai}/experimental/paths/benchmarks/__init__.py +0 -0
  792. {v0/relationalai → relationalai}/experimental/paths/path_algorithms/__init__.py +0 -0
  793. {v0/relationalai → relationalai}/experimental/paths/rpq/__init__.py +0 -0
  794. {v0/relationalai → relationalai}/experimental/paths/rpq/filter.py +0 -0
  795. {v0/relationalai → relationalai}/experimental/paths/rpq/glushkov.py +0 -0
  796. {v0/relationalai → relationalai}/experimental/paths/rpq/transition.py +0 -0
  797. {v0/relationalai → relationalai}/experimental/paths/utilities/__init__.py +0 -0
  798. {v0/relationalai → relationalai}/experimental/paths/utilities/utilities.py +0 -0
  799. {v0/relationalai/early_access/dsl/serialize → relationalai/loaders}/__init__.py +0 -0
  800. {v0/relationalai → relationalai}/metagen.py +0 -0
  801. {v0/relationalai → relationalai}/metamodel.py +0 -0
  802. {v0/relationalai → relationalai}/rel.py +0 -0
  803. {v0/relationalai → relationalai}/semantics/devtools/__init__.py +0 -0
  804. {v0/relationalai → relationalai}/semantics/internal/__init__.py +0 -0
  805. {v0/relationalai → relationalai}/semantics/internal/annotations.py +0 -0
  806. {v0/relationalai → relationalai}/semantics/lqp/__init__.py +0 -0
  807. {v0/relationalai → relationalai}/semantics/lqp/ir.py +0 -0
  808. {v0/relationalai → relationalai}/semantics/lqp/pragmas.py +0 -0
  809. {v0/relationalai → relationalai}/semantics/lqp/rewrite/__init__.py +0 -0
  810. {v0/relationalai → relationalai}/semantics/metamodel/dataflow.py +0 -0
  811. {v0/relationalai → relationalai}/semantics/metamodel/ir.py +0 -0
  812. {v0/relationalai → relationalai}/semantics/metamodel/rewrite/__init__.py +0 -0
  813. {v0/relationalai → relationalai}/semantics/metamodel/typer/__init__.py +0 -0
  814. {v0/relationalai → relationalai}/semantics/metamodel/types.py +0 -0
  815. {v0/relationalai → relationalai}/semantics/metamodel/visitor.py +0 -0
  816. {v0/relationalai → relationalai}/semantics/reasoners/experimental/__init__.py +0 -0
  817. {v0/relationalai → relationalai}/semantics/rel/__init__.py +0 -0
  818. {v0/relationalai → relationalai}/semantics/sql/__init__.py +0 -0
  819. {v0/relationalai → relationalai}/semantics/sql/executor/__init__.py +0 -0
  820. {v0/relationalai → relationalai}/semantics/sql/rewrite/__init__.py +0 -0
  821. {v0/relationalai/early_access/dsl/snow → relationalai/semantics/tests}/__init__.py +0 -0
  822. {v0/relationalai → relationalai}/semantics/tests/logging.py +0 -0
  823. {v0/relationalai → relationalai}/std/aggregates.py +0 -0
  824. {v0/relationalai → relationalai}/std/dates.py +0 -0
  825. {v0/relationalai → relationalai}/std/graphs.py +0 -0
  826. {v0/relationalai → relationalai}/std/inspect.py +0 -0
  827. {v0/relationalai → relationalai}/std/math.py +0 -0
  828. {v0/relationalai → relationalai}/std/re.py +0 -0
  829. {v0/relationalai → relationalai}/std/strings.py +0 -0
  830. {v0/relationalai/loaders → relationalai/tools}/__init__.py +0 -0
  831. {v0/relationalai → relationalai}/tools/cleanup_snapshots.py +0 -0
  832. {v0/relationalai → relationalai}/tools/constants.py +0 -0
  833. {v0/relationalai → relationalai}/tools/query_utils.py +0 -0
  834. {v0/relationalai → relationalai}/tools/snapshot_viewer.py +0 -0
  835. {v0/relationalai → relationalai}/util/__init__.py +0 -0
  836. {v0/relationalai → relationalai}/util/constants.py +0 -0
  837. {v0/relationalai → relationalai}/util/graph.py +0 -0
  838. {v0/relationalai → relationalai}/util/timeout.py +0 -0
@@ -0,0 +1,3185 @@
1
+ # pyright: reportUnusedExpression=false
2
+ from __future__ import annotations
3
+ import base64
4
+ import importlib.resources
5
+ import io
6
+ import re
7
+ import json
8
+ import time
9
+ import textwrap
10
+ import ast
11
+ import uuid
12
+ import warnings
13
+ import atexit
14
+ import hashlib
15
+ from dataclasses import dataclass
16
+
17
+ from ....auth.token_handler import TokenHandler
18
+ from relationalai.clients.exec_txn_poller import ExecTxnPoller, query_complete_message
19
+ import snowflake.snowpark
20
+
21
+ from ....rel_utils import sanitize_identifier, to_fqn_relation_name
22
+ from ....tools.constants import FIELD_PLACEHOLDER, SNOWFLAKE_AUTHS, USE_GRAPH_INDEX, DEFAULT_QUERY_TIMEOUT_MINS, WAIT_FOR_STREAM_SYNC, Generation
23
+ from .... import std
24
+ from collections import defaultdict
25
+ import requests
26
+ import snowflake.connector
27
+ import pyarrow as pa
28
+
29
+ from snowflake.snowpark import Session
30
+ from snowflake.snowpark.context import get_active_session
31
+ from ... import result_helpers
32
+ from .... import debugging
33
+ from typing import Any, Dict, Iterable, Tuple, List, Literal, cast
34
+
35
+ from pandas import DataFrame
36
+
37
+ from ....tools.cli_controls import Spinner
38
+ from ...types import AvailableModel, EngineState, Import, ImportSource, ImportSourceTable, ImportsStatus, SourceInfo, TransactionAsyncResponse
39
+ from ...config import Config
40
+ from ...client import Client, ExportParams, ProviderBase, ResourcesBase
41
+ from ...util import IdentityParser, escape_for_f_string, get_pyrel_version, get_with_retries, poll_with_specified_overhead, safe_json_loads, sanitize_module_name, scrub_exception, wrap_with_request_id, normalize_datetime
42
+ from .engine_service import EngineServiceSQL, EngineType
43
+ from .util import (
44
+ collect_error_messages,
45
+ process_jinja_template,
46
+ type_to_sql,
47
+ type_to_snowpark,
48
+ sanitize_user_name as _sanitize_user_name,
49
+ normalize_params,
50
+ format_sproc_name,
51
+ is_azure_url,
52
+ is_container_runtime,
53
+ imports_to_dicts,
54
+ txn_list_to_dicts,
55
+ decrypt_artifact,
56
+ )
57
+ from ....environments import runtime_env, HexEnvironment, SnowbookEnvironment
58
+ from .... import dsl, rel, metamodel as m
59
+ from ....errors import EngineProvisioningFailed, EngineNameValidationException, Errors, GuardRailsException, InvalidAliasError, InvalidEngineSizeError, InvalidSourceTypeWarning, RAIException, HexSessionException, SnowflakeChangeTrackingNotEnabledException, SnowflakeDatabaseException, SnowflakeImportMissingException, SnowflakeInvalidSource, SnowflakeMissingConfigValuesException, SnowflakeProxyAPIDeprecationWarning, SnowflakeProxySourceError, ModelNotFoundException, UnknownSourceWarning, RowsDroppedFromTargetTableWarning, QueryTimeoutExceededException
60
+ from concurrent.futures import ThreadPoolExecutor
61
+ from datetime import datetime, timedelta
62
+ from snowflake.snowpark.types import StringType, StructField, StructType
63
+ # Import error handlers and constants
64
+ from .error_handlers import (
65
+ ErrorHandler,
66
+ DuoSecurityErrorHandler,
67
+ AppMissingErrorHandler,
68
+ AppFunctionMissingErrorHandler,
69
+ DatabaseErrorsHandler,
70
+ EngineErrorsHandler,
71
+ ServiceNotStartedErrorHandler,
72
+ TransactionAbortedErrorHandler,
73
+ )
74
+ # Import engine state handlers
75
+ from .engine_state_handlers import (
76
+ EngineStateHandler,
77
+ EngineContext,
78
+ SyncPendingStateHandler,
79
+ SyncSuspendedStateHandler,
80
+ SyncReadyStateHandler,
81
+ SyncGoneStateHandler,
82
+ SyncMissingEngineHandler,
83
+ AsyncPendingStateHandler,
84
+ AsyncSuspendedStateHandler,
85
+ AsyncReadyStateHandler,
86
+ AsyncGoneStateHandler,
87
+ AsyncMissingEngineHandler,
88
+ )
89
+
90
+
91
+ #--------------------------------------------------
92
+ # Constants
93
+ #--------------------------------------------------
94
+
95
+ # transaction list and get return different fields (duration vs timings)
96
+ LIST_TXN_SQL_FIELDS = ["id", "database_name", "engine_name", "state", "abort_reason", "read_only","created_by", "created_on", "finished_at", "duration"]
97
+ GET_TXN_SQL_FIELDS = ["id", "database", "engine", "state", "abort_reason", "read_only","created_by", "created_on", "finished_at", "timings"]
98
+ VALID_ENGINE_STATES = ["READY", "PENDING"]
99
+ # Note: ENGINE_ERRORS, ENGINE_NOT_READY_MSGS, DATABASE_ERRORS moved to util.py
100
+ PYREL_ROOT_DB = 'pyrel_root_db'
101
+
102
+ TERMINAL_TXN_STATES = ["COMPLETED", "ABORTED"]
103
+
104
+ TXN_ABORT_REASON_TIMEOUT = "transaction timeout"
105
+ GUARDRAILS_ABORT_REASON = "guard rail violation"
106
+
107
+ PRINT_TXN_PROGRESS_FLAG = "print_txn_progress"
108
+
109
+ #--------------------------------------------------
110
+ # Helpers
111
+ #--------------------------------------------------
112
+
113
+ def should_print_txn_progress(config) -> bool:
114
+ return bool(config.get(PRINT_TXN_PROGRESS_FLAG, False))
115
+
116
+ #--------------------------------------------------
117
+ # Resources
118
+ #--------------------------------------------------
119
+
120
+ APP_NAME = "___RAI_APP___"
121
+
122
+ @dataclass
123
+ class ExecContext:
124
+ """Execution context for SQL queries, containing all parameters needed for execution and retry."""
125
+ code: str
126
+ params: List[Any] | None = None
127
+ raw: bool = False
128
+ help: bool = True
129
+ skip_engine_db_error_retry: bool = False
130
+
131
+ def re_execute(self, resources: 'Resources') -> Any:
132
+ """Re-execute this context's query using the provided resources instance."""
133
+ return resources._exec(
134
+ code=self.code,
135
+ params=self.params,
136
+ raw=self.raw,
137
+ help=self.help,
138
+ skip_engine_db_error_retry=self.skip_engine_db_error_retry
139
+ )
140
+
141
+
142
+ @dataclass
143
+ class TxnCreationResult:
144
+ """Result of creating a transaction via _create_v2_txn.
145
+
146
+ This standardizes the response format between different implementations
147
+ (SQL stored procedure vs HTTP direct access).
148
+ """
149
+ txn_id: str
150
+ state: str
151
+ artifact_info: Dict[str, Dict] # Populated if fast-path (state is COMPLETED/ABORTED)
152
+
153
+
154
+ class Resources(ResourcesBase):
155
+ def __init__(
156
+ self,
157
+ profile: str | None = None,
158
+ config: Config | None = None,
159
+ connection: Session | None = None,
160
+ dry_run: bool = False,
161
+ reset_session: bool = False,
162
+ generation: Generation | None = None,
163
+ language: str = "rel", # Accepted for backward compatibility, but not stored in base class
164
+ ):
165
+ super().__init__(profile, config=config)
166
+ self._token_handler: TokenHandler | None = None
167
+ self._session = connection
168
+ self.generation = generation
169
+ if self._session is None and not dry_run:
170
+ try:
171
+ # we may still be constructing the config, so this can fail now,
172
+ # if so we'll create later
173
+ self._session = self.get_sf_session(reset_session)
174
+ except Exception:
175
+ pass
176
+ self._pending_transactions: list[str] = []
177
+ self._ns_cache = {}
178
+ # self.sources contains fully qualified Snowflake table/view names
179
+ self.sources: set[str] = set()
180
+ self._sproc_models = None
181
+ # Store language for backward compatibility (used by child classes for use_index polling)
182
+ self.language = language
183
+ # Engine subsystem (composition: keeps engine CRUD isolated from the core Resources class)
184
+ self._engines = EngineServiceSQL(self)
185
+ # Register error and state handlers
186
+ self._register_handlers()
187
+ # Register atexit callback to cancel pending transactions
188
+ atexit.register(self.cancel_pending_transactions)
189
+
190
+ @property
191
+ def engines(self) -> EngineServiceSQL:
192
+ return self._engines
193
+
194
+ #--------------------------------------------------
195
+ # Initialization & Properties
196
+ #--------------------------------------------------
197
+
198
+ def _register_handlers(self) -> None:
199
+ """Register error and engine state handlers for processing."""
200
+ # Register base handlers using getter methods that subclasses can override
201
+ # Use defensive copying to ensure each instance has its own handler lists
202
+ # and prevent cross-instance contamination from subclass mutations
203
+ self._error_handlers = list(self._get_error_handlers())
204
+ self._sync_engine_state_handlers = list(self._get_engine_state_handlers(is_async=False))
205
+ self._async_engine_state_handlers = list(self._get_engine_state_handlers(is_async=True))
206
+
207
+ def _get_error_handlers(self) -> list[ErrorHandler]:
208
+ """Get list of error handlers. Subclasses can override to add custom handlers.
209
+
210
+ Returns:
211
+ List of error handlers for standard error processing using Strategy Pattern.
212
+
213
+ Example:
214
+ def _get_error_handlers(self) -> list[ErrorHandler]:
215
+ # Get base handlers
216
+ handlers = super()._get_error_handlers()
217
+ # Add custom handler
218
+ handlers.append(MyCustomErrorHandler())
219
+ return handlers
220
+ """
221
+ return [
222
+ AppMissingErrorHandler(),
223
+ AppFunctionMissingErrorHandler(),
224
+ ServiceNotStartedErrorHandler(),
225
+ DuoSecurityErrorHandler(),
226
+ DatabaseErrorsHandler(),
227
+ EngineErrorsHandler(),
228
+ TransactionAbortedErrorHandler(),
229
+ ]
230
+
231
+ def _get_engine_state_handlers(self, is_async: bool = False) -> list[EngineStateHandler]:
232
+ """Get list of engine state handlers. Subclasses can override.
233
+
234
+ Args:
235
+ is_async: If True, returns async handlers; if False, returns sync handlers.
236
+
237
+ Returns:
238
+ List of engine state handlers for processing engine states.
239
+
240
+ Example:
241
+ def _get_engine_state_handlers(self, is_async: bool = False) -> list[EngineStateHandler]:
242
+ # Get base handlers
243
+ handlers = super()._get_engine_state_handlers(is_async)
244
+ # Add custom handler
245
+ handlers.append(MyCustomStateHandler())
246
+ return handlers
247
+ """
248
+ if is_async:
249
+ return [
250
+ AsyncPendingStateHandler(),
251
+ AsyncSuspendedStateHandler(),
252
+ AsyncReadyStateHandler(),
253
+ AsyncGoneStateHandler(),
254
+ AsyncMissingEngineHandler(),
255
+ ]
256
+ else:
257
+ return [
258
+ SyncPendingStateHandler(),
259
+ SyncSuspendedStateHandler(),
260
+ SyncReadyStateHandler(),
261
+ SyncGoneStateHandler(),
262
+ SyncMissingEngineHandler(),
263
+ ]
264
+
265
+ @property
266
+ def token_handler(self) -> TokenHandler:
267
+ if not self._token_handler:
268
+ self._token_handler = TokenHandler.from_config(self.config)
269
+ return self._token_handler
270
+
271
+ def reset(self):
272
+ """Reset the session."""
273
+ self._session = None
274
+
275
+ #--------------------------------------------------
276
+ # Session Management
277
+ #--------------------------------------------------
278
+
279
+ def is_erp_running(self, app_name: str) -> bool:
280
+ """Check if the ERP is running. The app.service_status() returns single row/column containing an array of JSON service status objects."""
281
+ query = f"CALL {app_name}.app.service_status();"
282
+ try:
283
+ result = self._exec(query)
284
+ # The result is a list of dictionaries, each with a "STATUS" key
285
+ # The column name containing the result is "SERVICE_STATUS"
286
+ services_status = json.loads(result[0]["SERVICE_STATUS"])
287
+ # Find the dictionary with "name" of "main" and check if its "status" is "READY"
288
+ for service in services_status:
289
+ if service.get("name") == "main" and service.get("status") == "READY":
290
+ return True
291
+ return False
292
+ except Exception:
293
+ return False
294
+
295
+ def get_sf_session(self, reset_session: bool = False):
296
+ if self._session:
297
+ return self._session
298
+
299
+ if isinstance(runtime_env, HexEnvironment):
300
+ raise HexSessionException()
301
+ if isinstance(runtime_env, SnowbookEnvironment):
302
+ return get_active_session()
303
+ else:
304
+ # if there's already been a session created, try using that
305
+ # if reset_session is true always try to get the new session
306
+ if not reset_session:
307
+ try:
308
+ return get_active_session()
309
+ except Exception:
310
+ pass
311
+
312
+ # otherwise, create a new session
313
+ missing_keys = []
314
+ connection_parameters = {}
315
+
316
+ authenticator = self.config.get('authenticator', None)
317
+ passcode = self.config.get("passcode", "")
318
+ private_key_file = self.config.get("private_key_file", "")
319
+
320
+ # If the authenticator is not set, we need to set it based on the provided parameters
321
+ if authenticator is None:
322
+ if private_key_file != "":
323
+ authenticator = "snowflake_jwt"
324
+ elif passcode != "":
325
+ authenticator = "username_password_mfa"
326
+ else:
327
+ authenticator = "snowflake"
328
+ # set the default authenticator in the config so we can skip it when we check for missing keys
329
+ self.config.set("authenticator", authenticator)
330
+
331
+ if authenticator in SNOWFLAKE_AUTHS:
332
+ required_keys = {
333
+ key for key, value in SNOWFLAKE_AUTHS[authenticator].items() if value.get("required", True)
334
+ }
335
+ for key in required_keys:
336
+ if self.config.get(key, None) is None:
337
+ default = SNOWFLAKE_AUTHS[authenticator][key].get("value", None)
338
+ if default is None or default == FIELD_PLACEHOLDER:
339
+ # No default value and no value in the config, add to missing keys
340
+ missing_keys.append(key)
341
+ else:
342
+ # Set the default value in the config from the auth defaults
343
+ self.config.set(key, default)
344
+ if missing_keys:
345
+ profile = getattr(self.config, 'profile', None)
346
+ config_file_path = getattr(self.config, 'file_path', None)
347
+ raise SnowflakeMissingConfigValuesException(missing_keys, profile, config_file_path)
348
+ for key in SNOWFLAKE_AUTHS[authenticator]:
349
+ connection_parameters[key] = self.config.get(key, None)
350
+ else:
351
+ raise ValueError(f'Authenticator "{authenticator}" not supported')
352
+
353
+ return self._build_snowflake_session(connection_parameters)
354
+
355
+ def _build_snowflake_session(self, connection_parameters: Dict[str, Any]) -> Session:
356
+ try:
357
+ tmp = {
358
+ "client_session_keep_alive": True,
359
+ "client_session_keep_alive_heartbeat_frequency": 60 * 5,
360
+ }
361
+ tmp.update(connection_parameters)
362
+ connection_parameters = tmp
363
+ # authenticator programmatic access token needs to be upper cased to work...
364
+ connection_parameters["authenticator"] = connection_parameters["authenticator"].upper()
365
+ if "authenticator" in connection_parameters and connection_parameters["authenticator"] == "OAUTH_AUTHORIZATION_CODE":
366
+ # we are replicating OAUTH_AUTHORIZATION_CODE by first retrieving the token
367
+ # and then authenticating with the token via the OAUTH authenticator
368
+ connection_parameters["token"] = self.token_handler.get_session_login_token()
369
+ connection_parameters["authenticator"] = "OAUTH"
370
+ return Session.builder.configs(connection_parameters).create()
371
+ except snowflake.connector.errors.Error as e:
372
+ raise SnowflakeDatabaseException(e)
373
+ except Exception as e:
374
+ raise e
375
+
376
+ #--------------------------------------------------
377
+ # Core Execution Methods
378
+ #--------------------------------------------------
379
+
380
+ def _exec_sql(self, code: str, params: List[Any] | None, raw=False):
381
+ """
382
+ Lowest-level SQL execution method.
383
+
384
+ Directly executes SQL via the Snowflake session. This is the foundation
385
+ for all other execution methods. It:
386
+ - Replaces APP_NAME placeholder with actual app name
387
+ - Executes SQL with optional parameters
388
+ - Returns either raw session results or collected results
389
+
390
+ Args:
391
+ code: SQL code to execute (may contain APP_NAME placeholder)
392
+ params: Optional SQL parameters
393
+ raw: If True, return raw session results; if False, collect results
394
+
395
+ Returns:
396
+ Raw session results if raw=True, otherwise collected results
397
+ """
398
+ assert self._session is not None
399
+ sess_results = self._session.sql(
400
+ code.replace(APP_NAME, self.get_app_name()),
401
+ params
402
+ )
403
+ if raw:
404
+ return sess_results
405
+ return sess_results.collect()
406
+
407
+ def _exec(
408
+ self,
409
+ code: str,
410
+ params: List[Any] | Any | None = None,
411
+ raw: bool = False,
412
+ help: bool = True,
413
+ skip_engine_db_error_retry: bool = False
414
+ ) -> Any:
415
+ """
416
+ Mid-level SQL execution method with error handling.
417
+
418
+ This is the primary method for executing SQL queries. It wraps _exec_sql
419
+ with comprehensive error handling and parameter normalization. Used
420
+ extensively throughout the codebase for direct SQL operations like:
421
+ - SHOW commands (warehouses, databases, etc.)
422
+ - CALL statements to RAI app stored procedures
423
+ - Transaction management queries
424
+
425
+ The error handling flow:
426
+ 1. Normalizes parameters and creates execution context
427
+ 2. Calls _exec_sql to execute the query
428
+ 3. On error, uses standard error handling (Strategy Pattern), which subclasses
429
+ can influence via `_get_error_handlers()` or by overriding `_handle_standard_exec_errors()`
430
+
431
+ Args:
432
+ code: SQL code to execute
433
+ params: Optional SQL parameters (normalized to list if needed)
434
+ raw: If True, return raw session results; if False, collect results
435
+ help: If True, enable error handling; if False, raise errors immediately
436
+ skip_engine_db_error_retry: If True, skip use_index retry logic in error handlers
437
+
438
+ Returns:
439
+ Query results (collected or raw depending on 'raw' parameter)
440
+ """
441
+ # print(f"\n--- sql---\n{code}\n--- end sql---\n")
442
+ # Ensure session is initialized
443
+ if not self._session:
444
+ self._session = self.get_sf_session()
445
+
446
+ # Normalize parameters
447
+ normalized_params = normalize_params(params)
448
+
449
+ # Create execution context
450
+ ctx = ExecContext(
451
+ code=code,
452
+ params=normalized_params,
453
+ raw=raw,
454
+ help=help,
455
+ skip_engine_db_error_retry=skip_engine_db_error_retry
456
+ )
457
+
458
+ # Execute SQL
459
+ try:
460
+ return self._exec_sql(ctx.code, ctx.params, raw=ctx.raw)
461
+ except Exception as e:
462
+ if not ctx.help:
463
+ raise e
464
+
465
+ # Handle standard errors
466
+ result = self._handle_standard_exec_errors(e, ctx)
467
+ if result is not None:
468
+ return result
469
+
470
+ #--------------------------------------------------
471
+ # Error Handling
472
+ #--------------------------------------------------
473
+
474
+ def _handle_standard_exec_errors(self, e: Exception, ctx: ExecContext) -> Any | None:
475
+ """
476
+ Handle standard Snowflake/RAI errors using Strategy Pattern.
477
+
478
+ Each error type has a dedicated handler class that encapsulates
479
+ the detection logic and exception creation. Handlers are processed
480
+ in order until one matches and handles the error.
481
+ """
482
+ message = str(e).lower()
483
+
484
+ # Try each handler in order until one matches
485
+ for handler in self._error_handlers:
486
+ if handler.matches(e, message, ctx, self):
487
+ result = handler.handle(e, ctx, self)
488
+ if result is not None:
489
+ return result
490
+ return # Handler raised exception, we're done
491
+
492
+ # Fallback: transform to RAIException
493
+ raise RAIException(str(e))
494
+
495
+ #--------------------------------------------------
496
+ # Feature Detection & Configuration
497
+ #--------------------------------------------------
498
+
499
+ def is_direct_access_enabled(self) -> bool:
500
+ try:
501
+ feature_enabled = self._exec(
502
+ f"call {APP_NAME}.APP.DIRECT_INGRESS_ENABLED();"
503
+ )
504
+ if not feature_enabled:
505
+ return False
506
+
507
+ # Even if the feature is enabled, customers still need to reactivate ERP to ensure the endpoint is available.
508
+ endpoint = self._exec(
509
+ f"call {APP_NAME}.APP.SERVICE_ENDPOINT(true);"
510
+ )
511
+ if not endpoint or endpoint[0][0] is None:
512
+ return False
513
+
514
+ return feature_enabled[0][0]
515
+ except Exception as e:
516
+ raise Exception(f"Unable to determine if direct access is enabled. Details error: {e}") from e
517
+
518
+
519
+ def is_account_flag_set(self, flag: str) -> bool:
520
+ results = self._exec(
521
+ f"SHOW PARAMETERS LIKE '%{flag}%' IN ACCOUNT;"
522
+ )
523
+ if not results:
524
+ return False
525
+ return results[0]["value"] == "true"
526
+
527
+ #--------------------------------------------------
528
+ # Databases
529
+ #--------------------------------------------------
530
+
531
+ def get_database(self, database: str):
532
+ try:
533
+ results = self._exec(
534
+ f"call {APP_NAME}.api.get_database('{database}');"
535
+ )
536
+ except Exception as e:
537
+ messages = collect_error_messages(e)
538
+ if any("database does not exist" in msg for msg in messages):
539
+ return None
540
+ raise e
541
+
542
+ if not results:
543
+ return None
544
+ db = results[0]
545
+ if not db:
546
+ return None
547
+ return {
548
+ "id": db["ID"],
549
+ "name": db["NAME"],
550
+ "created_by": db["CREATED_BY"],
551
+ "created_on": db["CREATED_ON"],
552
+ "deleted_by": db["DELETED_BY"],
553
+ "deleted_on": db["DELETED_ON"],
554
+ "state": db["STATE"],
555
+ }
556
+
557
+ def get_installed_packages(self, database: str) -> Dict | None:
558
+ query = f"call {APP_NAME}.api.get_installed_package_versions('{database}');"
559
+ try:
560
+ results = self._exec(query)
561
+ except Exception as e:
562
+ messages = collect_error_messages(e)
563
+ if any("database does not exist" in msg for msg in messages):
564
+ return None
565
+ # fallback to None for old sql-lib versions
566
+ if any("unknown user-defined function" in msg for msg in messages):
567
+ return None
568
+ raise e
569
+
570
+ if not results:
571
+ return None
572
+
573
+ row = results[0]
574
+ if not row:
575
+ return None
576
+
577
+ return safe_json_loads(row["PACKAGE_VERSIONS"])
578
+
579
+ #--------------------------------------------------
580
+ # Engines
581
+ #--------------------------------------------------
582
+
583
+ def _prepare_engine_params(
584
+ self,
585
+ name: str | None,
586
+ size: str | None,
587
+ use_default_size: bool = False
588
+ ) -> tuple[str, str | None]:
589
+ """
590
+ Prepare engine parameters by resolving and validating name and size.
591
+
592
+ Args:
593
+ name: Engine name (None to use default)
594
+ size: Engine size (None to use config or default)
595
+ use_default_size: If True and size is None, use get_default_engine_size()
596
+
597
+ Returns:
598
+ Tuple of (engine_name, engine_size)
599
+
600
+ Raises:
601
+ EngineNameValidationException: If engine name is invalid
602
+ Exception: If engine size is invalid
603
+ """
604
+ from relationalai.tools.cli_helpers import validate_engine_name
605
+
606
+ engine_name = name or self.get_default_engine_name()
607
+
608
+ # Resolve engine size
609
+ if size:
610
+ engine_size = size
611
+ else:
612
+ if use_default_size:
613
+ engine_size = self.config.get_default_engine_size()
614
+ else:
615
+ engine_size = self.config.get("engine_size", None)
616
+
617
+ # Validate engine size
618
+ if engine_size:
619
+ is_size_valid, sizes = self._engines.validate_engine_size(engine_size)
620
+ if not is_size_valid:
621
+ error_msg = f"Invalid engine size '{engine_size}'. Valid sizes are: {', '.join(sizes)}"
622
+ if use_default_size:
623
+ error_msg = f"Invalid engine size in config: '{engine_size}'. Valid sizes are: {', '.join(sizes)}"
624
+ raise Exception(error_msg)
625
+
626
+ # Validate engine name
627
+ is_name_valid, _ = validate_engine_name(engine_name)
628
+ if not is_name_valid:
629
+ raise EngineNameValidationException(engine_name)
630
+
631
+ return engine_name, engine_size
632
+
633
+ def _get_state_handler(self, state: str | None, handlers: list[EngineStateHandler]) -> EngineStateHandler:
634
+ """Find the appropriate state handler for the given state."""
635
+ for handler in handlers:
636
+ if handler.handles_state(state):
637
+ return handler
638
+ # Fallback to missing engine handler if no match
639
+ return handlers[-1] # Last handler should be MissingEngineHandler
640
+
641
+ def _process_engine_state(
642
+ self,
643
+ engine: EngineState | Dict[str, Any] | None,
644
+ context: EngineContext,
645
+ handlers: list[EngineStateHandler],
646
+ set_active_on_success: bool = False
647
+ ) -> EngineState | Dict[str, Any] | None:
648
+ """
649
+ Process engine state using appropriate state handler.
650
+
651
+ Args:
652
+ engine: Current engine state (or None if missing)
653
+ context: Engine context for state handling
654
+ handlers: List of state handlers to use (sync or async)
655
+ set_active_on_success: If True, set engine as active when handler returns engine
656
+
657
+ Returns:
658
+ Engine state after processing, or None if engine needs to be created
659
+ """
660
+ # Find and execute appropriate state handler
661
+ state = engine["state"] if engine else None
662
+ handler = self._get_state_handler(state, handlers)
663
+ engine = handler.handle(engine, context, self)
664
+
665
+ # If handler returned None and we didn't start with None state, engine needs to be created
666
+ # (e.g., GONE state deleted the engine, so we need to create a new one)
667
+ if not engine and state is not None:
668
+ handler = self._get_state_handler(None, handlers)
669
+ handler.handle(None, context, self)
670
+ elif set_active_on_success:
671
+ # Cast to EngineState for type safety (handlers return EngineDict which is compatible)
672
+ self._set_active_engine(cast(EngineState, engine))
673
+
674
+ return engine
675
+
676
+ def _handle_engine_creation_errors(self, error: Exception, engine_name: str, preserve_rai_exception: bool = False) -> None:
677
+ """
678
+ Handle errors during engine creation using error handlers.
679
+
680
+ Args:
681
+ error: The exception that occurred
682
+ engine_name: Name of the engine being created
683
+ preserve_rai_exception: If True, re-raise RAIException without wrapping
684
+
685
+ Raises:
686
+ RAIException: If preserve_rai_exception is True and error is RAIException
687
+ EngineProvisioningFailed: If error is not handled by error handlers
688
+ """
689
+ # Preserve RAIException passthrough if requested (for async mode)
690
+ if preserve_rai_exception and isinstance(error, RAIException):
691
+ raise error
692
+
693
+ # Check if this is a known error type that should be handled by error handlers
694
+ message = str(error).lower()
695
+ handled = False
696
+ # Engine creation isn't tied to a specific SQL ExecContext; pass a context that
697
+ # disables use_index retry behavior (and any future ctx-dependent handlers).
698
+ ctx = ExecContext(code="", help=True, skip_engine_db_error_retry=True)
699
+ for handler in self._error_handlers:
700
+ if handler.matches(error, message, ctx, self):
701
+ handler.handle(error, ctx, self)
702
+ handled = True
703
+ break # Handler raised exception, we're done
704
+
705
+ # If not handled by error handlers, wrap in EngineProvisioningFailed
706
+ if not handled:
707
+ raise EngineProvisioningFailed(engine_name, error) from error
708
+
709
+ def get_engine_sizes(self, cloud_provider: str|None=None):
710
+ return self._engines.get_engine_sizes(cloud_provider=cloud_provider)
711
+
712
+ def list_engines(
713
+ self,
714
+ state: str | None = None,
715
+ name: str | None = None,
716
+ type: str | None = None,
717
+ size: str | None = None,
718
+ created_by: str | None = None,
719
+ ):
720
+ return self._engines.list_engines(
721
+ state=state,
722
+ name=name,
723
+ type=type,
724
+ size=size,
725
+ created_by=created_by,
726
+ )
727
+
728
+ def get_engine(self, name: str, type: str):
729
+ return self._engines.get_engine(name, type)
730
+
731
+ def get_default_engine_name(self) -> str:
732
+ if self.config.get("engine_name", None) is not None:
733
+ profile = self.config.profile
734
+ raise InvalidAliasError(f"""
735
+ 'engine_name' is not a valid config option.
736
+ If you meant to use a specific engine, use 'engine' instead.
737
+ Otherwise, remove it from your '{profile}' configuration profile.
738
+ """)
739
+ engine = self.config.get("engine", None)
740
+ if not engine and self.config.get("user", None):
741
+ engine = _sanitize_user_name(str(self.config.get("user")))
742
+ if not engine:
743
+ engine = self.get_user_based_engine_name()
744
+ self.config.set("engine", engine)
745
+ return engine
746
+
747
+ def is_valid_engine_state(self, name:str):
748
+ return name in VALID_ENGINE_STATES
749
+
750
+ # Can be overridden by subclasses (e.g. DirectAccessResources)
751
+ def _create_engine(
752
+ self,
753
+ name: str,
754
+ type: str = EngineType.LOGIC,
755
+ size: str | None = None,
756
+ auto_suspend_mins: int | None= None,
757
+ is_async: bool = False,
758
+ headers: Dict | None = None,
759
+ settings: Dict[str, Any] | None = None,
760
+ ):
761
+ return self._engines._create_engine(
762
+ name=name,
763
+ type=type,
764
+ size=size,
765
+ auto_suspend_mins=auto_suspend_mins,
766
+ is_async=is_async,
767
+ headers=headers,
768
+ settings=settings,
769
+ )
770
+
771
+ def create_engine(
772
+ self,
773
+ name: str,
774
+ type: str | None = None,
775
+ size: str | None = None,
776
+ auto_suspend_mins: int | None = None,
777
+ headers: Dict | None = None,
778
+ settings: Dict[str, Any] | None = None,
779
+ ):
780
+ if type is None:
781
+ type = EngineType.LOGIC
782
+ # Route through _create_engine so subclasses (e.g. DirectAccessResources)
783
+ # can override engine creation behavior.
784
+ return self._create_engine(
785
+ name=name,
786
+ type=type,
787
+ size=size,
788
+ auto_suspend_mins=auto_suspend_mins,
789
+ is_async=False,
790
+ headers=headers,
791
+ settings=settings,
792
+ )
793
+
794
+ def create_engine_async(
795
+ self,
796
+ name: str,
797
+ type: str = EngineType.LOGIC,
798
+ size: str | None = None,
799
+ auto_suspend_mins: int | None = None,
800
+ ):
801
+ # Route through _create_engine so subclasses (e.g. DirectAccessResources)
802
+ # can override async engine creation behavior.
803
+ return self._create_engine(
804
+ name=name,
805
+ type=type,
806
+ size=size,
807
+ auto_suspend_mins=auto_suspend_mins,
808
+ is_async=True,
809
+ )
810
+
811
+ def delete_engine(self, name: str, type: str):
812
+ return self._engines.delete_engine(name, type)
813
+
814
+ def suspend_engine(self, name: str, type: str | None = None):
815
+ return self._engines.suspend_engine(name, type)
816
+
817
+ def resume_engine(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict:
818
+ return self._engines.resume_engine(name, type=type, headers=headers)
819
+
820
+ def resume_engine_async(self, name: str, type: str | None = None, headers: Dict | None = None) -> Dict:
821
+ return self._engines.resume_engine_async(name, type=type, headers=headers)
822
+
823
+ def alter_engine_pool(self, size:str|None=None, mins:int|None=None, maxs:int|None=None):
824
+ """Alter engine pool node limits for Snowflake."""
825
+ return self._engines.alter_engine_pool(size=size, mins=mins, maxs=maxs)
826
+
827
+ #--------------------------------------------------
828
+ # Graphs
829
+ #--------------------------------------------------
830
+
831
+ def list_graphs(self) -> List[AvailableModel]:
832
+ with debugging.span("list_models"):
833
+ query = textwrap.dedent(f"""
834
+ SELECT NAME, ID, CREATED_BY, CREATED_ON, STATE, DELETED_BY, DELETED_ON
835
+ FROM {APP_NAME}.api.databases
836
+ WHERE state <> 'DELETED'
837
+ ORDER BY NAME ASC;
838
+ """)
839
+ results = self._exec(query)
840
+ if not results:
841
+ return []
842
+ return [
843
+ {
844
+ "name": row["NAME"],
845
+ "id": row["ID"],
846
+ "created_by": row["CREATED_BY"],
847
+ "created_on": row["CREATED_ON"],
848
+ "state": row["STATE"],
849
+ "deleted_by": row["DELETED_BY"],
850
+ "deleted_on": row["DELETED_ON"],
851
+ }
852
+ for row in results
853
+ ]
854
+
855
+ def get_graph(self, name: str):
856
+ res = self.get_database(name)
857
+ if res and res.get("state") != "DELETED":
858
+ return res
859
+
860
+ def create_graph(self, name: str):
861
+ with debugging.span("create_model", name=name):
862
+ self._exec(f"call {APP_NAME}.api.create_database('{name}', false, {debugging.gen_current_propagation_headers()});")
863
+
864
+ def delete_graph(self, name:str, force=False, language:str="rel"):
865
+ prop_hdrs = debugging.gen_current_propagation_headers()
866
+ if self.config.get("use_graph_index", USE_GRAPH_INDEX):
867
+ keep_database = not force and self.config.get("reuse_model", True)
868
+ with debugging.span("release_index", name=name, keep_database=keep_database, language=language):
869
+ #TODO add headers to release_index
870
+ response = self._exec(f"call {APP_NAME}.api.release_index('{name}', OBJECT_CONSTRUCT('keep_database', {keep_database}, 'language', '{language}', 'user_agent', '{get_pyrel_version(self.generation)}'));")
871
+ if response:
872
+ result = next(iter(response))
873
+ obj = json.loads(result["RELEASE_INDEX"])
874
+ error = obj.get('error', None)
875
+ if error and "Model database not found" not in error:
876
+ raise Exception(f"Error releasing index: {error}")
877
+ else:
878
+ raise Exception("There was no response from the release index call.")
879
+ else:
880
+ with debugging.span("delete_model", name=name):
881
+ self._exec(f"call {APP_NAME}.api.delete_database('{name}', false, {prop_hdrs});")
882
+
883
+ def clone_graph(self, target_name:str, source_name:str, nowait_durable=True, force=False):
884
+ if force and self.get_graph(target_name):
885
+ self.delete_graph(target_name)
886
+ with debugging.span("clone_model", target_name=target_name, source_name=source_name):
887
+ # not a mistake: the clone_database argument order is indeed target then source:
888
+ headers = debugging.gen_current_propagation_headers()
889
+ self._exec(f"call {APP_NAME}.api.clone_database('{target_name}', '{source_name}', {nowait_durable}, {headers});")
890
+
891
+ def _poll_use_index(
892
+ self,
893
+ app_name: str,
894
+ sources: Iterable[str],
895
+ model: str,
896
+ engine_name: str,
897
+ engine_size: str | None = None,
898
+ program_span_id: str | None = None,
899
+ headers: Dict | None = None,
900
+ ) -> None:
901
+ """
902
+ Poll use_index to prepare indices for the given sources.
903
+
904
+ This is an optional interface method. Base Resources provides a no-op implementation.
905
+ UseIndexResources and DirectAccessResources override this to provide actual polling.
906
+
907
+ Returns:
908
+ None for base implementation. Child classes may return poller results.
909
+ """
910
+ return None
911
+
912
+ def maybe_poll_use_index(
913
+ self,
914
+ app_name: str,
915
+ sources: Iterable[str],
916
+ model: str,
917
+ engine_name: str,
918
+ engine_size: str | None = None,
919
+ program_span_id: str | None = None,
920
+ headers: Dict | None = None,
921
+ ) -> None:
922
+ """
923
+ Only call _poll_use_index if there are sources to process.
924
+
925
+ This is an optional interface method. Base Resources provides a no-op implementation.
926
+ UseIndexResources and DirectAccessResources override this to provide actual polling with caching.
927
+
928
+ Returns:
929
+ None for base implementation. Child classes may return poller results.
930
+ """
931
+ return None
932
+
933
+ #--------------------------------------------------
934
+ # Models
935
+ #--------------------------------------------------
936
+
937
+ def list_models(self, database: str, engine: str):
938
+ pass
939
+
940
+ def create_models(self, database: str, engine: str | None, models:List[Tuple[str, str]]) -> List[Any]:
941
+ rel_code = self.create_models_code(models)
942
+ self.exec_raw(database, engine, rel_code, readonly=False)
943
+ # TODO: handle SPCS errors once they're figured out
944
+ return []
945
+
946
+ def delete_model(self, database:str, engine:str | None, name:str):
947
+ self.exec_raw(database, engine, f"def delete[:rel, :catalog, :model, \"{name}\"]: rel[:catalog, :model, \"{name}\"]", readonly=False)
948
+
949
+ def create_models_code(self, models:List[Tuple[str, str]]) -> str:
950
+ lines = []
951
+ for (name, code) in models:
952
+ name = name.replace("\"", "\\\"")
953
+ assert "\"\"\"\"\"\"\"" not in code, "Code literals must use fewer than 7 quotes."
954
+
955
+ lines.append(textwrap.dedent(f"""
956
+ def delete[:rel, :catalog, :model, "{name}"]: rel[:catalog, :model, "{name}"]
957
+ def insert[:rel, :catalog, :model, "{name}"]: raw\"\"\"\"\"\"\"
958
+ """) + code + "\n\"\"\"\"\"\"\"")
959
+ rel_code = "\n\n".join(lines)
960
+ return rel_code
961
+
962
+ #--------------------------------------------------
963
+ # Exports
964
+ #--------------------------------------------------
965
+
966
+ def list_exports(self, database: str, engine: str):
967
+ return []
968
+
969
+
970
+ def get_export_code(self, params: ExportParams, all_installs):
971
+ sql_inputs = ", ".join([f"{name} {type_to_sql(type)}" for (name, _, type) in params.inputs])
972
+ input_names = [name for (name, *_) in params.inputs]
973
+ has_return_hint = params.out_fields and isinstance(params.out_fields[0], tuple)
974
+ if has_return_hint:
975
+ sql_out = ", ".join([f"\"{name}\" {type_to_sql(type)}" for (name, type) in params.out_fields])
976
+ sql_out_names = ", ".join([f"('{name}', '{type_to_sql(type)}')" for (ix, (name, type)) in enumerate(params.out_fields)])
977
+ py_outs = ", ".join([f"StructField(\"{name}\", {type_to_snowpark(type)})" for (name, type) in params.out_fields])
978
+ else:
979
+ sql_out = ""
980
+ sql_out_names = ", ".join([f"'{name}'" for name in params.out_fields])
981
+ py_outs = ", ".join([f"StructField(\"{name}\", {type_to_snowpark(str)})" for name in params.out_fields])
982
+ py_inputs = ", ".join([name for (name, *_) in params.inputs])
983
+ safe_rel = escape_for_f_string(params.code).strip()
984
+ clean_inputs = []
985
+ for (name, var, type) in params.inputs:
986
+ if type is str:
987
+ clean_inputs.append(f"{name} = '\"' + escape({name}) + '\"'")
988
+ # Replace `var` with `name` and keep the following non-word character unchanged
989
+ pattern = re.compile(re.escape(var) + r'(\W)')
990
+ value = format_sproc_name(name, type)
991
+ safe_rel = re.sub(pattern, rf"{{{value}}}\1", safe_rel)
992
+ if py_inputs:
993
+ py_inputs = f", {py_inputs}"
994
+ clean_inputs = ("\n").join(clean_inputs)
995
+ file = "export_procedure.py.jinja"
996
+ with importlib.resources.open_text(
997
+ "relationalai.clients.resources.snowflake", file
998
+ ) as f:
999
+ template = f.read()
1000
+ def quote(s: str, f = False) -> str:
1001
+ return '"' + s + '"' if not f else 'f"' + s + '"'
1002
+
1003
+ wait_for_stream_sync = self.config.get("wait_for_stream_sync", WAIT_FOR_STREAM_SYNC)
1004
+ # 1. Check the sources for staled sources
1005
+ # 2. Get the object references for the sources
1006
+ # TODO: this could be optimized to do it in the run time of the stored procedure
1007
+ # instead of doing it here. It will make it more reliable when sources are
1008
+ # modified after the stored procedure is created.
1009
+ checked_sources = self._check_source_updates(self.sources)
1010
+ source_obj_references = self._get_source_references(checked_sources)
1011
+
1012
+ # Escape double quotes in the source object references
1013
+ escaped_source_obj_references = [source.replace('"', '\\"') for source in source_obj_references]
1014
+ escaped_proc_database = params.proc_database.replace('"', '\\"')
1015
+
1016
+ normalized_func_name = IdentityParser(params.func_name).identity
1017
+ assert normalized_func_name is not None, "Function name must be set"
1018
+ skip_invalid_data = params.skip_invalid_data
1019
+ python_code = process_jinja_template(
1020
+ template,
1021
+ func_name=quote(normalized_func_name),
1022
+ database=quote(params.root_database),
1023
+ proc_database=quote(escaped_proc_database),
1024
+ engine=quote(params.engine),
1025
+ rel_code=quote(safe_rel, f=True),
1026
+ APP_NAME=quote(APP_NAME),
1027
+ input_names=input_names,
1028
+ outputs=sql_out,
1029
+ sql_out_names=sql_out_names,
1030
+ clean_inputs=clean_inputs,
1031
+ py_inputs=py_inputs,
1032
+ py_outs=py_outs,
1033
+ skip_invalid_data=skip_invalid_data,
1034
+ source_references=", ".join(escaped_source_obj_references),
1035
+ install_code=all_installs.replace("\\", "\\\\").replace("\n", "\\n"),
1036
+ has_return_hint=has_return_hint,
1037
+ wait_for_stream_sync=wait_for_stream_sync,
1038
+ ).strip()
1039
+ return_clause = f"TABLE({sql_out})" if sql_out else "STRING"
1040
+ destination_input = "" if sql_out else "save_as_table STRING DEFAULT NULL,"
1041
+ module_name = sanitize_module_name(normalized_func_name)
1042
+ stage = f"@{self.get_app_name()}.app_state.stored_proc_code_stage"
1043
+ file_loc = f"{stage}/{module_name}.py"
1044
+ python_code = python_code.replace(APP_NAME, self.get_app_name())
1045
+
1046
+ hash = hashlib.sha256()
1047
+ hash.update(python_code.encode('utf-8'))
1048
+ code_hash = hash.hexdigest()
1049
+ print(code_hash)
1050
+
1051
+ sql_code = textwrap.dedent(f"""
1052
+ CREATE OR REPLACE PROCEDURE {normalized_func_name}({sql_inputs}{sql_inputs and ',' or ''} {destination_input} engine STRING DEFAULT NULL)
1053
+ RETURNS {return_clause}
1054
+ LANGUAGE PYTHON
1055
+ RUNTIME_VERSION = '3.10'
1056
+ IMPORTS = ('{file_loc}')
1057
+ PACKAGES = ('snowflake-snowpark-python')
1058
+ HANDLER = 'checked_handle'
1059
+ EXECUTE AS CALLER
1060
+ AS
1061
+ $$
1062
+ import {module_name}
1063
+ import inspect, hashlib, os, sys
1064
+ def checked_handle(*args, **kwargs):
1065
+ import_dir = sys._xoptions["snowflake_import_directory"]
1066
+ wheel_path = os.path.join(import_dir, '{module_name}.py')
1067
+ h = hashlib.sha256()
1068
+ with open(wheel_path, 'rb') as f:
1069
+ for chunk in iter(lambda: f.read(1<<20), b''):
1070
+ h.update(chunk)
1071
+ code_hash = h.hexdigest()
1072
+ if code_hash != '{code_hash}':
1073
+ raise RuntimeError("Code hash mismatch. The code has been modified since it was uploaded.")
1074
+ # Call the handle function with the provided arguments
1075
+ return {module_name}.handle(*args, **kwargs)
1076
+
1077
+ $$;
1078
+ """)
1079
+ # print(f"\n--- python---\n{python_code}\n--- end python---\n")
1080
+ # This check helps catch invalid code early and for dry runs:
1081
+ try:
1082
+ ast.parse(python_code)
1083
+ except SyntaxError:
1084
+ raise ValueError(f"Internal error: invalid Python code generated:\n{python_code}")
1085
+ return (sql_code, python_code, file_loc)
1086
+
1087
+ def get_sproc_models(self, params: ExportParams):
1088
+ if self._sproc_models is not None:
1089
+ return self._sproc_models
1090
+
1091
+ with debugging.span("get_sproc_models"):
1092
+ code = """
1093
+ def output(name, model):
1094
+ rel(:catalog, :model, name, model)
1095
+ and not starts_with(name, "rel/")
1096
+ and not starts_with(name, "pkg/rel")
1097
+ and not starts_with(name, "pkg/std")
1098
+ and starts_with(name, "pkg/")
1099
+ """
1100
+ res = self.exec_raw(params.model_database, params.engine, code, readonly=True, nowait_durable=True)
1101
+ df, errors = result_helpers.format_results(res, None, ["name", "model"])
1102
+ models = []
1103
+ for row in df.itertuples():
1104
+ models.append((row.name, row.model))
1105
+ self._sproc_models = models
1106
+ return models
1107
+
1108
+ def create_export(self, params: ExportParams):
1109
+ with debugging.span("create_export") as span:
1110
+ if params.dry_run:
1111
+ (sql_code, python_code, file_loc) = self.get_export_code(params, params.install_code)
1112
+ span["sql"] = sql_code
1113
+ return
1114
+
1115
+ start = time.perf_counter()
1116
+ use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
1117
+ # for the non graph index case we need to create the cloned proc database
1118
+ if not use_graph_index:
1119
+ raise RAIException(
1120
+ "To ensure permissions are properly accounted for, stored procedures require using the graph index. "
1121
+ "Set use_graph_index=True in your config to proceed."
1122
+ )
1123
+
1124
+ models = self.get_sproc_models(params)
1125
+ lib_installs = self.create_models_code(models)
1126
+ all_installs = lib_installs + "\n\n" + params.install_code
1127
+
1128
+ (sql_code, python_code, file_loc) = self.get_export_code(params, all_installs)
1129
+
1130
+ span["sql"] = sql_code
1131
+ assert self._session
1132
+
1133
+ with debugging.span("upload_sproc_code"):
1134
+ code_bytes = python_code.encode('utf-8')
1135
+ code_stream = io.BytesIO(code_bytes)
1136
+ self._session.file.put_stream(code_stream, file_loc, auto_compress=False, overwrite=True)
1137
+
1138
+ with debugging.span("sql_install"):
1139
+ self._exec(sql_code)
1140
+
1141
+ debugging.time("export", time.perf_counter() - start, DataFrame(), code=sql_code.replace(APP_NAME, self.get_app_name()))
1142
+
1143
+
1144
+ def create_export_table(self, database: str, engine: str, table: str, relation: str, columns: Dict[str, str], code: str, refresh: str|None=None):
1145
+ print("Snowflake doesn't support creating export tables yet. Try creating the table manually first.")
1146
+ pass
1147
+
1148
+ def delete_export(self, database: str, engine: str, name: str):
1149
+ pass
1150
+
1151
+ #--------------------------------------------------
1152
+ # Imports
1153
+ #--------------------------------------------------
1154
+
1155
+
1156
+ def change_stream_status(self, stream_id: str, model:str, suspend: bool):
1157
+ if stream_id and model:
1158
+ if suspend:
1159
+ self._exec(f"CALL {APP_NAME}.api.suspend_data_stream('{stream_id}', '{model}');")
1160
+ else:
1161
+ self._exec(f"CALL {APP_NAME}.api.resume_data_stream('{stream_id}', '{model}');")
1162
+
1163
+ def change_imports_status(self, suspend: bool):
1164
+ if suspend:
1165
+ self._exec(f"CALL {APP_NAME}.app.suspend_cdc();")
1166
+ else:
1167
+ self._exec(f"CALL {APP_NAME}.app.resume_cdc();")
1168
+
1169
+ def get_imports_status(self) -> ImportsStatus|None:
1170
+ # NOTE: We expect there to only ever be one result?
1171
+ results = self._exec(f"CALL {APP_NAME}.app.cdc_status();")
1172
+ if results:
1173
+ result = next(iter(results))
1174
+ engine = result['CDC_ENGINE_NAME']
1175
+ engine_status = result['CDC_ENGINE_STATUS']
1176
+ engine_size = result['CDC_ENGINE_SIZE']
1177
+ task_status = result['CDC_TASK_STATUS']
1178
+ info = result['CDC_TASK_INFO']
1179
+ enabled = result['CDC_ENABLED']
1180
+ return {"engine": engine, "engine_size": engine_size, "engine_status": engine_status, "status": task_status, "enabled": enabled, "info": info }
1181
+ return None
1182
+
1183
+ def set_imports_engine_size(self, size:str):
1184
+ try:
1185
+ self._exec(f"CALL {APP_NAME}.app.alter_cdc_engine_size('{size}');")
1186
+ except Exception as e:
1187
+ raise e
1188
+
1189
+ def list_imports(
1190
+ self,
1191
+ id:str|None = None,
1192
+ name:str|None = None,
1193
+ model:str|None = None,
1194
+ status:str|None = None,
1195
+ creator:str|None = None,
1196
+ ) -> list[Import]:
1197
+ where = []
1198
+ if id and isinstance(id, str):
1199
+ where.append(f"LOWER(ID) = '{id.lower()}'")
1200
+ if name and isinstance(name, str):
1201
+ where.append(f"LOWER(FQ_OBJECT_NAME) = '{name.lower()}'")
1202
+ if model and isinstance(model, str):
1203
+ where.append(f"LOWER(RAI_DATABASE) = '{model.lower()}'")
1204
+ if creator and isinstance(creator, str):
1205
+ where.append(f"LOWER(CREATED_BY) = '{creator.lower()}'")
1206
+ if status and isinstance(status, str):
1207
+ where.append(f"LOWER(batch_status) = '{status.lower()}'")
1208
+ where_clause = " AND ".join(where)
1209
+
1210
+ # This is roughly inspired by the native app code because we don't have a way to
1211
+ # get the status of multiple streams at once and doing them individually is way
1212
+ # too slow. We use window functions to get the status of the stream and the batch
1213
+ # details.
1214
+ statement = f"""
1215
+ SELECT
1216
+ ID,
1217
+ RAI_DATABASE,
1218
+ FQ_OBJECT_NAME,
1219
+ CREATED_AT,
1220
+ CREATED_BY,
1221
+ CASE
1222
+ WHEN nextBatch.quarantined > 0 THEN 'quarantined'
1223
+ ELSE nextBatch.status
1224
+ END as batch_status,
1225
+ nextBatch.processing_errors,
1226
+ nextBatch.batches
1227
+ FROM {APP_NAME}.api.data_streams as ds
1228
+ LEFT JOIN (
1229
+ SELECT DISTINCT
1230
+ data_stream_id,
1231
+ -- Get status from the progress record using window functions
1232
+ FIRST_VALUE(status) OVER (
1233
+ PARTITION BY data_stream_id
1234
+ ORDER BY
1235
+ CASE WHEN unloaded IS NOT NULL THEN 1 ELSE 0 END DESC,
1236
+ unloaded ASC
1237
+ ) as status,
1238
+ -- Get batch_details from the same record
1239
+ FIRST_VALUE(batch_details) OVER (
1240
+ PARTITION BY data_stream_id
1241
+ ORDER BY
1242
+ CASE WHEN unloaded IS NOT NULL THEN 1 ELSE 0 END DESC,
1243
+ unloaded ASC
1244
+ ) as batch_details,
1245
+ -- Aggregate the other fields
1246
+ FIRST_VALUE(processing_details:processingErrors) OVER (
1247
+ PARTITION BY data_stream_id
1248
+ ORDER BY
1249
+ CASE WHEN unloaded IS NOT NULL THEN 1 ELSE 0 END DESC,
1250
+ unloaded ASC
1251
+ ) as processing_errors,
1252
+ MIN(unloaded) OVER (PARTITION BY data_stream_id) as unloaded,
1253
+ COUNT(*) OVER (PARTITION BY data_stream_id) as batches,
1254
+ COUNT_IF(status = 'quarantined') OVER (PARTITION BY data_stream_id) as quarantined
1255
+ FROM {APP_NAME}.api.data_stream_batches
1256
+ ) nextBatch
1257
+ ON ds.id = nextBatch.data_stream_id
1258
+ {f"where {where_clause}" if where_clause else ""}
1259
+ ORDER BY FQ_OBJECT_NAME ASC;
1260
+ """
1261
+ results = self._exec(statement)
1262
+ items = []
1263
+ if results:
1264
+ for stream in results:
1265
+ (id, db, name, created_at, created_by, status, processing_errors, batches) = stream
1266
+ if status and isinstance(status, str):
1267
+ status = status.upper()
1268
+ if processing_errors:
1269
+ if status in ["QUARANTINED", "PENDING"]:
1270
+ start = processing_errors.rfind("Error")
1271
+ if start != -1:
1272
+ processing_errors = processing_errors[start:-1]
1273
+ else:
1274
+ processing_errors = None
1275
+ items.append(cast(Import, {
1276
+ "id": id,
1277
+ "model": db,
1278
+ "name": name,
1279
+ "created": created_at,
1280
+ "creator": created_by,
1281
+ "status": status.upper() if status else None,
1282
+ "errors": processing_errors if processing_errors != "[]" else None,
1283
+ "batches": f"{batches}" if batches else "",
1284
+ }))
1285
+ return items
1286
+
1287
+ def poll_imports(self, sources:List[str], model:str):
1288
+ source_set = self._create_source_set(sources)
1289
+ def check_imports():
1290
+ imports = [
1291
+ import_
1292
+ for import_ in self.list_imports(model=model)
1293
+ if import_["name"] in source_set
1294
+ ]
1295
+ # loop through printing status for each in the format (index): (name) - (status)
1296
+ statuses = [import_["status"] for import_ in imports]
1297
+ if all(status == "LOADED" for status in statuses):
1298
+ return True
1299
+ if any(status == "QUARANTINED" for status in statuses):
1300
+ failed_imports = [import_["name"] for import_ in imports if import_["status"] == "QUARANTINED"]
1301
+ raise RAIException("Imports failed:" + ", ".join(failed_imports)) from None
1302
+ # this check is necessary in case some of the tables are empty;
1303
+ # such tables may be synced even though their status is None:
1304
+ def synced(import_):
1305
+ if import_["status"] == "LOADED":
1306
+ return True
1307
+ if import_["status"] is None:
1308
+ import_full_status = self.get_import_stream(import_["name"], model)
1309
+ if import_full_status and import_full_status[0]["data_sync_status"] == "SYNCED":
1310
+ return True
1311
+ return False
1312
+ if all(synced(import_) for import_ in imports):
1313
+ return True
1314
+ poll_with_specified_overhead(check_imports, overhead_rate=0.1, max_delay=10)
1315
+
1316
+ def _create_source_set(self, sources: List[str]) -> set:
1317
+ return {
1318
+ source.upper() if not IdentityParser(source).has_double_quoted_identifier else IdentityParser(source).identity
1319
+ for source in sources
1320
+ }
1321
+
1322
+ def get_import_stream(self, name:str|None, model:str|None):
1323
+ results = self._exec(f"CALL {APP_NAME}.api.get_data_stream('{name}', '{model}');")
1324
+ if not results:
1325
+ return None
1326
+ return imports_to_dicts(results)
1327
+
1328
+ def create_import_stream(self, source:ImportSource, model:str, rate = 1, options: dict|None = None):
1329
+ assert isinstance(source, ImportSourceTable), "Snowflake integration only supports loading from SF Tables. Try loading your data as a table via the Snowflake interface first."
1330
+ object = source.fqn
1331
+
1332
+ # Parse only to the schema level
1333
+ schemaParser = IdentityParser(f"{source.database}.{source.schema}")
1334
+
1335
+ if object.lower() in [x["name"].lower() for x in self.list_imports(model=model)]:
1336
+ return
1337
+
1338
+ query = f"SHOW OBJECTS LIKE '{source.table}' IN {schemaParser.identity}"
1339
+
1340
+ info = self._exec(query)
1341
+ if not info:
1342
+ raise ValueError(f"Object {source.table} not found in schema {schemaParser.identity}")
1343
+ else:
1344
+ data = info[0]
1345
+ if not data:
1346
+ raise ValueError(f"Object {source.table} not found in {schemaParser.identity}")
1347
+ # (time, name, db_name, schema_name, kind, *rest)
1348
+ kind = data["kind"]
1349
+
1350
+ relation_name = to_fqn_relation_name(object)
1351
+
1352
+ command = f"""call {APP_NAME}.api.create_data_stream(
1353
+ {APP_NAME}.api.object_reference('{kind}', '{object}'),
1354
+ '{model}',
1355
+ '{relation_name}');"""
1356
+
1357
+ def create_stream(tracking_just_changed=False):
1358
+ try:
1359
+ self._exec(command)
1360
+ except Exception as e:
1361
+ messages = collect_error_messages(e)
1362
+ if any("ensure that change_tracking is enabled on the source object" in msg for msg in messages):
1363
+ if self.config.get("ensure_change_tracking", False) and not tracking_just_changed:
1364
+ try:
1365
+ self._exec(f"ALTER {kind} {object} SET CHANGE_TRACKING = TRUE;")
1366
+ create_stream(tracking_just_changed=True)
1367
+ except Exception:
1368
+ pass
1369
+ else:
1370
+ print("\n")
1371
+ exception = SnowflakeChangeTrackingNotEnabledException((object, kind))
1372
+ raise exception from None
1373
+ elif any("database does not exist" in msg for msg in messages):
1374
+ print("\n")
1375
+ raise ModelNotFoundException(model) from None
1376
+ raise e
1377
+
1378
+ create_stream()
1379
+
1380
+ def create_import_snapshot(self, source:ImportSource, model:str, options: dict|None = None):
1381
+ raise Exception("Snowflake integration doesn't support snapshot imports yet")
1382
+
1383
+ def delete_import(self, import_name:str, model:str, force = False):
1384
+ engine = self.get_default_engine_name()
1385
+ rel_name = to_fqn_relation_name(import_name)
1386
+ try:
1387
+ self._exec(f"""call {APP_NAME}.api.delete_data_stream(
1388
+ '{import_name}',
1389
+ '{model}'
1390
+ );""")
1391
+ except RAIException as err:
1392
+ if "streams do not exist" not in str(err) or not force:
1393
+ raise
1394
+
1395
+ # if force is true, we delete the leftover relation to free up the name (in case the user re-creates the stream)
1396
+ if force:
1397
+ self.exec_raw(model, engine, f"""
1398
+ declare ::{rel_name}
1399
+ def delete[:\"{rel_name}\"]: {{ {rel_name} }}
1400
+ """, readonly=False, bypass_index=True)
1401
+
1402
+ #--------------------------------------------------
1403
+ # Exec Async
1404
+ #--------------------------------------------------
1405
+
1406
+ def _check_exec_async_status(self, txn_id: str, headers: Dict | None = None):
1407
+ """Check whether the given transaction has completed."""
1408
+ if headers is None:
1409
+ headers = {}
1410
+
1411
+ with debugging.span("check_status"):
1412
+ response = self._exec(f"CALL {APP_NAME}.api.get_transaction('{txn_id}',{headers});")
1413
+ assert response, f"No results from get_transaction('{txn_id}')"
1414
+
1415
+ response_row = next(iter(response)).asDict()
1416
+ status: str = response_row['STATE']
1417
+
1418
+ # remove the transaction from the pending list if it's completed or aborted
1419
+ if status in ["COMPLETED", "ABORTED"]:
1420
+ if txn_id in self._pending_transactions:
1421
+ self._pending_transactions.remove(txn_id)
1422
+
1423
+ if status == "ABORTED":
1424
+ if response_row.get("ABORT_REASON", "") == TXN_ABORT_REASON_TIMEOUT:
1425
+ config_file_path = getattr(self.config, 'file_path', None)
1426
+ # todo: use the timeout returned alongside the transaction as soon as it's exposed
1427
+ timeout_mins = int(self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS) or DEFAULT_QUERY_TIMEOUT_MINS)
1428
+ raise QueryTimeoutExceededException(
1429
+ timeout_mins=timeout_mins,
1430
+ query_id=txn_id,
1431
+ config_file_path=config_file_path,
1432
+ )
1433
+ elif response_row.get("ABORT_REASON", "") == GUARDRAILS_ABORT_REASON:
1434
+ raise GuardRailsException()
1435
+
1436
+ # @TODO: Find some way to tunnel the ABORT_REASON out. Azure doesn't have this, but it's handy
1437
+ return status == "COMPLETED" or status == "ABORTED"
1438
+
1439
+
1440
+ def _list_exec_async_artifacts(self, txn_id: str, headers: Dict | None = None) -> Dict[str, Dict]:
1441
+ """Grab the list of artifacts produced in the transaction and the URLs to retrieve their contents."""
1442
+ if headers is None:
1443
+ headers = {}
1444
+ with debugging.span("list_results"):
1445
+ response = self._exec(
1446
+ f"CALL {APP_NAME}.api.get_own_transaction_artifacts('{txn_id}',{headers});"
1447
+ )
1448
+ assert response, f"No results from get_own_transaction_artifacts('{txn_id}')"
1449
+ return {row["FILENAME"]: row for row in response}
1450
+
1451
+ def _fetch_exec_async_artifacts(
1452
+ self, artifact_info: Dict[str, Dict[str, Any]]
1453
+ ) -> Dict[str, Any]:
1454
+ """Grab the contents of the given artifacts from SF in parallel using threads."""
1455
+
1456
+ with requests.Session() as session:
1457
+ def _fetch_data(name_info):
1458
+ filename, metadata = name_info
1459
+
1460
+ try:
1461
+ # Extract the presigned URL and encryption material from metadata
1462
+ url_key = self.get_url_key(metadata)
1463
+ presigned_url = metadata[url_key]
1464
+ encryption_material = metadata["ENCRYPTION_MATERIAL"]
1465
+
1466
+ response = get_with_retries(session, presigned_url, config=self.config)
1467
+ response.raise_for_status() # Throw if something goes wrong
1468
+
1469
+ decrypted = self._maybe_decrypt(response.content, encryption_material)
1470
+ return (filename, decrypted)
1471
+
1472
+ except requests.RequestException as e:
1473
+ raise scrub_exception(wrap_with_request_id(e))
1474
+
1475
+ # Create a list of tuples for the map function
1476
+ name_info_pairs = list(artifact_info.items())
1477
+
1478
+ with ThreadPoolExecutor(max_workers=5) as executor:
1479
+ results = executor.map(_fetch_data, name_info_pairs)
1480
+
1481
+ return {name: data for (name, data) in results}
1482
+
1483
+ def _maybe_decrypt(self, content: bytes, encryption_material: str) -> bytes:
1484
+ # Decrypt if encryption material is present
1485
+ if encryption_material:
1486
+ # if there's no padding, the initial file was empty
1487
+ if len(content) == 0:
1488
+ return b""
1489
+
1490
+ return decrypt_artifact(content, encryption_material)
1491
+
1492
+ # otherwise, return content directly
1493
+ return content
1494
+
1495
+ def _parse_exec_async_results(self, arrow_files: List[Tuple[str, bytes]]):
1496
+ """Mimics the logic in _parse_arrow_results of railib/api.py#L303 without requiring a wrapping multipart form."""
1497
+ results = []
1498
+
1499
+ for file_name, file_content in arrow_files:
1500
+ with pa.ipc.open_stream(file_content) as reader:
1501
+ schema = reader.schema
1502
+ batches = [batch for batch in reader]
1503
+ table = pa.Table.from_batches(batches=batches, schema=schema)
1504
+ results.append({"relationId": file_name, "table": table})
1505
+
1506
+ return results
1507
+
1508
+ def _download_results(
1509
+ self, artifact_info: Dict[str, Dict], txn_id: str, state: str
1510
+ ) -> TransactionAsyncResponse:
1511
+ with debugging.span("download_results"):
1512
+ # Fetch artifacts
1513
+ artifacts = self._fetch_exec_async_artifacts(artifact_info)
1514
+
1515
+ # Directly use meta_json as it is fetched
1516
+ meta_json_bytes = artifacts["metadata.json"]
1517
+
1518
+ # Decode the bytes and parse the JSON
1519
+ meta_json_str = meta_json_bytes.decode('utf-8')
1520
+ meta_json = json.loads(meta_json_str) # Parse the JSON string
1521
+
1522
+ # Use the metadata to map arrow files to the relations they contain
1523
+ try:
1524
+ arrow_files_to_relations = {
1525
+ artifact["filename"]: artifact["relationId"]
1526
+ for artifact in meta_json
1527
+ }
1528
+ except KeyError:
1529
+ # TODO: Remove this fallback mechanism later once several engine versions are updated
1530
+ arrow_files_to_relations = {
1531
+ f"{ix}.arrow": artifact["relationId"]
1532
+ for ix, artifact in enumerate(meta_json)
1533
+ }
1534
+
1535
+ # Hydrate the arrow files into tables
1536
+ results = self._parse_exec_async_results(
1537
+ [
1538
+ (arrow_files_to_relations[name], content)
1539
+ for name, content in artifacts.items()
1540
+ if name.endswith(".arrow")
1541
+ ]
1542
+ )
1543
+
1544
+ # Create and return the response
1545
+ rsp = TransactionAsyncResponse()
1546
+ rsp.transaction = {
1547
+ "id": txn_id,
1548
+ "state": state,
1549
+ "response_format_version": None,
1550
+ }
1551
+ rsp.metadata = meta_json
1552
+ rsp.problems = artifacts.get(
1553
+ "problems.json"
1554
+ ) # Safely access possible missing keys
1555
+ rsp.results = results
1556
+ return rsp
1557
+
1558
+ def get_transaction_problems(self, txn_id: str) -> List[Dict[str, Any]]:
1559
+ with debugging.span("get_own_transaction_problems"):
1560
+ response = self._exec(
1561
+ f"select * from table({APP_NAME}.api.get_own_transaction_problems('{txn_id}'));"
1562
+ )
1563
+ if not response:
1564
+ return []
1565
+ return response
1566
+
1567
+ def get_url_key(self, metadata) -> str:
1568
+ # In Azure, there is only one type of URL, which is used for both internal and
1569
+ # external access; always use that one
1570
+ if is_azure_url(metadata['PRESIGNED_URL']):
1571
+ return 'PRESIGNED_URL'
1572
+
1573
+ configured = self.config.get("download_url_type", None)
1574
+ if configured == "internal":
1575
+ return 'PRESIGNED_URL_AP'
1576
+ elif configured == "external":
1577
+ return "PRESIGNED_URL"
1578
+
1579
+ if is_container_runtime():
1580
+ return 'PRESIGNED_URL_AP'
1581
+
1582
+ return 'PRESIGNED_URL'
1583
+
1584
+ def _exec_rai_app(
1585
+ self,
1586
+ database: str,
1587
+ engine: str | None,
1588
+ raw_code: str,
1589
+ inputs: Dict,
1590
+ readonly=True,
1591
+ nowait_durable=False,
1592
+ request_headers: Dict | None = None,
1593
+ bypass_index=False,
1594
+ language: str = "rel",
1595
+ query_timeout_mins: int | None = None,
1596
+ ):
1597
+ """
1598
+ High-level method to execute RAI app stored procedures.
1599
+
1600
+ Builds and executes SQL to call the RAI app's exec_async_v2 stored procedure.
1601
+ This method handles the SQL string construction for two different formats:
1602
+ 1. New format (with graph index): Uses object payload with parameterized query
1603
+ 2. Legacy format: Uses positional parameters
1604
+
1605
+ The choice between formats depends on the use_graph_index configuration.
1606
+ The new format allows the stored procedure to hash the model and username
1607
+ to determine the database, while the legacy format uses the passed database directly.
1608
+
1609
+ This method is called by _exec_async_v2 to create transactions. It skips
1610
+ use_index retry logic (skip_engine_db_error_retry=True) because that
1611
+ is handled at a higher level by exec_raw/exec_lqp.
1612
+
1613
+ Args:
1614
+ database: Database/model name
1615
+ engine: Engine name (optional)
1616
+ raw_code: Code to execute (REL, LQP, or SQL)
1617
+ inputs: Input parameters for the query
1618
+ readonly: Whether the transaction is read-only
1619
+ nowait_durable: Whether to wait for durable writes
1620
+ request_headers: Optional HTTP headers
1621
+ bypass_index: Whether to bypass graph index setup
1622
+ language: Query language ("rel" or "lqp")
1623
+ query_timeout_mins: Optional query timeout in minutes
1624
+
1625
+ Returns:
1626
+ Response from the stored procedure call (transaction creation result)
1627
+
1628
+ Raises:
1629
+ Exception: If transaction creation fails
1630
+ """
1631
+ assert language == "rel" or language == "lqp", "Only 'rel' and 'lqp' languages are supported"
1632
+ if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
1633
+ query_timeout_mins = int(timeout_value)
1634
+ # Depending on the shape of the input, the behavior of exec_async_v2 changes.
1635
+ # When using the new format (with an object), the function retrieves the
1636
+ # 'rai' database by hashing the model and username. In contrast, the
1637
+ # current version directly uses the passed database value.
1638
+ # Therefore, we must use the original exec_async_v2 when not using the
1639
+ # graph index to ensure the correct database is utilized.
1640
+ use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
1641
+ if use_graph_index and not bypass_index:
1642
+ payload = {
1643
+ 'database': database,
1644
+ 'engine': engine,
1645
+ 'inputs': inputs,
1646
+ 'readonly': readonly,
1647
+ 'nowait_durable': nowait_durable,
1648
+ 'language': language,
1649
+ 'headers': request_headers
1650
+ }
1651
+ if query_timeout_mins is not None:
1652
+ payload["timeout_mins"] = query_timeout_mins
1653
+ sql_string = f"CALL {APP_NAME}.api.exec_async_v2(?, {payload});"
1654
+ else:
1655
+ if query_timeout_mins is not None:
1656
+ sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {query_timeout_mins}, {request_headers});"
1657
+ else:
1658
+ sql_string = f"CALL {APP_NAME}.api.exec_async_v2('{database}','{engine}', ?, {inputs}, {readonly}, {nowait_durable}, '{language}', {request_headers});"
1659
+ # Don't let exec setup GI on failure, exec_raw and exec_lqp will do that and add the correct headers.
1660
+ response = self._exec(
1661
+ sql_string,
1662
+ raw_code,
1663
+ skip_engine_db_error_retry=True,
1664
+ )
1665
+ if not response:
1666
+ raise Exception("Failed to create transaction")
1667
+ return response
1668
+
1669
+ def _create_v2_txn(
1670
+ self,
1671
+ database: str,
1672
+ engine: str | None,
1673
+ raw_code: str,
1674
+ inputs: Dict,
1675
+ headers: Dict[str, str],
1676
+ readonly: bool,
1677
+ nowait_durable: bool,
1678
+ bypass_index: bool,
1679
+ language: str,
1680
+ query_timeout_mins: int | None,
1681
+ ) -> TxnCreationResult:
1682
+ """
1683
+ Create a transaction and return the result.
1684
+
1685
+ This method handles calling the RAI app stored procedure to create a transaction
1686
+ and parses the response into a standardized TxnCreationResult format.
1687
+
1688
+ This method can be overridden by subclasses (e.g., DirectAccessResources)
1689
+ to use different transport mechanisms (HTTP instead of SQL).
1690
+
1691
+ Args:
1692
+ database: Database/model name
1693
+ engine: Engine name (optional)
1694
+ raw_code: Code to execute (REL, LQP, or SQL)
1695
+ inputs: Input parameters for the query
1696
+ headers: HTTP headers (must be prepared by caller)
1697
+ readonly: Whether the transaction is read-only
1698
+ nowait_durable: Whether to wait for durable writes
1699
+ bypass_index: Whether to bypass graph index setup
1700
+ language: Query language ("rel" or "lqp")
1701
+ query_timeout_mins: Optional query timeout in minutes
1702
+
1703
+ Returns:
1704
+ TxnCreationResult containing txn_id, state, and artifact_info
1705
+ """
1706
+ response = self._exec_rai_app(
1707
+ database=database,
1708
+ engine=engine,
1709
+ raw_code=raw_code,
1710
+ inputs=inputs,
1711
+ readonly=readonly,
1712
+ nowait_durable=nowait_durable,
1713
+ request_headers=headers,
1714
+ bypass_index=bypass_index,
1715
+ language=language,
1716
+ query_timeout_mins=query_timeout_mins,
1717
+ )
1718
+
1719
+ rows = list(iter(response))
1720
+
1721
+ # process the first row since txn_id and state are the same for all rows
1722
+ first_row = rows[0]
1723
+ txn_id = first_row['ID']
1724
+ state = first_row['STATE']
1725
+
1726
+ # Build artifact_info if transaction completed immediately (fast path)
1727
+ artifact_info: Dict[str, Dict] = {}
1728
+ if state in ["COMPLETED", "ABORTED"]:
1729
+ for row in rows:
1730
+ filename = row['FILENAME']
1731
+ artifact_info[filename] = row
1732
+
1733
+ return TxnCreationResult(txn_id=txn_id, state=state, artifact_info=artifact_info)
1734
+
1735
+ def _exec_async_v2(
1736
+ self,
1737
+ database: str,
1738
+ engine: str | None,
1739
+ raw_code: str,
1740
+ inputs: Dict | None = None,
1741
+ readonly=True,
1742
+ nowait_durable=False,
1743
+ headers: Dict | None = None,
1744
+ bypass_index=False,
1745
+ language: str = "rel",
1746
+ query_timeout_mins: int | None = None,
1747
+ gi_setup_skipped: bool = False,
1748
+ ):
1749
+ """
1750
+ High-level async execution method with transaction polling and artifact management.
1751
+
1752
+ This is the core method for executing queries asynchronously. It:
1753
+ 1. Creates a transaction by calling _create_v2_txn
1754
+ 2. Handles two execution paths:
1755
+ - Fast path: Transaction completes immediately (COMPLETED/ABORTED)
1756
+ - Slow path: Transaction is pending, requires polling until completion
1757
+ 3. Manages pending transactions list
1758
+ 4. Downloads and returns query results/artifacts
1759
+
1760
+ This method is called by _execute_code (base implementation), and calls the
1761
+ following methods that can be overridden by child classes (e.g.,
1762
+ DirectAccessResources uses HTTP instead):
1763
+ - _create_v2_txn
1764
+ - _check_exec_async_status
1765
+ - _list_exec_async_artifacts
1766
+ - _download_results
1767
+
1768
+ Args:
1769
+ database: Database/model name
1770
+ engine: Engine name (optional)
1771
+ raw_code: Code to execute (REL, LQP, or SQL)
1772
+ inputs: Input parameters for the query
1773
+ readonly: Whether the transaction is read-only
1774
+ nowait_durable: Whether to wait for durable writes
1775
+ headers: Optional HTTP headers
1776
+ bypass_index: Whether to bypass graph index setup
1777
+ language: Query language ("rel" or "lqp")
1778
+ query_timeout_mins: Optional query timeout in minutes
1779
+ gi_setup_skipped: Whether graph index setup was skipped (for retry logic)
1780
+
1781
+ Returns:
1782
+ Query results (downloaded artifacts)
1783
+ """
1784
+ if inputs is None:
1785
+ inputs = {}
1786
+ request_headers = debugging.add_current_propagation_headers(headers)
1787
+ query_attrs_dict = json.loads(request_headers.get("X-Query-Attributes", "{}"))
1788
+
1789
+ with debugging.span("transaction", **query_attrs_dict) as txn_span:
1790
+ txn_start_time = time.time()
1791
+ with debugging.span("create_v2", **query_attrs_dict) as create_span:
1792
+ # Prepare headers for transaction creation
1793
+ request_headers['user-agent'] = get_pyrel_version(self.generation)
1794
+ request_headers['gi_setup_skipped'] = str(gi_setup_skipped)
1795
+ request_headers['pyrel_program_id'] = debugging.get_program_span_id() or ""
1796
+
1797
+ # Create the transaction
1798
+ result = self._create_v2_txn(
1799
+ database=database,
1800
+ engine=engine,
1801
+ raw_code=raw_code,
1802
+ inputs=inputs,
1803
+ headers=request_headers,
1804
+ readonly=readonly,
1805
+ nowait_durable=nowait_durable,
1806
+ bypass_index=bypass_index,
1807
+ language=language,
1808
+ query_timeout_mins=query_timeout_mins,
1809
+ )
1810
+
1811
+ txn_id = result.txn_id
1812
+ state = result.state
1813
+
1814
+ txn_span["txn_id"] = txn_id
1815
+ create_span["txn_id"] = txn_id
1816
+ debugging.event("transaction_created", txn_span, txn_id=txn_id)
1817
+
1818
+ print_txn_progress = should_print_txn_progress(self.config)
1819
+
1820
+ # fast path: transaction already finished
1821
+ if state in ["COMPLETED", "ABORTED"]:
1822
+ txn_end_time = time.time()
1823
+ if txn_id in self._pending_transactions:
1824
+ self._pending_transactions.remove(txn_id)
1825
+
1826
+ artifact_info = result.artifact_info
1827
+
1828
+ txn_duration = txn_end_time - txn_start_time
1829
+ if print_txn_progress:
1830
+ print(
1831
+ query_complete_message(txn_id, txn_duration, status_header=True)
1832
+ )
1833
+
1834
+ # Slow path: transaction not done yet; start polling
1835
+ else:
1836
+ self._pending_transactions.append(txn_id)
1837
+ # Use the interactive poller for transaction status
1838
+ with debugging.span("wait", txn_id=txn_id):
1839
+ if print_txn_progress:
1840
+ poller = ExecTxnPoller(resource=self, txn_id=txn_id, headers=request_headers, txn_start_time=txn_start_time)
1841
+ poller.poll()
1842
+ else:
1843
+ poll_with_specified_overhead(
1844
+ lambda: self._check_exec_async_status(txn_id, headers=request_headers), 0.1
1845
+ )
1846
+ artifact_info = self._list_exec_async_artifacts(txn_id, headers=request_headers)
1847
+
1848
+ with debugging.span("fetch"):
1849
+ return self._download_results(artifact_info, txn_id, state)
1850
+
1851
+ def get_user_based_engine_name(self):
1852
+ if not self._session:
1853
+ self._session = self.get_sf_session()
1854
+ user_table = self._session.sql("select current_user()").collect()
1855
+ user = user_table[0][0]
1856
+ assert isinstance(user, str), f"current_user() must return a string, not {type(user)}"
1857
+ return _sanitize_user_name(user)
1858
+
1859
+ def is_engine_ready(self, engine_name: str, type: str = EngineType.LOGIC):
1860
+ engine = self.get_engine(engine_name, type)
1861
+ return engine and engine["state"] == "READY"
1862
+
1863
+ def auto_create_engine(
1864
+ self,
1865
+ name: str | None = None,
1866
+ type: str = EngineType.LOGIC,
1867
+ size: str | None = None,
1868
+ headers: Dict | None = None,
1869
+ ):
1870
+ """Synchronously create/ensure an engine is ready, blocking until ready."""
1871
+ with debugging.span("auto_create_engine", active=self._active_engine) as span:
1872
+ active = self._get_active_engine()
1873
+ if active:
1874
+ return active
1875
+
1876
+ # Resolve and validate parameters
1877
+ name, size = self._prepare_engine_params(name, size)
1878
+
1879
+ try:
1880
+ # Get current engine state
1881
+ engine = self.get_engine(name, type)
1882
+ if engine:
1883
+ span.update(cast(dict, engine))
1884
+
1885
+ # Create context for state handling
1886
+ context = EngineContext(
1887
+ name=name,
1888
+ size=size,
1889
+ type=type,
1890
+ headers=headers,
1891
+ requested_size=size,
1892
+ span=span,
1893
+ )
1894
+
1895
+ # Process engine state using sync handlers
1896
+ self._process_engine_state(engine, context, self._sync_engine_state_handlers)
1897
+
1898
+ except Exception as e:
1899
+ self._handle_engine_creation_errors(e, name)
1900
+
1901
+ return name
1902
+
1903
+ def auto_create_engine_async(self, name: str | None = None, type: str | None = None):
1904
+ """Asynchronously create/ensure an engine, returns immediately."""
1905
+ if type is None:
1906
+ type = EngineType.LOGIC
1907
+ active = self._get_active_engine()
1908
+ if active and (active == name or name is None):
1909
+ return active
1910
+
1911
+ with Spinner(
1912
+ "Checking engine status",
1913
+ leading_newline=True,
1914
+ ) as spinner:
1915
+ with debugging.span("auto_create_engine_async", active=self._active_engine):
1916
+ # Resolve and validate parameters (use_default_size=True for async)
1917
+ name, size = self._prepare_engine_params(name, None, use_default_size=True)
1918
+
1919
+ try:
1920
+ # Get current engine state
1921
+ engine = self.get_engine(name, type)
1922
+
1923
+ # Create context for state handling
1924
+ context = EngineContext(
1925
+ name=name,
1926
+ size=size,
1927
+ type=type,
1928
+ headers=None,
1929
+ requested_size=None,
1930
+ spinner=spinner,
1931
+ )
1932
+
1933
+ # Process engine state using async handlers
1934
+ self._process_engine_state(engine, context, self._async_engine_state_handlers, set_active_on_success=True)
1935
+
1936
+ except Exception as e:
1937
+ spinner.update_messages({
1938
+ "finished_message": f"Failed to create engine {name}",
1939
+ })
1940
+ self._handle_engine_creation_errors(e, name, preserve_rai_exception=True)
1941
+
1942
+ return name
1943
+
1944
+ #--------------------------------------------------
1945
+ # Exec
1946
+ #--------------------------------------------------
1947
+
1948
+ def _execute_code(
1949
+ self,
1950
+ database: str,
1951
+ engine: str | None,
1952
+ raw_code: str,
1953
+ inputs: Dict | None,
1954
+ readonly: bool,
1955
+ nowait_durable: bool,
1956
+ headers: Dict | None,
1957
+ bypass_index: bool,
1958
+ language: str,
1959
+ query_timeout_mins: int | None,
1960
+ ) -> Any:
1961
+ """
1962
+ Template method for code execution - can be overridden by child classes.
1963
+
1964
+ This is a template method that provides a hook for child classes to add
1965
+ execution logic (like retry mechanisms). The base implementation simply
1966
+ calls _exec_async_v2 directly.
1967
+
1968
+ UseIndexResources overrides this method to use _exec_with_gi_retry, which
1969
+ adds automatic use_index polling on engine/database errors.
1970
+
1971
+ This method is called by exec_lqp() and exec_raw() to provide a single
1972
+ execution point that can be customized per resource class.
1973
+
1974
+ Args:
1975
+ database: Database/model name
1976
+ engine: Engine name (optional)
1977
+ raw_code: Code to execute (already processed/encoded)
1978
+ inputs: Input parameters for the query
1979
+ readonly: Whether the transaction is read-only
1980
+ nowait_durable: Whether to wait for durable writes
1981
+ headers: Optional HTTP headers
1982
+ bypass_index: Whether to bypass graph index setup
1983
+ language: Query language ("rel" or "lqp")
1984
+ query_timeout_mins: Optional query timeout in minutes
1985
+
1986
+ Returns:
1987
+ Query results
1988
+ """
1989
+ return self._exec_async_v2(
1990
+ database, engine, raw_code, inputs, readonly, nowait_durable,
1991
+ headers=headers, bypass_index=bypass_index, language=language,
1992
+ query_timeout_mins=query_timeout_mins, gi_setup_skipped=True,
1993
+ )
1994
+
1995
+ def exec_lqp(
1996
+ self,
1997
+ database: str,
1998
+ engine: str | None,
1999
+ raw_code: bytes,
2000
+ readonly=True,
2001
+ *,
2002
+ inputs: Dict | None = None,
2003
+ nowait_durable=False,
2004
+ headers: Dict | None = None,
2005
+ bypass_index=False,
2006
+ query_timeout_mins: int | None = None,
2007
+ ):
2008
+ """Execute LQP code."""
2009
+ raw_code_b64 = base64.b64encode(raw_code).decode("utf-8")
2010
+ return self._execute_code(
2011
+ database, engine, raw_code_b64, inputs, readonly, nowait_durable,
2012
+ headers, bypass_index, 'lqp', query_timeout_mins
2013
+ )
2014
+
2015
+ def exec_raw(
2016
+ self,
2017
+ database: str,
2018
+ engine: str | None,
2019
+ raw_code: str,
2020
+ readonly=True,
2021
+ *,
2022
+ inputs: Dict | None = None,
2023
+ nowait_durable=False,
2024
+ headers: Dict | None = None,
2025
+ bypass_index=False,
2026
+ query_timeout_mins: int | None = None,
2027
+ ):
2028
+ """Execute raw code."""
2029
+ raw_code = raw_code.replace("'", "\\'")
2030
+ return self._execute_code(
2031
+ database, engine, raw_code, inputs, readonly, nowait_durable,
2032
+ headers, bypass_index, 'rel', query_timeout_mins
2033
+ )
2034
+
2035
+
2036
+ def format_results(self, results, task:m.Task|None=None) -> Tuple[DataFrame, List[Any]]:
2037
+ return result_helpers.format_results(results, task)
2038
+
2039
+ #--------------------------------------------------
2040
+ # Exec format
2041
+ #--------------------------------------------------
2042
+
2043
+ def exec_format(
2044
+ self,
2045
+ database: str,
2046
+ engine: str,
2047
+ raw_code: str,
2048
+ cols: List[str],
2049
+ format: str,
2050
+ inputs: Dict | None = None,
2051
+ readonly=True,
2052
+ nowait_durable=False,
2053
+ skip_invalid_data=False,
2054
+ headers: Dict | None = None,
2055
+ query_timeout_mins: int | None = None,
2056
+ ):
2057
+ if inputs is None:
2058
+ inputs = {}
2059
+ if headers is None:
2060
+ headers = {}
2061
+ if 'user-agent' not in headers:
2062
+ headers['user-agent'] = get_pyrel_version(self.generation)
2063
+ if query_timeout_mins is None and (timeout_value := self.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
2064
+ query_timeout_mins = int(timeout_value)
2065
+ # TODO: add headers
2066
+ start = time.perf_counter()
2067
+ output_table = "out" + str(uuid.uuid4()).replace("-", "_")
2068
+ temp_table = f"temp_{output_table}"
2069
+ use_graph_index = self.config.get("use_graph_index", USE_GRAPH_INDEX)
2070
+ txn_id = None
2071
+ rejected_rows = None
2072
+ col_names_map = None
2073
+ artifacts = None
2074
+ assert self._session
2075
+ temp = self._session.createDataFrame([], StructType([StructField(name, StringType()) for name in cols]))
2076
+ with debugging.span("transaction") as txn_span:
2077
+ try:
2078
+ # In the graph index case we need to use the new exec_into_table proc as it obfuscates the db name
2079
+ with debugging.span("exec_format"):
2080
+ if use_graph_index:
2081
+ # we do not provide a default value for query_timeout_mins so that we can control the default on app level
2082
+ if query_timeout_mins is not None:
2083
+ res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, NULL, ?, {headers}, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data, query_timeout_mins])
2084
+ else:
2085
+ res = self._exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, NULL, ?, {headers}, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
2086
+ txn_id = json.loads(res[0]["EXEC_INTO_TABLE"])["rai_transaction_id"]
2087
+ rejected_rows = json.loads(res[0]["EXEC_INTO_TABLE"]).get("rejected_rows", [])
2088
+ rejected_rows_count = json.loads(res[0]["EXEC_INTO_TABLE"]).get("rejected_rows_count", 0)
2089
+ else:
2090
+ if query_timeout_mins is not None:
2091
+ res = self._exec(f"call {APP_NAME}.api.exec_into(?, ?, ?, ?, ?, {inputs}, ?, {headers}, ?, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data, query_timeout_mins])
2092
+ else:
2093
+ res = self._exec(f"call {APP_NAME}.api.exec_into(?, ?, ?, ?, ?, {inputs}, ?, {headers}, ?);", [database, engine, raw_code, output_table, readonly, nowait_durable, skip_invalid_data])
2094
+ txn_id = json.loads(res[0]["EXEC_INTO"])["rai_transaction_id"]
2095
+ rejected_rows = json.loads(res[0]["EXEC_INTO"]).get("rejected_rows", [])
2096
+ rejected_rows_count = json.loads(res[0]["EXEC_INTO"]).get("rejected_rows_count", 0)
2097
+ debugging.event("transaction_created", txn_span, txn_id=txn_id)
2098
+ debugging.time("exec_format", time.perf_counter() - start, DataFrame())
2099
+
2100
+ with debugging.span("temp_table_swap", txn_id=txn_id):
2101
+ out_sample = self._exec(f"select * from {APP_NAME}.results.{output_table} limit 1;")
2102
+ if out_sample:
2103
+ keys = set([k.lower() for k in out_sample[0].as_dict().keys()])
2104
+ col_names_map = {}
2105
+ for ix, name in enumerate(cols):
2106
+ col_key = f"col{ix:03}"
2107
+ if col_key in keys:
2108
+ col_names_map[col_key] = IdentityParser(name).identity
2109
+ else:
2110
+ col_names_map[col_key] = name
2111
+
2112
+ names = ", ".join([
2113
+ f"{col_key} as {alias}" if col_key in keys else f"NULL as {alias}"
2114
+ for col_key, alias in col_names_map.items()
2115
+ ])
2116
+ self._exec(f"CREATE TEMPORARY TABLE {APP_NAME}.results.{temp_table} AS SELECT {names} FROM {APP_NAME}.results.{output_table};")
2117
+ self._exec(f"call {APP_NAME}.api.drop_result_table(?)", [output_table])
2118
+ temp = cast(snowflake.snowpark.DataFrame, self._exec(f"select * from {APP_NAME}.results.{temp_table}", raw=True))
2119
+ if rejected_rows:
2120
+ debugging.warn(RowsDroppedFromTargetTableWarning(rejected_rows, rejected_rows_count, col_names_map))
2121
+ except Exception as e:
2122
+ messages = collect_error_messages(e)
2123
+ if any("no columns returned" in msg or "columns of results could not be determined" in msg for msg in messages):
2124
+ pass
2125
+ else:
2126
+ raise e
2127
+ if txn_id:
2128
+ artifact_info = self._list_exec_async_artifacts(txn_id)
2129
+ with debugging.span("fetch"):
2130
+ artifacts = self._download_results(artifact_info, txn_id, "ABORTED")
2131
+ return (temp, artifacts)
2132
+
2133
+ #--------------------------------------------------
2134
+ # Custom model types
2135
+ #--------------------------------------------------
2136
+
2137
+ def _get_ns(self, model:dsl.Graph):
2138
+ if model not in self._ns_cache:
2139
+ self._ns_cache[model] = _Snowflake(model)
2140
+ return self._ns_cache[model]
2141
+
2142
+ def to_model_type(self, model:dsl.Graph, name: str, source:str):
2143
+ parser = IdentityParser(source)
2144
+ if not parser.is_complete:
2145
+ raise SnowflakeInvalidSource(Errors.call_source(), source)
2146
+ ns = self._get_ns(model)
2147
+ # skip the last item in the list (the full identifier)
2148
+ for part in parser.to_list()[:-1]:
2149
+ ns = ns._safe_get(part)
2150
+ assert parser.identity, f"Error parsing source in to_model_type: {source}"
2151
+ self.sources.add(parser.identity)
2152
+ return ns
2153
+
2154
+ #--------------------------------------------------
2155
+ # Source Management
2156
+ #--------------------------------------------------
2157
+
2158
+ def _check_source_updates(self, sources: Iterable[str]):
2159
+ if not sources:
2160
+ return {}
2161
+ app_name = self.get_app_name()
2162
+
2163
+ source_types = dict[str, SourceInfo]()
2164
+ partitioned_sources: dict[str, dict[str, list[dict[str, str]]]] = defaultdict(
2165
+ lambda: defaultdict(list)
2166
+ )
2167
+ fqn_to_parts: dict[str, tuple[str, str, str]] = {}
2168
+
2169
+ for source in sources:
2170
+ parser = IdentityParser(source, True)
2171
+ parsed = parser.to_list()
2172
+ assert len(parsed) == 4, f"Invalid source: {source}"
2173
+ db, schema, entity, identity = parsed
2174
+ assert db and schema and entity and identity, f"Invalid source: {source}"
2175
+ source_types[identity] = cast(
2176
+ SourceInfo,
2177
+ {
2178
+ "type": None,
2179
+ "state": "",
2180
+ "columns_hash": None,
2181
+ "table_created_at": None,
2182
+ "stream_created_at": None,
2183
+ "last_ddl": None,
2184
+ },
2185
+ )
2186
+ partitioned_sources[db][schema].append({"entity": entity, "identity": identity})
2187
+ fqn_to_parts[identity] = (db, schema, entity)
2188
+
2189
+ if not partitioned_sources:
2190
+ return source_types
2191
+
2192
+ state_queries: list[str] = []
2193
+ for db, schemas in partitioned_sources.items():
2194
+ select_rows: list[str] = []
2195
+ for schema, tables in schemas.items():
2196
+ for table_info in tables:
2197
+ select_rows.append(
2198
+ "SELECT "
2199
+ f"{IdentityParser.to_sql_value(db)} AS catalog_name, "
2200
+ f"{IdentityParser.to_sql_value(schema)} AS schema_name, "
2201
+ f"{IdentityParser.to_sql_value(table_info['entity'])} AS table_name"
2202
+ )
2203
+
2204
+ if not select_rows:
2205
+ continue
2206
+
2207
+ target_entities_clause = "\n UNION ALL\n ".join(select_rows)
2208
+ # Main query:
2209
+ # 1. Enumerate the target tables via target_entities.
2210
+ # 2. Pull their metadata (last_altered, type) from INFORMATION_SCHEMA.TABLES.
2211
+ # 3. Look up the most recent stream activity for those FQNs only.
2212
+ # 4. Capture creation timestamps and use last_ddl vs created_at to classify each target,
2213
+ # so we mark tables as stale when they were recreated even if column hashes still match.
2214
+ state_queries.append(
2215
+ f"""WITH target_entities AS (
2216
+ {target_entities_clause}
2217
+ ),
2218
+ table_info AS (
2219
+ SELECT
2220
+ {app_name}.api.normalize_fq_ids(
2221
+ ARRAY_CONSTRUCT(
2222
+ CASE
2223
+ WHEN t.table_catalog = UPPER(t.table_catalog) THEN t.table_catalog
2224
+ ELSE '"' || t.table_catalog || '"'
2225
+ END || '.' ||
2226
+ CASE
2227
+ WHEN t.table_schema = UPPER(t.table_schema) THEN t.table_schema
2228
+ ELSE '"' || t.table_schema || '"'
2229
+ END || '.' ||
2230
+ CASE
2231
+ WHEN t.table_name = UPPER(t.table_name) THEN t.table_name
2232
+ ELSE '"' || t.table_name || '"'
2233
+ END
2234
+ )
2235
+ )[0]:identifier::string AS fqn,
2236
+ CONVERT_TIMEZONE('UTC', t.last_altered) AS last_ddl,
2237
+ CONVERT_TIMEZONE('UTC', t.created) AS table_created_at,
2238
+ t.table_type AS kind
2239
+ FROM {db}.INFORMATION_SCHEMA.tables t
2240
+ JOIN target_entities te
2241
+ ON t.table_catalog = te.catalog_name
2242
+ AND t.table_schema = te.schema_name
2243
+ AND t.table_name = te.table_name
2244
+ ),
2245
+ stream_activity AS (
2246
+ SELECT
2247
+ sa.fqn,
2248
+ MAX(sa.created_at) AS created_at
2249
+ FROM (
2250
+ SELECT
2251
+ {app_name}.api.normalize_fq_ids(ARRAY_CONSTRUCT(fq_object_name))[0]:identifier::string AS fqn,
2252
+ created_at
2253
+ FROM {app_name}.api.data_streams
2254
+ WHERE rai_database = '{PYREL_ROOT_DB}'
2255
+ ) sa
2256
+ JOIN table_info ti
2257
+ ON sa.fqn = ti.fqn
2258
+ GROUP BY sa.fqn
2259
+ )
2260
+ SELECT
2261
+ ti.fqn,
2262
+ ti.kind,
2263
+ ti.last_ddl,
2264
+ ti.table_created_at,
2265
+ sa.created_at AS stream_created_at,
2266
+ IFF(
2267
+ DATEDIFF(second, sa.created_at::timestamp, ti.last_ddl::timestamp) > 0,
2268
+ 'STALE',
2269
+ 'CURRENT'
2270
+ ) AS state
2271
+ FROM table_info ti
2272
+ LEFT JOIN stream_activity sa
2273
+ ON sa.fqn = ti.fqn
2274
+ """
2275
+ )
2276
+
2277
+ stale_fqns: list[str] = []
2278
+ for state_query in state_queries:
2279
+ for row in self._exec(state_query):
2280
+ row_dict = row.as_dict() if hasattr(row, "as_dict") else dict(row)
2281
+ row_fqn = row_dict["FQN"]
2282
+ parser = IdentityParser(row_fqn, True)
2283
+ fqn = parser.identity
2284
+ assert fqn, f"Error parsing returned FQN: {row_fqn}"
2285
+
2286
+ source_types[fqn]["type"] = (
2287
+ "TABLE" if row_dict["KIND"] == "BASE TABLE" else row_dict["KIND"]
2288
+ )
2289
+ source_types[fqn]["state"] = row_dict["STATE"]
2290
+ source_types[fqn]["last_ddl"] = normalize_datetime(row_dict.get("LAST_DDL"))
2291
+ source_types[fqn]["table_created_at"] = normalize_datetime(row_dict.get("TABLE_CREATED_AT"))
2292
+ source_types[fqn]["stream_created_at"] = normalize_datetime(row_dict.get("STREAM_CREATED_AT"))
2293
+ if row_dict["STATE"] == "STALE":
2294
+ stale_fqns.append(fqn)
2295
+
2296
+ if not stale_fqns:
2297
+ return source_types
2298
+
2299
+ # We batch stale tables by database/schema so each Snowflake query can hash
2300
+ # multiple objects at once instead of issuing one statement per table.
2301
+ stale_partitioned: dict[str, dict[str, list[dict[str, str]]]] = defaultdict(
2302
+ lambda: defaultdict(list)
2303
+ )
2304
+ for fqn in stale_fqns:
2305
+ db, schema, table = fqn_to_parts[fqn]
2306
+ stale_partitioned[db][schema].append({"table": table, "identity": fqn})
2307
+
2308
+ # Build one hash query per database, grouping schemas/tables inside so we submit
2309
+ # at most a handful of set-based statements to Snowflake.
2310
+ for db, schemas in stale_partitioned.items():
2311
+ column_select_rows: list[str] = []
2312
+ for schema, tables in schemas.items():
2313
+ for table_info in tables:
2314
+ # Build the literal rows for this db/schema so we can join back
2315
+ # against INFORMATION_SCHEMA.COLUMNS in a single statement.
2316
+ column_select_rows.append(
2317
+ "SELECT "
2318
+ f"{IdentityParser.to_sql_value(db)} AS catalog_name, "
2319
+ f"{IdentityParser.to_sql_value(schema)} AS schema_name, "
2320
+ f"{IdentityParser.to_sql_value(table_info['table'])} AS table_name"
2321
+ )
2322
+
2323
+ if not column_select_rows:
2324
+ continue
2325
+
2326
+ target_entities_clause = "\n UNION ALL\n ".join(column_select_rows)
2327
+ # Main query: compute deterministic column hashes for every stale table
2328
+ # in this database/schema batch so we can compare schemas without a round trip per table.
2329
+ column_query = f"""WITH target_entities AS (
2330
+ {target_entities_clause}
2331
+ ),
2332
+ column_info AS (
2333
+ SELECT
2334
+ {app_name}.api.normalize_fq_ids(
2335
+ ARRAY_CONSTRUCT(
2336
+ CASE
2337
+ WHEN c.table_catalog = UPPER(c.table_catalog) THEN c.table_catalog
2338
+ ELSE '"' || c.table_catalog || '"'
2339
+ END || '.' ||
2340
+ CASE
2341
+ WHEN c.table_schema = UPPER(c.table_schema) THEN c.table_schema
2342
+ ELSE '"' || c.table_schema || '"'
2343
+ END || '.' ||
2344
+ CASE
2345
+ WHEN c.table_name = UPPER(c.table_name) THEN c.table_name
2346
+ ELSE '"' || c.table_name || '"'
2347
+ END
2348
+ )
2349
+ )[0]:identifier::string AS fqn,
2350
+ c.column_name,
2351
+ CASE
2352
+ WHEN c.numeric_precision IS NOT NULL AND c.numeric_scale IS NOT NULL
2353
+ THEN c.data_type || '(' || c.numeric_precision || ',' || c.numeric_scale || ')'
2354
+ WHEN c.datetime_precision IS NOT NULL
2355
+ THEN c.data_type || '(0,' || c.datetime_precision || ')'
2356
+ WHEN c.character_maximum_length IS NOT NULL
2357
+ THEN c.data_type || '(' || c.character_maximum_length || ')'
2358
+ ELSE c.data_type
2359
+ END AS type_signature,
2360
+ IFF(c.is_nullable = 'YES', 'YES', 'NO') AS nullable_flag
2361
+ FROM {db}.INFORMATION_SCHEMA.COLUMNS c
2362
+ JOIN target_entities te
2363
+ ON c.table_catalog = te.catalog_name
2364
+ AND c.table_schema = te.schema_name
2365
+ AND c.table_name = te.table_name
2366
+ )
2367
+ SELECT
2368
+ fqn,
2369
+ HEX_ENCODE(
2370
+ HASH_AGG(
2371
+ HASH(
2372
+ column_name,
2373
+ type_signature,
2374
+ nullable_flag
2375
+ )
2376
+ )
2377
+ ) AS columns_hash
2378
+ FROM column_info
2379
+ GROUP BY fqn
2380
+ """
2381
+
2382
+ for row in self._exec(column_query):
2383
+ row_fqn = row["FQN"]
2384
+ parser = IdentityParser(row_fqn, True)
2385
+ fqn = parser.identity
2386
+ assert fqn, f"Error parsing returned FQN: {row_fqn}"
2387
+ source_types[fqn]["columns_hash"] = row["COLUMNS_HASH"]
2388
+
2389
+ return source_types
2390
+
2391
+ def _get_source_references(self, source_info: dict[str, SourceInfo]):
2392
+ app_name = self.get_app_name()
2393
+ missing_sources = []
2394
+ invalid_sources = {}
2395
+ source_references = []
2396
+ for source, info in source_info.items():
2397
+ source_type = info.get("type")
2398
+ if source_type is None:
2399
+ missing_sources.append(source)
2400
+ elif source_type not in ("TABLE", "VIEW"):
2401
+ invalid_sources[source] = source_type
2402
+ else:
2403
+ source_references.append(f"{app_name}.api.object_reference('{source_type}', '{source}')")
2404
+
2405
+ if missing_sources:
2406
+ current_role = self.get_sf_session().get_current_role()
2407
+ if current_role is None:
2408
+ current_role = self.config.get("role", None)
2409
+ debugging.warn(UnknownSourceWarning(missing_sources, current_role))
2410
+
2411
+ if invalid_sources:
2412
+ debugging.warn(InvalidSourceTypeWarning(invalid_sources))
2413
+
2414
+ self.source_references = source_references
2415
+ return source_references
2416
+
2417
+ #--------------------------------------------------
2418
+ # Transactions
2419
+ #--------------------------------------------------
2420
+
2421
+ def get_transaction(self, transaction_id):
2422
+ results = self._exec(
2423
+ f"CALL {APP_NAME}.api.get_transaction(?);", [transaction_id])
2424
+ if not results:
2425
+ return None
2426
+
2427
+ results = txn_list_to_dicts(results)
2428
+
2429
+ txn = {field: results[0][field] for field in GET_TXN_SQL_FIELDS}
2430
+
2431
+ state = txn.get("state")
2432
+ created_on = txn.get("created_on")
2433
+ finished_at = txn.get("finished_at")
2434
+ if created_on:
2435
+ # Transaction is still running
2436
+ if state not in TERMINAL_TXN_STATES:
2437
+ tz_info = created_on.tzinfo
2438
+ txn['duration'] = datetime.now(tz_info) - created_on
2439
+ # Transaction is terminal
2440
+ elif finished_at:
2441
+ txn['duration'] = finished_at - created_on
2442
+ # Transaction is still running and we have no state or finished_at
2443
+ else:
2444
+ txn['duration'] = timedelta(0)
2445
+ return txn
2446
+
2447
+ def list_transactions(self, **kwargs):
2448
+ id = kwargs.get("id", None)
2449
+ state = kwargs.get("state", None)
2450
+ engine = kwargs.get("engine", None)
2451
+ limit = kwargs.get("limit", 100)
2452
+ all_users = kwargs.get("all_users", False)
2453
+ created_by = kwargs.get("created_by", None)
2454
+ only_active = kwargs.get("only_active", False)
2455
+ where_clause_arr = []
2456
+
2457
+ if id:
2458
+ where_clause_arr.append(f"id = '{id}'")
2459
+ if state:
2460
+ where_clause_arr.append(f"state = '{state.upper()}'")
2461
+ if engine:
2462
+ where_clause_arr.append(f"LOWER(engine_name) = '{engine.lower()}'")
2463
+ else:
2464
+ if only_active:
2465
+ where_clause_arr.append("state in ('CREATED', 'RUNNING', 'PENDING')")
2466
+ if not all_users and created_by is not None:
2467
+ where_clause_arr.append(f"LOWER(created_by) = '{created_by.lower()}'")
2468
+
2469
+ if len(where_clause_arr):
2470
+ where_clause = f'WHERE {" AND ".join(where_clause_arr)}'
2471
+ else:
2472
+ where_clause = ""
2473
+
2474
+ sql_fields = ", ".join(LIST_TXN_SQL_FIELDS)
2475
+ query = f"SELECT {sql_fields} from {APP_NAME}.api.transactions {where_clause} ORDER BY created_on DESC LIMIT ?"
2476
+ results = self._exec(query, [limit])
2477
+ if not results:
2478
+ return []
2479
+ return txn_list_to_dicts(results)
2480
+
2481
+ def cancel_transaction(self, transaction_id):
2482
+ self._exec(f"CALL {APP_NAME}.api.cancel_own_transaction(?);", [transaction_id])
2483
+ if transaction_id in self._pending_transactions:
2484
+ self._pending_transactions.remove(transaction_id)
2485
+
2486
+ def cancel_pending_transactions(self):
2487
+ for txn_id in self._pending_transactions:
2488
+ self.cancel_transaction(txn_id)
2489
+
2490
+ def get_transaction_events(self, transaction_id: str, continuation_token:str=''):
2491
+ results = self._exec(
2492
+ f"SELECT {APP_NAME}.api.get_own_transaction_events(?, ?);",
2493
+ [transaction_id, continuation_token],
2494
+ )
2495
+ if not results:
2496
+ return {
2497
+ "events": [],
2498
+ "continuation_token": None
2499
+ }
2500
+ row = results[0][0]
2501
+ return json.loads(row)
2502
+
2503
+ #--------------------------------------------------
2504
+ # Snowflake specific
2505
+ #--------------------------------------------------
2506
+
2507
+ def get_version(self):
2508
+ results = self._exec(f"SELECT {APP_NAME}.app.get_release()")
2509
+ if not results:
2510
+ return None
2511
+ return results[0][0]
2512
+
2513
+ # CLI methods (list_warehouses, list_compute_pools, list_roles, list_apps,
2514
+ # list_databases, list_sf_schemas, list_tables) are now in CLIResources class
2515
+ # schema_info is kept in base Resources class since it's used by SnowflakeSchema._fetch_info()
2516
+
2517
+ def schema_info(self, database: str, schema: str, tables: Iterable[str]):
2518
+ """Get detailed schema information including primary keys, foreign keys, and columns."""
2519
+ app_name = self.get_app_name()
2520
+ # Only pass the db + schema as the identifier so that the resulting identity is correct
2521
+ parser = IdentityParser(f"{database}.{schema}")
2522
+
2523
+ with debugging.span("schema_info"):
2524
+ with debugging.span("primary_keys") as span:
2525
+ pk_query = f"SHOW PRIMARY KEYS IN SCHEMA {parser.identity};"
2526
+ pks = self._exec(pk_query)
2527
+ span["sql"] = pk_query
2528
+
2529
+ with debugging.span("foreign_keys") as span:
2530
+ fk_query = f"SHOW IMPORTED KEYS IN SCHEMA {parser.identity};"
2531
+ fks = self._exec(fk_query)
2532
+ span["sql"] = fk_query
2533
+
2534
+ # IdentityParser will parse a single value (with no ".") and store it in this case in the db field
2535
+ with debugging.span("columns") as span:
2536
+ tables_str = ", ".join([f"'{IdentityParser(t).db}'" for t in tables])
2537
+ query = textwrap.dedent(f"""
2538
+ begin
2539
+ SHOW COLUMNS IN SCHEMA {parser.identity};
2540
+ let r resultset := (
2541
+ SELECT
2542
+ CASE
2543
+ WHEN "table_name" = UPPER("table_name") THEN "table_name"
2544
+ ELSE '"' || "table_name" || '"'
2545
+ END as "table_name",
2546
+ "column_name",
2547
+ "data_type",
2548
+ CASE
2549
+ WHEN ARRAY_CONTAINS(PARSE_JSON("data_type"):"type", {app_name}.app.get_supported_column_types()) THEN TRUE
2550
+ ELSE FALSE
2551
+ END as "supported_type"
2552
+ FROM table(result_scan(-1)) as t
2553
+ WHERE "table_name" in ({tables_str})
2554
+ );
2555
+ return table(r);
2556
+ end;
2557
+ """)
2558
+ span["sql"] = query
2559
+ columns = self._exec(query)
2560
+
2561
+ results = defaultdict(lambda: {"pks": [], "fks": {}, "columns": {}, "invalid_columns": {}})
2562
+ if pks:
2563
+ for row in pks:
2564
+ results[row[3]]["pks"].append(row[4]) # type: ignore
2565
+ if fks:
2566
+ for row in fks:
2567
+ results[row[7]]["fks"][row[8]] = row[3]
2568
+ if columns:
2569
+ # It seems that a SF parameter (QUOTED_IDENTIFIERS_IGNORE_CASE) can control
2570
+ # whether snowflake will ignore case on `row.data_type`,
2571
+ # so we have to use column indexes instead :(
2572
+ for row in columns:
2573
+ table_name = row[0]
2574
+ column_name = row[1]
2575
+ data_type = row[2]
2576
+ supported_type = row[3]
2577
+ # Filter out unsupported types
2578
+ if supported_type:
2579
+ results[table_name]["columns"][column_name] = data_type
2580
+ else:
2581
+ results[table_name]["invalid_columns"][column_name] = data_type
2582
+ return results
2583
+
2584
+ def get_cloud_provider(self) -> str:
2585
+ """
2586
+ Detect whether this is Snowflake on Azure, or AWS using Snowflake's CURRENT_REGION().
2587
+ Returns 'azure' or 'aws'.
2588
+ """
2589
+ if self._session:
2590
+ try:
2591
+ # Query Snowflake's current region using the built-in function
2592
+ result = self._session.sql("SELECT CURRENT_REGION()").collect()
2593
+ if result:
2594
+ region_info = result[0][0]
2595
+ # Check if the region string contains the cloud provider name
2596
+ if isinstance(region_info, str):
2597
+ region_str = region_info.lower()
2598
+ # Check for cloud providers in the region string
2599
+ if 'azure' in region_str:
2600
+ return 'azure'
2601
+ else:
2602
+ return 'aws'
2603
+ except Exception:
2604
+ pass
2605
+
2606
+ # Fallback to AWS as default if detection fails
2607
+ return 'aws'
2608
+
2609
+ #--------------------------------------------------
2610
+ # Snowflake Wrapper
2611
+ #--------------------------------------------------
2612
+
2613
+ class PrimaryKey:
2614
+ pass
2615
+
2616
+ class _Snowflake:
2617
+ def __init__(self, model, auto_import=False):
2618
+ self._model = model
2619
+ self._auto_import = auto_import
2620
+ if not isinstance(model._client.resources, Resources):
2621
+ raise ValueError("Snowflake model must be used with a snowflake config")
2622
+ self._dbs = {}
2623
+ imports = model._client.resources.list_imports(model=model.name)
2624
+ self._import_structure(imports)
2625
+
2626
+ def _import_structure(self, imports: list[Import]):
2627
+ tree = self._dbs
2628
+ # pre-create existing imports
2629
+ schemas = set()
2630
+ for item in imports:
2631
+ parser = IdentityParser(item["name"])
2632
+ database_name, schema_name, table_name = parser.to_list()[:-1]
2633
+ database = getattr(self, database_name)
2634
+ schema = getattr(database, schema_name)
2635
+ schemas.add(schema)
2636
+ schema._add(table_name, is_imported=True)
2637
+ return tree
2638
+
2639
+ def _safe_get(self, name:str) -> 'SnowflakeDB':
2640
+ name = name
2641
+ if name in self._dbs:
2642
+ return self._dbs[name]
2643
+ self._dbs[name] = SnowflakeDB(self, name)
2644
+ return self._dbs[name]
2645
+
2646
+ def __getattr__(self, name: str) -> 'SnowflakeDB':
2647
+ return self._safe_get(name)
2648
+
2649
+
2650
+ class Snowflake(_Snowflake):
2651
+ def __init__(self, model: dsl.Graph, auto_import=False):
2652
+ if model._config.get_bool("use_graph_index", USE_GRAPH_INDEX):
2653
+ raise SnowflakeProxySourceError()
2654
+ else:
2655
+ debugging.warn(SnowflakeProxyAPIDeprecationWarning())
2656
+
2657
+ super().__init__(model, auto_import)
2658
+
2659
+ class SnowflakeDB:
2660
+ def __init__(self, parent, name):
2661
+ self._name = name
2662
+ self._parent = parent
2663
+ self._model = parent._model
2664
+ self._schemas = {}
2665
+
2666
+ def _safe_get(self, name: str) -> 'SnowflakeSchema':
2667
+ name = name
2668
+ if name in self._schemas:
2669
+ return self._schemas[name]
2670
+ self._schemas[name] = SnowflakeSchema(self, name)
2671
+ return self._schemas[name]
2672
+
2673
+ def __getattr__(self, name: str) -> 'SnowflakeSchema':
2674
+ return self._safe_get(name)
2675
+
2676
+ class SnowflakeSchema:
2677
+ def __init__(self, parent, name):
2678
+ self._name = name
2679
+ self._parent = parent
2680
+ self._model = parent._model
2681
+ self._tables = {}
2682
+ self._imported = set()
2683
+ self._table_info = defaultdict(lambda: {"pks": [], "fks": {}, "columns": {}, "invalid_columns": {}})
2684
+ self._dirty = True
2685
+
2686
+ def _fetch_info(self):
2687
+ if not self._dirty:
2688
+ return
2689
+ self._table_info = self._model._client.resources.schema_info(self._parent._name, self._name, list(self._tables.keys()))
2690
+
2691
+ check_column_types = self._model._config.get("check_column_types", True)
2692
+
2693
+ if check_column_types:
2694
+ self._check_and_confirm_invalid_columns()
2695
+
2696
+ self._dirty = False
2697
+
2698
+ def _check_and_confirm_invalid_columns(self):
2699
+ """Check for invalid columns across the schema's tables."""
2700
+ tables_with_invalid_columns = {}
2701
+ for table_name, table_info in self._table_info.items():
2702
+ if table_info.get("invalid_columns"):
2703
+ tables_with_invalid_columns[table_name] = table_info["invalid_columns"]
2704
+
2705
+ if tables_with_invalid_columns:
2706
+ from relationalai.errors import UnsupportedColumnTypesWarning
2707
+ UnsupportedColumnTypesWarning(tables_with_invalid_columns)
2708
+
2709
+ def _add(self, name, is_imported=False):
2710
+ if name in self._tables:
2711
+ return self._tables[name]
2712
+ self._dirty = True
2713
+ if is_imported:
2714
+ self._imported.add(name)
2715
+ else:
2716
+ self._tables[name] = SnowflakeTable(self, name)
2717
+ return self._tables.get(name)
2718
+
2719
+ def _safe_get(self, name: str) -> 'SnowflakeTable | None':
2720
+ table = self._add(name)
2721
+ return table
2722
+
2723
+ def __getattr__(self, name: str) -> 'SnowflakeTable | None':
2724
+ return self._safe_get(name)
2725
+
2726
+
2727
+ class SnowflakeTable(dsl.Type):
2728
+ def __init__(self, parent, name):
2729
+ super().__init__(parent._model, f"sf_{name}")
2730
+ # hack to make this work for pathfinder
2731
+ self._type.parents.append(m.Builtins.PQFilterAnnotation)
2732
+ self._name = name
2733
+ self._model = parent._model
2734
+ self._parent = parent
2735
+ self._aliases = {}
2736
+ self._finalzed = False
2737
+ self._source = runtime_env.get_source()
2738
+ relation_name = to_fqn_relation_name(self.fqname())
2739
+ self._model.install_raw(f"declare {relation_name}")
2740
+
2741
+ def __call__(self, *args, **kwargs):
2742
+ self._lazy_init()
2743
+ return super().__call__(*args, **kwargs)
2744
+
2745
+ def add(self, *args, **kwargs):
2746
+ self._lazy_init()
2747
+ return super().add(*args, **kwargs)
2748
+
2749
+ def extend(self, *args, **kwargs):
2750
+ self._lazy_init()
2751
+ return super().extend(*args, **kwargs)
2752
+
2753
+ def known_properties(self):
2754
+ self._lazy_init()
2755
+ return super().known_properties()
2756
+
2757
+ def _lazy_init(self):
2758
+ if self._finalzed:
2759
+ return
2760
+
2761
+ parent = self._parent
2762
+ name = self._name
2763
+ use_graph_index = self._model._config.get("use_graph_index", USE_GRAPH_INDEX)
2764
+
2765
+ if not use_graph_index and name not in parent._imported:
2766
+ if self._parent._parent._parent._auto_import:
2767
+ with Spinner(f"Creating stream for {self.fqname()}", f"Stream for {self.fqname()} created successfully"):
2768
+ db_name = parent._parent._name
2769
+ schema_name = parent._name
2770
+ self._model._client.resources.create_import_stream(ImportSourceTable(db_name, schema_name, name), self._model.name)
2771
+ print("")
2772
+ parent._imported.add(name)
2773
+ else:
2774
+ imports = self._model._client.resources.list_imports(model=self._model.name)
2775
+ for item in imports:
2776
+ cur_name = item["name"].lower().split(".")[-1]
2777
+ parent._imported.add(cur_name)
2778
+ if name not in parent._imported:
2779
+ exception = SnowflakeImportMissingException(runtime_env.get_source(), self.fqname(), self._model.name)
2780
+ raise exception from None
2781
+
2782
+ parent._fetch_info()
2783
+ self._finalize()
2784
+
2785
+ def _finalize(self):
2786
+ if self._finalzed:
2787
+ return
2788
+
2789
+ self._finalzed = True
2790
+ self._schema = self._parent._table_info[self._name]
2791
+
2792
+ # Set the relation name to the sanitized version of the fully qualified name
2793
+ relation_name = to_fqn_relation_name(self.fqname())
2794
+
2795
+ model:dsl.Graph = self._model
2796
+ edb = getattr(std.rel, relation_name)
2797
+ edb._rel.parents.append(m.Builtins.EDB)
2798
+ id_rel = getattr(std.rel, f"{relation_name}_pyrel_id")
2799
+
2800
+ with model.rule(globalize=True, source=self._source):
2801
+ id, val = dsl.create_vars(2)
2802
+ edb(dsl.Symbol("METADATA$ROW_ID"), id, val)
2803
+ std.rel.SHA1(id)
2804
+ id_rel.add(id)
2805
+
2806
+ with model.rule(dynamic=True, globalize=True, source=self._source):
2807
+ prop, id, val = dsl.create_vars(3)
2808
+ id_rel(id)
2809
+ std.rel.SHA1(id)
2810
+ self.add(snowflake_id=id)
2811
+
2812
+ for prop, prop_type in self._schema["columns"].items():
2813
+ _prop = prop
2814
+ if _prop.startswith("_"):
2815
+ _prop = "col" + prop
2816
+
2817
+ prop_ident = sanitize_identifier(_prop.lower())
2818
+
2819
+ with model.rule(dynamic=True, globalize=True, source=self._source):
2820
+ id, val = dsl.create_vars(2)
2821
+ edb(dsl.Symbol(prop), id, val)
2822
+ std.rel.SHA1(id)
2823
+ _prop = getattr(self, prop_ident)
2824
+ if not _prop:
2825
+ raise ValueError(f"Property {_prop} couldn't be accessed on {self.fqname()}")
2826
+ if _prop.is_multi_valued:
2827
+ inst = self(snowflake_id=id)
2828
+ getattr(inst, prop_ident).add(val)
2829
+ else:
2830
+ self(snowflake_id=id).set(**{prop_ident: val})
2831
+
2832
+ # Because we're bypassing a bunch of the normal Type.add machinery here,
2833
+ # we need to manually account for the case where people are using value types.
2834
+ def wrapped(x):
2835
+ if not model._config.get("compiler.use_value_types", False):
2836
+ return x
2837
+ other_id = dsl.create_var()
2838
+ model._action(dsl.build.construct(self._type, [x, other_id]))
2839
+ return other_id
2840
+
2841
+ # new UInt128 schema mapping rules
2842
+ with model.rule(dynamic=True, globalize=True, source=self._source):
2843
+ id = dsl.create_var()
2844
+ # This will generate an arity mismatch warning when used with the old SHA-1 Data Streams.
2845
+ # Ideally we have the `@no_diagnostics(:ARITY_MISMATCH)` attribute on the relation using
2846
+ # the METADATA$KEY column but that ended up being a more involved change then expected
2847
+ # for avoiding a non-blocking warning
2848
+ edb(dsl.Symbol("METADATA$KEY"), id)
2849
+ std.rel.UInt128(id)
2850
+ self.add(wrapped(id), snowflake_id=id)
2851
+
2852
+ for prop, prop_type in self._schema["columns"].items():
2853
+ _prop = prop
2854
+ if _prop.startswith("_"):
2855
+ _prop = "col" + prop
2856
+
2857
+ prop_ident = sanitize_identifier(_prop.lower())
2858
+ with model.rule(dynamic=True, globalize=True, source=self._source):
2859
+ id, val = dsl.create_vars(2)
2860
+ edb(dsl.Symbol(prop), id, val)
2861
+ std.rel.UInt128(id)
2862
+ _prop = getattr(self, prop_ident)
2863
+ if not _prop:
2864
+ raise ValueError(f"Property {_prop} couldn't be accessed on {self.fqname()}")
2865
+ if _prop.is_multi_valued:
2866
+ inst = self(id)
2867
+ getattr(inst, prop_ident).add(val)
2868
+ else:
2869
+ model._check_property(_prop._prop)
2870
+ raw_relation = getattr(std.rel, prop_ident)
2871
+ dsl.tag(raw_relation, dsl.Builtins.FunctionAnnotation)
2872
+ raw_relation.add(wrapped(id), val)
2873
+
2874
+ def namespace(self):
2875
+ return f"{self._parent._parent._name}.{self._parent._name}"
2876
+
2877
+ def fqname(self):
2878
+ return f"{self.namespace()}.{self._name}"
2879
+
2880
+ def describe(self, **kwargs):
2881
+ model = self._model
2882
+ for k, v in kwargs.items():
2883
+ if v is PrimaryKey:
2884
+ self._schema["pks"] = [k]
2885
+ elif isinstance(v, tuple):
2886
+ (table, name) = v
2887
+ if isinstance(table, SnowflakeTable):
2888
+ fk_table = table
2889
+ pk = fk_table._schema["pks"]
2890
+ with model.rule():
2891
+ inst = fk_table()
2892
+ me = self()
2893
+ getattr(inst, pk[0]) == getattr(me, k)
2894
+ if getattr(self, name).is_multi_valued:
2895
+ getattr(me, name).add(inst)
2896
+ else:
2897
+ me.set(**{name: inst})
2898
+ else:
2899
+ raise ValueError(f"Invalid foreign key {v}")
2900
+ else:
2901
+ raise ValueError(f"Invalid column {k}={v}")
2902
+ return self
2903
+
2904
+ class Provider(ProviderBase):
2905
+ def __init__(
2906
+ self,
2907
+ profile: str | None = None,
2908
+ config: Config | None = None,
2909
+ resources: Resources | None = None,
2910
+ generation: Generation | None = None,
2911
+ ):
2912
+ if resources:
2913
+ self.resources = resources
2914
+ else:
2915
+ from .resources_factory import create_resources_instance
2916
+ self.resources = create_resources_instance(
2917
+ config=config,
2918
+ profile=profile,
2919
+ generation=generation or Generation.V0,
2920
+ dry_run=False,
2921
+ language="rel",
2922
+ )
2923
+
2924
+ def list_streams(self, model:str):
2925
+ return self.resources.list_imports(model=model)
2926
+
2927
+ def create_streams(self, sources:List[str], model:str, force=False):
2928
+ if not self.resources.get_graph(model):
2929
+ self.resources.create_graph(model)
2930
+ def parse_source(raw:str):
2931
+ parser = IdentityParser(raw)
2932
+ assert parser.is_complete, "Snowflake table imports must be in `database.schema.table` format"
2933
+ return ImportSourceTable(*parser.to_list())
2934
+ for source in sources:
2935
+ source_table = parse_source(source)
2936
+ try:
2937
+ with Spinner(f"Creating stream for {source_table.name}", f"Stream for {source_table.name} created successfully"):
2938
+ if force:
2939
+ self.resources.delete_import(source_table.name, model, True)
2940
+ self.resources.create_import_stream(source_table, model)
2941
+ except Exception as e:
2942
+ if "stream already exists" in f"{e}":
2943
+ raise Exception(f"\n\nStream'{source_table.name.upper()}' already exists.")
2944
+ elif "engine not found" in f"{e}":
2945
+ raise Exception("\n\nNo engines found in a READY state. Please use `engines:create` to create an engine that will be used to initialize the target relation.")
2946
+ else:
2947
+ raise e
2948
+ with Spinner("Waiting for imports to complete", "Imports complete"):
2949
+ self.resources.poll_imports(sources, model)
2950
+
2951
+ def delete_stream(self, stream_id: str, model: str):
2952
+ return self.resources.delete_import(stream_id, model)
2953
+
2954
+ def sql(self, query:str, params:List[Any]=[], format:Literal["list", "pandas", "polars", "lazy"]="list"):
2955
+ # note: default format cannot be pandas because .to_pandas() only works on SELECT queries
2956
+ result = self.resources._exec(query, params, raw=True, help=False)
2957
+ if format == "lazy":
2958
+ return cast(snowflake.snowpark.DataFrame, result)
2959
+ elif format == "list":
2960
+ return cast(list, result.collect())
2961
+ elif format == "pandas":
2962
+ import pandas as pd
2963
+ try:
2964
+ # use to_pandas for SELECT queries
2965
+ return cast(pd.DataFrame, result.to_pandas())
2966
+ except Exception:
2967
+ # handle non-SELECT queries like SHOW
2968
+ return pd.DataFrame(result.collect())
2969
+ elif format == "polars":
2970
+ import polars as pl # type: ignore
2971
+ return pl.DataFrame(
2972
+ [row.as_dict() for row in result.collect()],
2973
+ orient="row",
2974
+ strict=False,
2975
+ infer_schema_length=None
2976
+ )
2977
+ else:
2978
+ raise ValueError(f"Invalid format {format}. Should be one of 'list', 'pandas', 'polars', 'lazy'")
2979
+
2980
+ def activate(self):
2981
+ with Spinner("Activating RelationalAI app...", "RelationalAI app activated"):
2982
+ self.sql("CALL RELATIONALAI.APP.ACTIVATE();")
2983
+
2984
+ def deactivate(self):
2985
+ with Spinner("Deactivating RelationalAI app...", "RelationalAI app deactivated"):
2986
+ self.sql("CALL RELATIONALAI.APP.DEACTIVATE();")
2987
+
2988
+ def drop_service(self):
2989
+ warnings.warn(
2990
+ "The drop_service method has been deprecated in favor of deactivate",
2991
+ DeprecationWarning,
2992
+ stacklevel=2,
2993
+ )
2994
+ self.deactivate()
2995
+
2996
+ def resume_service(self):
2997
+ warnings.warn(
2998
+ "The resume_service method has been deprecated in favor of activate",
2999
+ DeprecationWarning,
3000
+ stacklevel=2,
3001
+ )
3002
+ self.activate()
3003
+
3004
+
3005
+ #--------------------------------------------------
3006
+ # SnowflakeClient
3007
+ #--------------------------------------------------
3008
+ class SnowflakeClient(Client):
3009
+ def create_database(self, isolated=True, nowait_durable=True, headers: Dict | None = None):
3010
+ from relationalai.tools.cli_helpers import validate_engine_name
3011
+
3012
+ assert isinstance(self.resources, Resources)
3013
+
3014
+ if self.last_database_version == len(self.resources.sources):
3015
+ return
3016
+
3017
+ model = self._source_database
3018
+ app_name = self.resources.get_app_name()
3019
+ engine_name = self.resources.get_default_engine_name()
3020
+ engine_size = self.resources.config.get_default_engine_size()
3021
+
3022
+ # Validate engine name
3023
+ is_name_valid, _ = validate_engine_name(engine_name)
3024
+ if not is_name_valid:
3025
+ raise EngineNameValidationException(engine_name)
3026
+
3027
+ # Validate engine size
3028
+ valid_sizes = self.resources.get_engine_sizes()
3029
+ if not isinstance(engine_size, str) or engine_size not in valid_sizes:
3030
+ raise InvalidEngineSizeError(str(engine_size), valid_sizes)
3031
+
3032
+ program_span_id = debugging.get_program_span_id()
3033
+
3034
+ query_attrs_dict = json.loads(headers.get("X-Query-Attributes", "{}")) if headers else {}
3035
+ with debugging.span("poll_use_index", sources=self.resources.sources, model=model, engine=engine_name, **query_attrs_dict):
3036
+ self.maybe_poll_use_index(
3037
+ app_name=app_name,
3038
+ sources=self.resources.sources,
3039
+ model=model,
3040
+ engine_name=engine_name,
3041
+ engine_size=engine_size,
3042
+ program_span_id=program_span_id,
3043
+ headers=headers
3044
+ )
3045
+
3046
+ self.last_database_version = len(self.resources.sources)
3047
+ self._manage_packages()
3048
+
3049
+ if isolated and not self.keep_model:
3050
+ atexit.register(self.delete_database)
3051
+
3052
+ def maybe_poll_use_index(
3053
+ self,
3054
+ app_name: str,
3055
+ sources: Iterable[str],
3056
+ model: str,
3057
+ engine_name: str,
3058
+ engine_size: str | None = None,
3059
+ program_span_id: str | None = None,
3060
+ headers: Dict | None = None,
3061
+ ):
3062
+ """Only call _poll_use_index if there are sources to process."""
3063
+ assert isinstance(self.resources, Resources)
3064
+ return self.resources.maybe_poll_use_index(
3065
+ app_name=app_name,
3066
+ sources=sources,
3067
+ model=model,
3068
+ engine_name=engine_name,
3069
+ engine_size=engine_size,
3070
+ program_span_id=program_span_id,
3071
+ headers=headers
3072
+ )
3073
+
3074
+
3075
+ #--------------------------------------------------
3076
+ # Graph
3077
+ #--------------------------------------------------
3078
+
3079
+ def Graph(
3080
+ name,
3081
+ *,
3082
+ profile: str | None = None,
3083
+ config: Config,
3084
+ dry_run: bool = False,
3085
+ isolated: bool = True,
3086
+ connection: Session | None = None,
3087
+ keep_model: bool = False,
3088
+ nowait_durable: bool = True,
3089
+ format: str = "default",
3090
+ ):
3091
+ from .resources_factory import create_resources_instance
3092
+ from .use_index_resources import UseIndexResources
3093
+
3094
+ use_graph_index = config.get("use_graph_index", USE_GRAPH_INDEX)
3095
+ use_monotype_operators = config.get("compiler.use_monotype_operators", False)
3096
+
3097
+ # Create resources instance using factory
3098
+ resources = create_resources_instance(
3099
+ config=config,
3100
+ profile=profile,
3101
+ connection=connection,
3102
+ generation=Generation.V0,
3103
+ dry_run=False, # Resources instance dry_run is separate from client dry_run
3104
+ language="rel",
3105
+ )
3106
+
3107
+ # Determine client class based on resources type and config
3108
+ # SnowflakeClient is used for resources that support use_index functionality
3109
+ if use_graph_index or isinstance(resources, UseIndexResources):
3110
+ client_class = SnowflakeClient
3111
+ else:
3112
+ client_class = Client
3113
+
3114
+ client = client_class(
3115
+ resources,
3116
+ rel.Compiler(config),
3117
+ name,
3118
+ config,
3119
+ dry_run=dry_run,
3120
+ isolated=isolated,
3121
+ keep_model=keep_model,
3122
+ nowait_durable=nowait_durable
3123
+ )
3124
+ base_rel = """
3125
+ @inline
3126
+ def make_identity(x..., z):
3127
+ rel_primitive_hash_tuple_uint128(x..., z)
3128
+
3129
+ @inline
3130
+ def pyrel_default({F}, c, k..., v):
3131
+ F(k..., v) or (not F(k..., _) and v = c)
3132
+
3133
+ @inline
3134
+ def pyrel_unwrap(x in UInt128, y): y = x
3135
+
3136
+ @inline
3137
+ def pyrel_dates_period_days(x in Date, y in Date, z in Int):
3138
+ exists((u) | dates_period_days(x, y , u) and u = ::std::common::^Day[z])
3139
+
3140
+ @inline
3141
+ def pyrel_datetimes_period_milliseconds(x in DateTime, y in DateTime, z in Int):
3142
+ exists((u) | datetimes_period_milliseconds(x, y , u) and u = ^Millisecond[z])
3143
+
3144
+ @inline
3145
+ def pyrel_bool_filter(a, b, {F}, z): { z = if_then_else[F(a, b), boolean_true, boolean_false] }
3146
+
3147
+ @inline
3148
+ def pyrel_strftime(v, fmt, tz in String, s in String):
3149
+ (Date(v) and s = format_date[v, fmt])
3150
+ or (DateTime(v) and s = format_datetime[v, fmt, tz])
3151
+
3152
+ @inline
3153
+ def pyrel_regex_match_all(pattern, string in String, pos in Int, offset in Int, match in String):
3154
+ regex_match_all(pattern, string, offset, match) and offset >= pos
3155
+
3156
+ @inline
3157
+ def pyrel_regex_match(pattern, string in String, pos in Int, offset in Int, match in String):
3158
+ pyrel_regex_match_all(pattern, string, pos, offset, match) and offset = pos
3159
+
3160
+ @inline
3161
+ def pyrel_regex_search(pattern, string in String, pos in Int, offset in Int, match in String):
3162
+ enumerate(pyrel_regex_match_all[pattern, string, pos], 1, offset, match)
3163
+
3164
+ @inline
3165
+ def pyrel_regex_sub(pattern, repl in String, string in String, result in String):
3166
+ string_replace_multiple(string, {(last[regex_match_all[pattern, string]], repl)}, result)
3167
+
3168
+ @inline
3169
+ def pyrel_capture_group(regex in Pattern, string in String, pos in Int, index, match in String):
3170
+ (Integer(index) and capture_group_by_index(regex, string, pos, index, match)) or
3171
+ (String(index) and capture_group_by_name(regex, string, pos, index, match))
3172
+
3173
+ declare __resource
3174
+ declare __compiled_patterns
3175
+ """
3176
+ if use_monotype_operators:
3177
+ base_rel += """
3178
+
3179
+ // use monotyped operators
3180
+ from ::std::monotype import +, -, *, /, <, <=, >, >=
3181
+ """
3182
+ pyrel_base = dsl.build.raw_task(base_rel)
3183
+ debugging.set_source(pyrel_base)
3184
+ client.install("pyrel_base", pyrel_base)
3185
+ return dsl.Graph(client, name, format=format)