mage-ai 0.9.69__py3-none-any.whl → 0.9.71__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.

Potentially problematic release.


This version of mage-ai might be problematic. Click here for more details.

Files changed (624) hide show
  1. mage_ai/ai/utils/xgboost.py +222 -0
  2. mage_ai/api/errors.py +37 -25
  3. mage_ai/api/operations/base.py +13 -1
  4. mage_ai/api/parsers/PipelineScheduleParser.py +1 -1
  5. mage_ai/api/policies/BackfillPolicy.py +1 -0
  6. mage_ai/api/policies/BlockOutputPolicy.py +40 -17
  7. mage_ai/api/policies/GlobalDataProductPolicy.py +91 -41
  8. mage_ai/api/policies/KernelPolicy.py +55 -32
  9. mage_ai/api/policies/KernelProcessPolicy.py +56 -0
  10. mage_ai/api/policies/OutputPolicy.py +73 -41
  11. mage_ai/api/policies/PipelinePolicy.py +206 -137
  12. mage_ai/api/policies/WorkspacePolicy.py +1 -0
  13. mage_ai/api/presenters/BackfillPresenter.py +1 -0
  14. mage_ai/api/presenters/BlockLayoutItemPresenter.py +9 -7
  15. mage_ai/api/presenters/BlockPresenter.py +1 -1
  16. mage_ai/api/presenters/GlobalDataProductPresenter.py +6 -1
  17. mage_ai/api/presenters/KernelPresenter.py +5 -26
  18. mage_ai/api/presenters/KernelProcessPresenter.py +28 -0
  19. mage_ai/api/presenters/PipelinePresenter.py +18 -5
  20. mage_ai/api/presenters/StatusPresenter.py +2 -0
  21. mage_ai/api/presenters/SyncPresenter.py +25 -0
  22. mage_ai/api/resources/AutocompleteItemResource.py +1 -1
  23. mage_ai/api/resources/BlockLayoutItemResource.py +90 -44
  24. mage_ai/api/resources/BlockOutputResource.py +42 -9
  25. mage_ai/api/resources/BlockResource.py +4 -3
  26. mage_ai/api/resources/BlockRunResource.py +27 -22
  27. mage_ai/api/resources/ClusterResource.py +4 -1
  28. mage_ai/api/resources/CustomTemplateResource.py +34 -14
  29. mage_ai/api/resources/DataProviderResource.py +1 -1
  30. mage_ai/api/resources/ExecutionStateResource.py +3 -1
  31. mage_ai/api/resources/FileContentResource.py +8 -2
  32. mage_ai/api/resources/FileResource.py +10 -4
  33. mage_ai/api/resources/FileVersionResource.py +3 -1
  34. mage_ai/api/resources/GitBranchResource.py +101 -31
  35. mage_ai/api/resources/GitCustomBranchResource.py +29 -1
  36. mage_ai/api/resources/GlobalDataProductResource.py +44 -7
  37. mage_ai/api/resources/GlobalHookResource.py +4 -1
  38. mage_ai/api/resources/IntegrationDestinationResource.py +6 -2
  39. mage_ai/api/resources/IntegrationSourceResource.py +8 -4
  40. mage_ai/api/resources/IntegrationSourceStreamResource.py +6 -2
  41. mage_ai/api/resources/KernelProcessResource.py +44 -0
  42. mage_ai/api/resources/KernelResource.py +25 -3
  43. mage_ai/api/resources/OauthResource.py +1 -1
  44. mage_ai/api/resources/OutputResource.py +33 -11
  45. mage_ai/api/resources/PageBlockLayoutResource.py +34 -23
  46. mage_ai/api/resources/PipelineInteractionResource.py +31 -15
  47. mage_ai/api/resources/PipelineResource.py +258 -125
  48. mage_ai/api/resources/PipelineRunResource.py +52 -7
  49. mage_ai/api/resources/PipelineScheduleResource.py +11 -2
  50. mage_ai/api/resources/PipelineTriggerResource.py +6 -1
  51. mage_ai/api/resources/ProjectResource.py +18 -7
  52. mage_ai/api/resources/PullRequestResource.py +6 -4
  53. mage_ai/api/resources/SecretResource.py +1 -1
  54. mage_ai/api/resources/SeedResource.py +8 -1
  55. mage_ai/api/resources/StatusResource.py +21 -6
  56. mage_ai/api/resources/SyncResource.py +6 -8
  57. mage_ai/api/resources/VariableResource.py +46 -26
  58. mage_ai/api/resources/VersionControlProjectResource.py +9 -2
  59. mage_ai/api/resources/WidgetResource.py +1 -1
  60. mage_ai/api/resources/WorkspaceResource.py +6 -5
  61. mage_ai/api/views.py +47 -40
  62. mage_ai/authentication/permissions/seed.py +16 -2
  63. mage_ai/authentication/providers/oidc.py +21 -1
  64. mage_ai/autocomplete/utils.py +13 -9
  65. mage_ai/cache/base.py +1 -1
  66. mage_ai/cache/block.py +18 -12
  67. mage_ai/cache/block_action_object/__init__.py +33 -5
  68. mage_ai/cache/file.py +22 -19
  69. mage_ai/cache/pipeline.py +18 -12
  70. mage_ai/cli/main.py +1 -0
  71. mage_ai/cluster_manager/aws/emr_cluster_manager.py +9 -5
  72. mage_ai/cluster_manager/config.py +2 -2
  73. mage_ai/cluster_manager/kubernetes/workload_manager.py +52 -1
  74. mage_ai/cluster_manager/manage.py +1 -1
  75. mage_ai/cluster_manager/workspace/base.py +7 -1
  76. mage_ai/cluster_manager/workspace/kubernetes.py +22 -1
  77. mage_ai/command_center/applications/factory.py +10 -7
  78. mage_ai/command_center/applications/utils.py +2 -2
  79. mage_ai/command_center/files/factory.py +17 -15
  80. mage_ai/command_center/presenters/text.py +1 -1
  81. mage_ai/command_center/utils.py +25 -13
  82. mage_ai/data/__init__.py +0 -0
  83. mage_ai/data/constants.py +45 -0
  84. mage_ai/data/models/__init__.py +0 -0
  85. mage_ai/data/models/base.py +119 -0
  86. mage_ai/data/models/constants.py +1 -0
  87. mage_ai/data/models/generator.py +115 -0
  88. mage_ai/data/models/manager.py +168 -0
  89. mage_ai/data/models/pyarrow/__init__.py +0 -0
  90. mage_ai/data/models/pyarrow/record_batch.py +55 -0
  91. mage_ai/data/models/pyarrow/shared.py +21 -0
  92. mage_ai/data/models/pyarrow/table.py +8 -0
  93. mage_ai/data/models/reader.py +103 -0
  94. mage_ai/data/models/utils.py +59 -0
  95. mage_ai/data/models/writer.py +91 -0
  96. mage_ai/data/tabular/__init__.py +0 -0
  97. mage_ai/data/tabular/constants.py +23 -0
  98. mage_ai/data/tabular/mocks.py +19 -0
  99. mage_ai/data/tabular/models.py +126 -0
  100. mage_ai/data/tabular/reader.py +602 -0
  101. mage_ai/data/tabular/utils.py +102 -0
  102. mage_ai/data/tabular/writer.py +266 -0
  103. mage_ai/data/variables/__init__.py +0 -0
  104. mage_ai/data/variables/wrapper.py +54 -0
  105. mage_ai/data_cleaner/analysis/charts.py +61 -39
  106. mage_ai/data_cleaner/column_types/column_type_detector.py +53 -31
  107. mage_ai/data_cleaner/estimators/encoders.py +5 -2
  108. mage_ai/data_integrations/utils/scheduler.py +16 -11
  109. mage_ai/data_preparation/decorators.py +1 -0
  110. mage_ai/data_preparation/executors/block_executor.py +237 -155
  111. mage_ai/data_preparation/executors/k8s_block_executor.py +30 -7
  112. mage_ai/data_preparation/executors/k8s_pipeline_executor.py +30 -7
  113. mage_ai/data_preparation/executors/streaming_pipeline_executor.py +2 -2
  114. mage_ai/data_preparation/git/__init__.py +77 -29
  115. mage_ai/data_preparation/git/api.py +69 -8
  116. mage_ai/data_preparation/git/utils.py +64 -34
  117. mage_ai/data_preparation/logging/logger_manager.py +4 -3
  118. mage_ai/data_preparation/models/block/__init__.py +1562 -879
  119. mage_ai/data_preparation/models/block/data_integration/mixins.py +4 -3
  120. mage_ai/data_preparation/models/block/dynamic/__init__.py +17 -6
  121. mage_ai/data_preparation/models/block/dynamic/child.py +41 -102
  122. mage_ai/data_preparation/models/block/dynamic/constants.py +1 -0
  123. mage_ai/data_preparation/models/block/dynamic/counter.py +296 -0
  124. mage_ai/data_preparation/models/block/dynamic/data.py +16 -0
  125. mage_ai/data_preparation/models/block/dynamic/factory.py +163 -0
  126. mage_ai/data_preparation/models/block/dynamic/models.py +19 -0
  127. mage_ai/data_preparation/models/block/dynamic/shared.py +92 -0
  128. mage_ai/data_preparation/models/block/dynamic/utils.py +295 -167
  129. mage_ai/data_preparation/models/block/dynamic/variables.py +384 -144
  130. mage_ai/data_preparation/models/block/dynamic/wrappers.py +77 -0
  131. mage_ai/data_preparation/models/block/extension/utils.py +10 -1
  132. mage_ai/data_preparation/models/block/global_data_product/__init__.py +35 -3
  133. mage_ai/data_preparation/models/block/integration/__init__.py +6 -2
  134. mage_ai/data_preparation/models/block/outputs.py +722 -0
  135. mage_ai/data_preparation/models/block/platform/mixins.py +7 -8
  136. mage_ai/data_preparation/models/block/r/__init__.py +56 -38
  137. mage_ai/data_preparation/models/block/remote/__init__.py +0 -0
  138. mage_ai/data_preparation/models/block/remote/models.py +58 -0
  139. mage_ai/data_preparation/models/block/settings/__init__.py +0 -0
  140. mage_ai/data_preparation/models/block/settings/dynamic/__init__.py +0 -0
  141. mage_ai/data_preparation/models/block/settings/dynamic/constants.py +7 -0
  142. mage_ai/data_preparation/models/block/settings/dynamic/mixins.py +118 -0
  143. mage_ai/data_preparation/models/block/settings/dynamic/models.py +31 -0
  144. mage_ai/data_preparation/models/block/settings/global_data_products/__init__.py +0 -0
  145. mage_ai/data_preparation/models/block/settings/global_data_products/mixins.py +20 -0
  146. mage_ai/data_preparation/models/block/settings/global_data_products/models.py +46 -0
  147. mage_ai/data_preparation/models/block/settings/variables/__init__.py +0 -0
  148. mage_ai/data_preparation/models/block/settings/variables/mixins.py +74 -0
  149. mage_ai/data_preparation/models/block/settings/variables/models.py +49 -0
  150. mage_ai/data_preparation/models/block/spark/mixins.py +2 -1
  151. mage_ai/data_preparation/models/block/sql/__init__.py +30 -5
  152. mage_ai/data_preparation/models/block/sql/utils/shared.py +21 -3
  153. mage_ai/data_preparation/models/block/utils.py +164 -69
  154. mage_ai/data_preparation/models/constants.py +21 -14
  155. mage_ai/data_preparation/models/custom_templates/custom_block_template.py +18 -13
  156. mage_ai/data_preparation/models/custom_templates/custom_pipeline_template.py +33 -16
  157. mage_ai/data_preparation/models/custom_templates/utils.py +1 -1
  158. mage_ai/data_preparation/models/file.py +41 -28
  159. mage_ai/data_preparation/models/global_data_product/__init__.py +100 -58
  160. mage_ai/data_preparation/models/global_hooks/models.py +1 -0
  161. mage_ai/data_preparation/models/interfaces.py +29 -0
  162. mage_ai/data_preparation/models/pipeline.py +374 -185
  163. mage_ai/data_preparation/models/pipelines/integration_pipeline.py +1 -2
  164. mage_ai/data_preparation/models/pipelines/seed.py +1 -1
  165. mage_ai/data_preparation/models/project/__init__.py +66 -18
  166. mage_ai/data_preparation/models/project/constants.py +2 -0
  167. mage_ai/data_preparation/models/triggers/__init__.py +124 -26
  168. mage_ai/data_preparation/models/utils.py +467 -17
  169. mage_ai/data_preparation/models/variable.py +1028 -137
  170. mage_ai/data_preparation/models/variables/__init__.py +0 -0
  171. mage_ai/data_preparation/models/variables/cache.py +149 -0
  172. mage_ai/data_preparation/models/variables/constants.py +72 -0
  173. mage_ai/data_preparation/models/variables/summarizer.py +336 -0
  174. mage_ai/data_preparation/models/variables/utils.py +77 -0
  175. mage_ai/data_preparation/models/widget/__init__.py +63 -41
  176. mage_ai/data_preparation/models/widget/charts.py +40 -27
  177. mage_ai/data_preparation/models/widget/constants.py +2 -0
  178. mage_ai/data_preparation/models/widget/utils.py +3 -3
  179. mage_ai/data_preparation/preferences.py +3 -3
  180. mage_ai/data_preparation/repo_manager.py +55 -21
  181. mage_ai/data_preparation/storage/base_storage.py +2 -2
  182. mage_ai/data_preparation/storage/gcs_storage.py +7 -4
  183. mage_ai/data_preparation/storage/local_storage.py +18 -9
  184. mage_ai/data_preparation/storage/s3_storage.py +5 -2
  185. mage_ai/data_preparation/templates/data_exporters/streaming/oracledb.yaml +8 -0
  186. mage_ai/data_preparation/variable_manager.py +281 -76
  187. mage_ai/io/base.py +3 -2
  188. mage_ai/io/bigquery.py +1 -0
  189. mage_ai/io/redshift.py +7 -5
  190. mage_ai/kernels/__init__.py +0 -0
  191. mage_ai/kernels/models.py +188 -0
  192. mage_ai/kernels/utils.py +169 -0
  193. mage_ai/orchestration/concurrency.py +6 -2
  194. mage_ai/orchestration/db/__init__.py +1 -0
  195. mage_ai/orchestration/db/migrations/versions/0227396a216c_add_userproject_table.py +38 -0
  196. mage_ai/orchestration/db/migrations/versions/42a14d6143f1_update_token_column_type.py +54 -0
  197. mage_ai/orchestration/db/models/dynamic/__init__.py +0 -0
  198. mage_ai/orchestration/db/models/dynamic/controller.py +67 -0
  199. mage_ai/orchestration/db/models/oauth.py +12 -18
  200. mage_ai/orchestration/db/models/projects.py +10 -0
  201. mage_ai/orchestration/db/models/schedules.py +225 -187
  202. mage_ai/orchestration/db/models/schedules_project_platform.py +18 -12
  203. mage_ai/orchestration/db/models/utils.py +46 -5
  204. mage_ai/orchestration/metrics/pipeline_run.py +8 -9
  205. mage_ai/orchestration/notification/sender.py +38 -15
  206. mage_ai/orchestration/pipeline_scheduler_original.py +64 -33
  207. mage_ai/orchestration/pipeline_scheduler_project_platform.py +1 -1
  208. mage_ai/orchestration/run_status_checker.py +11 -4
  209. mage_ai/orchestration/triggers/api.py +41 -2
  210. mage_ai/orchestration/triggers/global_data_product.py +9 -4
  211. mage_ai/orchestration/triggers/utils.py +10 -1
  212. mage_ai/orchestration/utils/resources.py +3 -0
  213. mage_ai/presenters/charts/data_sources/base.py +4 -2
  214. mage_ai/presenters/charts/data_sources/block.py +15 -9
  215. mage_ai/presenters/charts/data_sources/chart_code.py +8 -5
  216. mage_ai/presenters/charts/data_sources/constants.py +1 -0
  217. mage_ai/presenters/charts/data_sources/system_metrics.py +22 -0
  218. mage_ai/presenters/interactions/models.py +11 -7
  219. mage_ai/presenters/pages/loaders/pipelines.py +5 -3
  220. mage_ai/presenters/pages/models/page_components/pipeline_schedules.py +3 -1
  221. mage_ai/presenters/utils.py +2 -0
  222. mage_ai/server/api/blocks.py +2 -1
  223. mage_ai/server/api/downloads.py +9 -2
  224. mage_ai/server/api/runs.py +151 -0
  225. mage_ai/server/api/triggers.py +3 -1
  226. mage_ai/server/constants.py +1 -1
  227. mage_ai/server/frontend_dist/404.html +8 -8
  228. mage_ai/server/frontend_dist/_next/static/UZLabyPgcxtZvp0O0EUUS/_buildManifest.js +1 -0
  229. mage_ai/server/frontend_dist/_next/static/chunks/1376-22de38b4ad008d8a.js +1 -0
  230. mage_ai/server/frontend_dist/_next/static/chunks/1557-25a7d985d5564fd3.js +1 -0
  231. mage_ai/server/frontend_dist/_next/static/chunks/1668-30b4619b9534519b.js +1 -0
  232. mage_ai/server/frontend_dist/_next/static/chunks/1799-c42db95a015689ee.js +1 -0
  233. mage_ai/server/frontend_dist/_next/static/chunks/2996-2108b53b9d371d8d.js +1 -0
  234. mage_ai/server/frontend_dist/_next/static/chunks/{3548-fa0792ddb88f4646.js → 3548-9d26185b3fb663b1.js} +1 -1
  235. mage_ai/server/frontend_dist/_next/static/chunks/{3763-61b542dafdbf5754.js → 3763-40780c6d1e4b261d.js} +1 -1
  236. mage_ai/server/frontend_dist/_next/static/chunks/3782-129dd2a2448a2e36.js +1 -0
  237. mage_ai/server/frontend_dist/_next/static/chunks/3958-bcdfa414ccfa1eb2.js +1 -0
  238. mage_ai/server/frontend_dist/_next/static/chunks/4168-97fd1578d1a38315.js +1 -0
  239. mage_ai/server/frontend_dist/_next/static/chunks/4982-fa5a238b139fbdd2.js +1 -0
  240. mage_ai/server/frontend_dist/_next/static/chunks/5699-176f445e1313f001.js +1 -0
  241. mage_ai/server/frontend_dist/_next/static/chunks/7162-7dd03f0f605de721.js +1 -0
  242. mage_ai/server/frontend_dist/_next/static/chunks/7779-68d2b72a90c5f925.js +1 -0
  243. mage_ai/server/frontend_dist/_next/static/chunks/7966-5446a8e43711e2f9.js +1 -0
  244. mage_ai/server/frontend_dist/_next/static/chunks/8023-6c2f172f48dcb99b.js +1 -0
  245. mage_ai/server/frontend_dist/_next/static/chunks/8095-c351b8a735d73e0c.js +1 -0
  246. mage_ai/server/frontend_dist/_next/static/chunks/9624-8b8e100079ab69e1.js +1 -0
  247. mage_ai/server/frontend_dist/_next/static/chunks/{main-77fe248a6fbd12d8.js → main-b99d4e30a88d9dc7.js} +1 -1
  248. mage_ai/server/frontend_dist/_next/static/chunks/pages/_app-9fe2d9d07c94e968.js +1 -0
  249. mage_ai/server/frontend_dist/_next/static/chunks/pages/{block-layout-14f952f66964022f.js → block-layout-7f4b735c67115df5.js} +1 -1
  250. mage_ai/server/frontend_dist/_next/static/chunks/pages/global-data-products/[...slug]-e7d48e6b0c3068ac.js +1 -0
  251. mage_ai/server/frontend_dist/_next/static/chunks/pages/global-data-products-b943f31f050fc3a4.js +1 -0
  252. mage_ai/server/frontend_dist/_next/static/chunks/pages/manage-4bfc84ff07d7656f.js +1 -0
  253. mage_ai/server/frontend_dist/_next/static/chunks/pages/{overview-597b74828bf105db.js → overview-9f1ac4ec003884f3.js} +1 -1
  254. mage_ai/server/{frontend_dist_base_path_template/_next/static/chunks/pages/pipeline-runs-3edc6270c5b0e962.js → frontend_dist/_next/static/chunks/pages/pipeline-runs-6d183f91a2ff6668.js} +1 -1
  255. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/backfills/[...slug]-7e737f6fc7e83e9b.js +1 -0
  256. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/backfills-38e1fbcfbfc1014e.js +1 -0
  257. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/dashboard-d94488e3f2eeef36.js +1 -0
  258. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-cc641a7fa8473796.js +1 -0
  259. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors/{block-runs-a5c0362763a21fa8.js → block-runs-284309877f3c5a5a.js} +1 -1
  260. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runtime-26250e5335194ade.js +1 -0
  261. mage_ai/server/{frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/monitors-7acc7afc00df17c2.js → frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors-5f4c8128b2413fd8.js} +1 -1
  262. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/runs/[run]-4ebfc8e400315dda.js +1 -0
  263. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/settings-e5e0150a256aadb3.js +1 -0
  264. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/triggers/[...slug]-eb11c5390c982b49.js +1 -0
  265. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/{triggers-1bdfda8edc9cf4a8.js → triggers-4612d15a65c35912.js} +1 -1
  266. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/account/{profile-3f0df3decc856ee9.js → profile-3ae43c932537b254.js} +1 -1
  267. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/platform/preferences-b603d7fe4b175256.js +1 -0
  268. mage_ai/server/{frontend_dist_base_path_template/_next/static/chunks/pages/settings/platform/settings-c2e9ef989c8bfa73.js → frontend_dist/_next/static/chunks/pages/settings/platform/settings-319ddbabc239e91b.js} +1 -1
  269. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/permissions/{[...slug]-47b64ced27c24985.js → [...slug]-5c360f72e4498855.js} +1 -1
  270. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/{permissions-e5a4d3d815cec25d.js → permissions-fb29fa6c2bd90bb0.js} +1 -1
  271. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/preferences-3b76fa959ffa09d3.js +1 -0
  272. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/roles/{[...slug]-379e1ee292504842.js → [...slug]-3b787b42f1093b1f.js} +1 -1
  273. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/roles-0b83fbdd39e85f5b.js +1 -0
  274. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/sync-data-a1e6950974d643a8.js +1 -0
  275. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/users/{[...slug]-2af9afbe727d88aa.js → [...slug]-0aa019d87db8b0b8.js} +1 -1
  276. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/{users-a4db8710f703c729.js → users-88c694d19207f2ec.js} +1 -1
  277. mage_ai/server/frontend_dist/_next/static/chunks/pages/{triggers-9cba3211434a8966.js → triggers-a599c6ac89be8c8d.js} +1 -1
  278. mage_ai/server/frontend_dist/_next/static/chunks/pages/version-control-31d0d50f7f30462b.js +1 -0
  279. mage_ai/server/frontend_dist/_next/static/chunks/{webpack-d079359c241db804.js → webpack-ac7fdc472bedf682.js} +1 -1
  280. mage_ai/server/frontend_dist/block-layout.html +3 -3
  281. mage_ai/server/frontend_dist/compute.html +6 -6
  282. mage_ai/server/frontend_dist/files.html +6 -6
  283. mage_ai/server/frontend_dist/global-data-products/[...slug].html +6 -6
  284. mage_ai/server/frontend_dist/global-data-products.html +6 -6
  285. mage_ai/server/frontend_dist/global-hooks/[...slug].html +6 -6
  286. mage_ai/server/frontend_dist/global-hooks.html +6 -6
  287. mage_ai/server/frontend_dist/index.html +3 -3
  288. mage_ai/server/frontend_dist/manage/files.html +6 -6
  289. mage_ai/server/frontend_dist/manage/settings.html +6 -6
  290. mage_ai/server/frontend_dist/manage/users/[user].html +6 -6
  291. mage_ai/server/frontend_dist/manage/users/new.html +6 -6
  292. mage_ai/server/frontend_dist/manage/users.html +6 -6
  293. mage_ai/server/frontend_dist/manage.html +6 -6
  294. mage_ai/server/frontend_dist/oauth.html +5 -5
  295. mage_ai/server/frontend_dist/overview.html +6 -6
  296. mage_ai/server/frontend_dist/pipeline-runs.html +6 -6
  297. mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills/[...slug].html +6 -6
  298. mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills.html +6 -6
  299. mage_ai/server/frontend_dist/pipelines/[pipeline]/dashboard.html +6 -6
  300. mage_ai/server/frontend_dist/pipelines/[pipeline]/edit.html +3 -3
  301. mage_ai/server/frontend_dist/pipelines/[pipeline]/logs.html +6 -6
  302. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runs.html +6 -6
  303. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runtime.html +6 -6
  304. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors.html +6 -6
  305. mage_ai/server/frontend_dist/pipelines/[pipeline]/runs/[run].html +6 -6
  306. mage_ai/server/frontend_dist/pipelines/[pipeline]/runs.html +6 -6
  307. mage_ai/server/frontend_dist/pipelines/[pipeline]/settings.html +6 -6
  308. mage_ai/server/frontend_dist/pipelines/[pipeline]/syncs.html +6 -6
  309. mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers/[...slug].html +6 -6
  310. mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers.html +6 -6
  311. mage_ai/server/frontend_dist/pipelines/[pipeline].html +3 -3
  312. mage_ai/server/frontend_dist/pipelines.html +6 -6
  313. mage_ai/server/frontend_dist/platform/global-hooks/[...slug].html +6 -6
  314. mage_ai/server/frontend_dist/platform/global-hooks.html +6 -6
  315. mage_ai/server/frontend_dist/settings/account/profile.html +6 -6
  316. mage_ai/server/frontend_dist/settings/platform/preferences.html +6 -6
  317. mage_ai/server/frontend_dist/settings/platform/settings.html +6 -6
  318. mage_ai/server/frontend_dist/settings/workspace/permissions/[...slug].html +6 -6
  319. mage_ai/server/frontend_dist/settings/workspace/permissions.html +6 -6
  320. mage_ai/server/frontend_dist/settings/workspace/preferences.html +6 -6
  321. mage_ai/server/frontend_dist/settings/workspace/roles/[...slug].html +6 -6
  322. mage_ai/server/frontend_dist/settings/workspace/roles.html +6 -6
  323. mage_ai/server/frontend_dist/settings/workspace/sync-data.html +6 -6
  324. mage_ai/server/frontend_dist/settings/workspace/users/[...slug].html +6 -6
  325. mage_ai/server/frontend_dist/settings/workspace/users.html +6 -6
  326. mage_ai/server/frontend_dist/settings.html +3 -3
  327. mage_ai/server/frontend_dist/sign-in.html +15 -15
  328. mage_ai/server/frontend_dist/templates/[...slug].html +6 -6
  329. mage_ai/server/frontend_dist/templates.html +6 -6
  330. mage_ai/server/frontend_dist/terminal.html +6 -6
  331. mage_ai/server/frontend_dist/test.html +3 -3
  332. mage_ai/server/frontend_dist/triggers.html +6 -6
  333. mage_ai/server/frontend_dist/version-control.html +6 -6
  334. mage_ai/server/frontend_dist_base_path_template/404.html +8 -8
  335. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/1376-22de38b4ad008d8a.js +1 -0
  336. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/1557-25a7d985d5564fd3.js +1 -0
  337. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/1668-30b4619b9534519b.js +1 -0
  338. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/1799-c42db95a015689ee.js +1 -0
  339. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/2996-2108b53b9d371d8d.js +1 -0
  340. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/{3548-fa0792ddb88f4646.js → 3548-9d26185b3fb663b1.js} +1 -1
  341. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/{3763-61b542dafdbf5754.js → 3763-40780c6d1e4b261d.js} +1 -1
  342. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/3782-129dd2a2448a2e36.js +1 -0
  343. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/3958-bcdfa414ccfa1eb2.js +1 -0
  344. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/4168-97fd1578d1a38315.js +1 -0
  345. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/4982-fa5a238b139fbdd2.js +1 -0
  346. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/5699-176f445e1313f001.js +1 -0
  347. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/7162-7dd03f0f605de721.js +1 -0
  348. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/7779-68d2b72a90c5f925.js +1 -0
  349. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/7966-5446a8e43711e2f9.js +1 -0
  350. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/8023-6c2f172f48dcb99b.js +1 -0
  351. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/8095-c351b8a735d73e0c.js +1 -0
  352. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/9624-8b8e100079ab69e1.js +1 -0
  353. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/{main-70b78159c2bb3fe1.js → main-384298e9133cec76.js} +1 -1
  354. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/_app-13a578bce3b7f30c.js +1 -0
  355. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/{block-layout-14f952f66964022f.js → block-layout-7f4b735c67115df5.js} +1 -1
  356. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/global-data-products/[...slug]-e7d48e6b0c3068ac.js +1 -0
  357. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/global-data-products-b943f31f050fc3a4.js +1 -0
  358. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/manage-4bfc84ff07d7656f.js +1 -0
  359. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/{overview-597b74828bf105db.js → overview-9f1ac4ec003884f3.js} +1 -1
  360. mage_ai/server/{frontend_dist/_next/static/chunks/pages/pipeline-runs-3edc6270c5b0e962.js → frontend_dist_base_path_template/_next/static/chunks/pages/pipeline-runs-6d183f91a2ff6668.js} +1 -1
  361. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/backfills/[...slug]-7e737f6fc7e83e9b.js +1 -0
  362. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/backfills-38e1fbcfbfc1014e.js +1 -0
  363. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/dashboard-d94488e3f2eeef36.js +1 -0
  364. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/edit-cc641a7fa8473796.js +1 -0
  365. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/monitors/{block-runs-a5c0362763a21fa8.js → block-runs-284309877f3c5a5a.js} +1 -1
  366. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runtime-26250e5335194ade.js +1 -0
  367. mage_ai/server/{frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors-7acc7afc00df17c2.js → frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/monitors-5f4c8128b2413fd8.js} +1 -1
  368. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/runs/[run]-4ebfc8e400315dda.js +1 -0
  369. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/settings-e5e0150a256aadb3.js +1 -0
  370. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/triggers/[...slug]-eb11c5390c982b49.js +1 -0
  371. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/{triggers-1bdfda8edc9cf4a8.js → triggers-4612d15a65c35912.js} +1 -1
  372. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/account/{profile-3f0df3decc856ee9.js → profile-3ae43c932537b254.js} +1 -1
  373. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/platform/preferences-b603d7fe4b175256.js +1 -0
  374. mage_ai/server/{frontend_dist/_next/static/chunks/pages/settings/platform/settings-c2e9ef989c8bfa73.js → frontend_dist_base_path_template/_next/static/chunks/pages/settings/platform/settings-319ddbabc239e91b.js} +1 -1
  375. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/permissions/{[...slug]-47b64ced27c24985.js → [...slug]-5c360f72e4498855.js} +1 -1
  376. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/{permissions-e5a4d3d815cec25d.js → permissions-fb29fa6c2bd90bb0.js} +1 -1
  377. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/preferences-3b76fa959ffa09d3.js +1 -0
  378. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/roles/{[...slug]-379e1ee292504842.js → [...slug]-3b787b42f1093b1f.js} +1 -1
  379. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/roles-0b83fbdd39e85f5b.js +1 -0
  380. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/sync-data-a1e6950974d643a8.js +1 -0
  381. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/users/{[...slug]-2af9afbe727d88aa.js → [...slug]-0aa019d87db8b0b8.js} +1 -1
  382. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/{users-a4db8710f703c729.js → users-88c694d19207f2ec.js} +1 -1
  383. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/{triggers-9cba3211434a8966.js → triggers-a599c6ac89be8c8d.js} +1 -1
  384. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/version-control-31d0d50f7f30462b.js +1 -0
  385. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/{webpack-68c003fb6a175cd7.js → webpack-481689d9989710cd.js} +1 -1
  386. mage_ai/server/frontend_dist_base_path_template/_next/static/kcptwoOU-JJJg6Vwpkfmx/_buildManifest.js +1 -0
  387. mage_ai/server/frontend_dist_base_path_template/block-layout.html +3 -3
  388. mage_ai/server/frontend_dist_base_path_template/compute.html +6 -6
  389. mage_ai/server/frontend_dist_base_path_template/files.html +6 -6
  390. mage_ai/server/frontend_dist_base_path_template/global-data-products/[...slug].html +6 -6
  391. mage_ai/server/frontend_dist_base_path_template/global-data-products.html +6 -6
  392. mage_ai/server/frontend_dist_base_path_template/global-hooks/[...slug].html +6 -6
  393. mage_ai/server/frontend_dist_base_path_template/global-hooks.html +6 -6
  394. mage_ai/server/frontend_dist_base_path_template/index.html +3 -3
  395. mage_ai/server/frontend_dist_base_path_template/manage/files.html +6 -6
  396. mage_ai/server/frontend_dist_base_path_template/manage/settings.html +6 -6
  397. mage_ai/server/frontend_dist_base_path_template/manage/users/[user].html +6 -6
  398. mage_ai/server/frontend_dist_base_path_template/manage/users/new.html +6 -6
  399. mage_ai/server/frontend_dist_base_path_template/manage/users.html +6 -6
  400. mage_ai/server/frontend_dist_base_path_template/manage.html +6 -6
  401. mage_ai/server/frontend_dist_base_path_template/oauth.html +5 -5
  402. mage_ai/server/frontend_dist_base_path_template/overview.html +6 -6
  403. mage_ai/server/frontend_dist_base_path_template/pipeline-runs.html +6 -6
  404. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/backfills/[...slug].html +6 -6
  405. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/backfills.html +6 -6
  406. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/dashboard.html +6 -6
  407. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/edit.html +3 -3
  408. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/logs.html +6 -6
  409. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/monitors/block-runs.html +6 -6
  410. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/monitors/block-runtime.html +6 -6
  411. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/monitors.html +6 -6
  412. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/runs/[run].html +6 -6
  413. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/runs.html +6 -6
  414. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/settings.html +6 -6
  415. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/syncs.html +6 -6
  416. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/triggers/[...slug].html +6 -6
  417. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline]/triggers.html +6 -6
  418. mage_ai/server/frontend_dist_base_path_template/pipelines/[pipeline].html +3 -3
  419. mage_ai/server/frontend_dist_base_path_template/pipelines.html +6 -6
  420. mage_ai/server/frontend_dist_base_path_template/platform/global-hooks/[...slug].html +6 -6
  421. mage_ai/server/frontend_dist_base_path_template/platform/global-hooks.html +6 -6
  422. mage_ai/server/frontend_dist_base_path_template/settings/account/profile.html +6 -6
  423. mage_ai/server/frontend_dist_base_path_template/settings/platform/preferences.html +6 -6
  424. mage_ai/server/frontend_dist_base_path_template/settings/platform/settings.html +6 -6
  425. mage_ai/server/frontend_dist_base_path_template/settings/workspace/permissions/[...slug].html +6 -6
  426. mage_ai/server/frontend_dist_base_path_template/settings/workspace/permissions.html +6 -6
  427. mage_ai/server/frontend_dist_base_path_template/settings/workspace/preferences.html +6 -6
  428. mage_ai/server/frontend_dist_base_path_template/settings/workspace/roles/[...slug].html +6 -6
  429. mage_ai/server/frontend_dist_base_path_template/settings/workspace/roles.html +6 -6
  430. mage_ai/server/frontend_dist_base_path_template/settings/workspace/sync-data.html +6 -6
  431. mage_ai/server/frontend_dist_base_path_template/settings/workspace/users/[...slug].html +6 -6
  432. mage_ai/server/frontend_dist_base_path_template/settings/workspace/users.html +6 -6
  433. mage_ai/server/frontend_dist_base_path_template/settings.html +3 -3
  434. mage_ai/server/frontend_dist_base_path_template/sign-in.html +15 -15
  435. mage_ai/server/frontend_dist_base_path_template/templates/[...slug].html +6 -6
  436. mage_ai/server/frontend_dist_base_path_template/templates.html +6 -6
  437. mage_ai/server/frontend_dist_base_path_template/terminal.html +6 -6
  438. mage_ai/server/frontend_dist_base_path_template/test.html +3 -3
  439. mage_ai/server/frontend_dist_base_path_template/triggers.html +6 -6
  440. mage_ai/server/frontend_dist_base_path_template/version-control.html +6 -6
  441. mage_ai/server/kernel_output_parser.py +4 -1
  442. mage_ai/server/scheduler_manager.py +12 -1
  443. mage_ai/server/server.py +69 -42
  444. mage_ai/server/utils/custom_output.py +284 -0
  445. mage_ai/server/utils/execute_custom_code.py +245 -0
  446. mage_ai/server/utils/output_display.py +123 -289
  447. mage_ai/server/websocket_server.py +116 -69
  448. mage_ai/services/aws/ecs/ecs.py +1 -0
  449. mage_ai/services/k8s/config.py +27 -4
  450. mage_ai/services/k8s/job_manager.py +6 -1
  451. mage_ai/services/k8s/utils.py +97 -0
  452. mage_ai/services/ssh/aws/emr/utils.py +8 -8
  453. mage_ai/settings/keys/auth.py +1 -0
  454. mage_ai/settings/platform/__init__.py +159 -38
  455. mage_ai/settings/platform/constants.py +5 -0
  456. mage_ai/settings/platform/utils.py +53 -10
  457. mage_ai/settings/repo.py +26 -12
  458. mage_ai/settings/server.py +128 -37
  459. mage_ai/shared/array.py +24 -1
  460. mage_ai/shared/complex.py +45 -0
  461. mage_ai/shared/config.py +2 -1
  462. mage_ai/shared/custom_logger.py +11 -0
  463. mage_ai/shared/dates.py +10 -6
  464. mage_ai/shared/files.py +63 -8
  465. mage_ai/shared/hash.py +33 -9
  466. mage_ai/shared/io.py +9 -5
  467. mage_ai/shared/models.py +82 -24
  468. mage_ai/shared/outputs.py +87 -0
  469. mage_ai/shared/parsers.py +144 -13
  470. mage_ai/shared/path_fixer.py +11 -7
  471. mage_ai/shared/singletons/__init__.py +0 -0
  472. mage_ai/shared/singletons/base.py +47 -0
  473. mage_ai/shared/singletons/memory.py +38 -0
  474. mage_ai/shared/strings.py +34 -1
  475. mage_ai/shared/yaml.py +24 -0
  476. mage_ai/streaming/sinks/oracledb.py +57 -0
  477. mage_ai/streaming/sinks/sink_factory.py +4 -0
  478. mage_ai/system/__init__.py +0 -0
  479. mage_ai/system/constants.py +14 -0
  480. mage_ai/system/memory/__init__.py +0 -0
  481. mage_ai/system/memory/constants.py +1 -0
  482. mage_ai/system/memory/manager.py +174 -0
  483. mage_ai/system/memory/presenters.py +158 -0
  484. mage_ai/system/memory/process.py +216 -0
  485. mage_ai/system/memory/samples.py +13 -0
  486. mage_ai/system/memory/utils.py +656 -0
  487. mage_ai/system/memory/wrappers.py +177 -0
  488. mage_ai/system/models.py +58 -0
  489. mage_ai/system/storage/__init__.py +0 -0
  490. mage_ai/system/storage/utils.py +29 -0
  491. mage_ai/tests/api/endpoints/mixins.py +2 -2
  492. mage_ai/tests/api/endpoints/test_blocks.py +2 -1
  493. mage_ai/tests/api/endpoints/test_custom_designs.py +4 -4
  494. mage_ai/tests/api/endpoints/test_pipeline_runs.py +2 -2
  495. mage_ai/tests/api/endpoints/test_projects.py +2 -1
  496. mage_ai/tests/api/operations/base/mixins.py +1 -1
  497. mage_ai/tests/api/operations/base/test_base.py +27 -27
  498. mage_ai/tests/api/operations/base/test_base_with_user_authentication.py +27 -27
  499. mage_ai/tests/api/operations/base/test_base_with_user_permissions.py +23 -23
  500. mage_ai/tests/api/operations/test_syncs.py +6 -4
  501. mage_ai/tests/api/resources/test_pipeline_resource.py +11 -4
  502. mage_ai/tests/authentication/oauth/test_utils.py +1 -1
  503. mage_ai/tests/authentication/providers/test_oidc.py +59 -0
  504. mage_ai/tests/base_test.py +2 -2
  505. mage_ai/tests/data/__init__.py +0 -0
  506. mage_ai/tests/data/models/__init__.py +0 -0
  507. mage_ai/tests/data_preparation/executors/test_block_executor.py +23 -16
  508. mage_ai/tests/data_preparation/git/test_git.py +4 -1
  509. mage_ai/tests/data_preparation/models/block/dynamic/test_combos.py +305 -0
  510. mage_ai/tests/data_preparation/models/block/dynamic/test_counter.py +212 -0
  511. mage_ai/tests/data_preparation/models/block/dynamic/test_factory.py +360 -0
  512. mage_ai/tests/data_preparation/models/block/dynamic/test_variables.py +332 -0
  513. mage_ai/tests/data_preparation/models/block/hook/test_hook_block.py +2 -2
  514. mage_ai/tests/data_preparation/models/block/platform/test_mixins.py +1 -1
  515. mage_ai/tests/data_preparation/models/block/sql/utils/test_shared.py +26 -1
  516. mage_ai/tests/data_preparation/models/block/test_global_data_product.py +5 -2
  517. mage_ai/tests/data_preparation/models/custom_templates/test_utils.py +5 -4
  518. mage_ai/tests/data_preparation/models/global_hooks/test_hook.py +3 -0
  519. mage_ai/tests/data_preparation/models/global_hooks/test_predicates.py +9 -3
  520. mage_ai/tests/data_preparation/models/test_block.py +115 -120
  521. mage_ai/tests/data_preparation/models/test_blocks_helper.py +114 -0
  522. mage_ai/tests/data_preparation/models/test_global_data_product.py +41 -24
  523. mage_ai/tests/data_preparation/models/test_pipeline.py +9 -6
  524. mage_ai/tests/data_preparation/models/test_project.py +4 -1
  525. mage_ai/tests/data_preparation/models/test_utils.py +80 -0
  526. mage_ai/tests/data_preparation/models/test_variable.py +242 -69
  527. mage_ai/tests/data_preparation/models/variables/__init__.py +0 -0
  528. mage_ai/tests/data_preparation/models/variables/test_summarizer.py +481 -0
  529. mage_ai/tests/data_preparation/storage/shared/__init__.py +0 -0
  530. mage_ai/tests/data_preparation/test_repo_manager.py +6 -7
  531. mage_ai/tests/data_preparation/test_variable_manager.py +57 -48
  532. mage_ai/tests/factory.py +64 -43
  533. mage_ai/tests/orchestration/db/models/test_schedules.py +3 -3
  534. mage_ai/tests/orchestration/db/models/test_schedules_dynamic_blocks.py +279 -0
  535. mage_ai/tests/orchestration/test_pipeline_scheduler.py +1 -0
  536. mage_ai/tests/orchestration/triggers/test_global_data_product.py +141 -138
  537. mage_ai/tests/orchestration/triggers/test_utils.py +3 -2
  538. mage_ai/tests/server/test_server.py +19 -0
  539. mage_ai/tests/services/k8s/test_job_manager.py +27 -6
  540. mage_ai/tests/streaming/sinks/test_oracledb.py +38 -0
  541. mage_ai/tests/test_shared.py +61 -0
  542. mage_ai/usage_statistics/logger.py +7 -2
  543. mage_ai/utils/code.py +33 -19
  544. mage_ai/version_control/branch/utils.py +2 -1
  545. mage_ai/version_control/models.py +3 -2
  546. {mage_ai-0.9.69.dist-info → mage_ai-0.9.71.dist-info}/METADATA +6 -3
  547. {mage_ai-0.9.69.dist-info → mage_ai-0.9.71.dist-info}/RECORD +555 -454
  548. mage_ai/data_preparation/models/global_data_product/constants.py +0 -6
  549. mage_ai/server/frontend_dist/_next/static/_krrrgup_C-dPOpX36S8I/_buildManifest.js +0 -1
  550. mage_ai/server/frontend_dist/_next/static/chunks/1557-df144fbd8b2208c3.js +0 -1
  551. mage_ai/server/frontend_dist/_next/static/chunks/2631-b9f9bea3f1cf906d.js +0 -1
  552. mage_ai/server/frontend_dist/_next/static/chunks/3782-ef4cd4f0b52072d0.js +0 -1
  553. mage_ai/server/frontend_dist/_next/static/chunks/4783-422429203610c318.js +0 -1
  554. mage_ai/server/frontend_dist/_next/static/chunks/5699-6d708c6b2153ea08.js +0 -1
  555. mage_ai/server/frontend_dist/_next/static/chunks/635-0d6b7c8804bcd2dc.js +0 -1
  556. mage_ai/server/frontend_dist/_next/static/chunks/7022-0d52dd8868621fb0.js +0 -1
  557. mage_ai/server/frontend_dist/_next/static/chunks/7361-8a23dd8360593e7a.js +0 -1
  558. mage_ai/server/frontend_dist/_next/static/chunks/7966-f07b2913f7326b50.js +0 -1
  559. mage_ai/server/frontend_dist/_next/static/chunks/8095-bdce03896ef9639a.js +0 -1
  560. mage_ai/server/frontend_dist/_next/static/chunks/8146-6bed4e7401e067e6.js +0 -1
  561. mage_ai/server/frontend_dist/_next/static/chunks/9265-d2a1aaec75ec69b8.js +0 -1
  562. mage_ai/server/frontend_dist/_next/static/chunks/9440-4069842b90d4b801.js +0 -1
  563. mage_ai/server/frontend_dist/_next/static/chunks/9624-59b2f803f9c88cd6.js +0 -1
  564. mage_ai/server/frontend_dist/_next/static/chunks/9832-67896490f6e8a014.js +0 -1
  565. mage_ai/server/frontend_dist/_next/static/chunks/pages/_app-d9c89527266296f7.js +0 -1
  566. mage_ai/server/frontend_dist/_next/static/chunks/pages/global-data-products/[...slug]-591abd392dc50ed4.js +0 -1
  567. mage_ai/server/frontend_dist/_next/static/chunks/pages/global-data-products-78e8e88f2a757a18.js +0 -1
  568. mage_ai/server/frontend_dist/_next/static/chunks/pages/manage-852d403c7bda21b3.js +0 -1
  569. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/backfills/[...slug]-ff4bd7a8ec3bab40.js +0 -1
  570. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/backfills-a8b61d8d239fd16f.js +0 -1
  571. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/dashboard-95ffcd3e2b27e567.js +0 -1
  572. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-e1dd1ed71d26c10d.js +0 -1
  573. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runtime-1ed9045b2f1dfd65.js +0 -1
  574. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/runs/[run]-1417ad1c821d720a.js +0 -1
  575. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/settings-59aca25a5b1d3998.js +0 -1
  576. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/triggers/[...slug]-f028ef3880ed856c.js +0 -1
  577. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/platform/preferences-503049734a8b082f.js +0 -1
  578. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/preferences-5b26eeda8aed8a7b.js +0 -1
  579. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/roles-36fa165a48af586b.js +0 -1
  580. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/sync-data-8b793b3b696a2cd3.js +0 -1
  581. mage_ai/server/frontend_dist/_next/static/chunks/pages/version-control-5753fac7c1bfdc88.js +0 -1
  582. mage_ai/server/frontend_dist_base_path_template/_next/static/KLL5mirre9d7_ZeEpaw3s/_buildManifest.js +0 -1
  583. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/1557-df144fbd8b2208c3.js +0 -1
  584. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/2631-b9f9bea3f1cf906d.js +0 -1
  585. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/3782-ef4cd4f0b52072d0.js +0 -1
  586. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/4783-422429203610c318.js +0 -1
  587. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/5699-6d708c6b2153ea08.js +0 -1
  588. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/635-0d6b7c8804bcd2dc.js +0 -1
  589. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/7022-0d52dd8868621fb0.js +0 -1
  590. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/7361-8a23dd8360593e7a.js +0 -1
  591. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/7966-f07b2913f7326b50.js +0 -1
  592. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/8095-bdce03896ef9639a.js +0 -1
  593. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/8146-6bed4e7401e067e6.js +0 -1
  594. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/9265-d2a1aaec75ec69b8.js +0 -1
  595. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/9440-4069842b90d4b801.js +0 -1
  596. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/9624-59b2f803f9c88cd6.js +0 -1
  597. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/9832-67896490f6e8a014.js +0 -1
  598. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/_app-d9c89527266296f7.js +0 -1
  599. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/global-data-products/[...slug]-591abd392dc50ed4.js +0 -1
  600. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/global-data-products-78e8e88f2a757a18.js +0 -1
  601. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/manage-852d403c7bda21b3.js +0 -1
  602. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/backfills/[...slug]-ff4bd7a8ec3bab40.js +0 -1
  603. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/backfills-a8b61d8d239fd16f.js +0 -1
  604. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/dashboard-95ffcd3e2b27e567.js +0 -1
  605. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/edit-e1dd1ed71d26c10d.js +0 -1
  606. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runtime-1ed9045b2f1dfd65.js +0 -1
  607. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/runs/[run]-1417ad1c821d720a.js +0 -1
  608. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/settings-59aca25a5b1d3998.js +0 -1
  609. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/pipelines/[pipeline]/triggers/[...slug]-f028ef3880ed856c.js +0 -1
  610. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/platform/preferences-503049734a8b082f.js +0 -1
  611. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/preferences-5b26eeda8aed8a7b.js +0 -1
  612. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/roles-36fa165a48af586b.js +0 -1
  613. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/settings/workspace/sync-data-8b793b3b696a2cd3.js +0 -1
  614. mage_ai/server/frontend_dist_base_path_template/_next/static/chunks/pages/version-control-5753fac7c1bfdc88.js +0 -1
  615. mage_ai/shared/memory.py +0 -90
  616. mage_ai/tests/data_preparation/models/block/dynamic/test_dynamic_helpers.py +0 -48
  617. /mage_ai/{tests/data_preparation/shared → ai/utils}/__init__.py +0 -0
  618. /mage_ai/server/frontend_dist/_next/static/{_krrrgup_C-dPOpX36S8I → UZLabyPgcxtZvp0O0EUUS}/_ssgManifest.js +0 -0
  619. /mage_ai/server/frontend_dist_base_path_template/_next/static/{KLL5mirre9d7_ZeEpaw3s → kcptwoOU-JJJg6Vwpkfmx}/_ssgManifest.js +0 -0
  620. /mage_ai/tests/data_preparation/{shared → storage/shared}/test_secrets.py +0 -0
  621. {mage_ai-0.9.69.dist-info → mage_ai-0.9.71.dist-info}/LICENSE +0 -0
  622. {mage_ai-0.9.69.dist-info → mage_ai-0.9.71.dist-info}/WHEEL +0 -0
  623. {mage_ai-0.9.69.dist-info → mage_ai-0.9.71.dist-info}/entry_points.txt +0 -0
  624. {mage_ai-0.9.69.dist-info → mage_ai-0.9.71.dist-info}/top_level.txt +0 -0
@@ -14,7 +14,18 @@ from inspect import Parameter, isfunction, signature
14
14
  from logging import Logger
15
15
  from pathlib import Path
16
16
  from queue import Queue
17
- from typing import Any, Callable, Dict, Generator, List, Set, Tuple, Union
17
+ from typing import (
18
+ Any,
19
+ Callable,
20
+ Dict,
21
+ Generator,
22
+ Iterable,
23
+ List,
24
+ Optional,
25
+ Set,
26
+ Tuple,
27
+ Union,
28
+ )
18
29
 
19
30
  import inflection
20
31
  import pandas as pd
@@ -25,7 +36,8 @@ from jinja2 import Template
25
36
 
26
37
  import mage_ai.data_preparation.decorators
27
38
  from mage_ai.cache.block import BlockCache
28
- from mage_ai.data_cleaner.shared.utils import is_geo_dataframe, is_spark_dataframe
39
+ from mage_ai.data.constants import InputDataType
40
+ from mage_ai.data.tabular.models import BatchSettings
29
41
  from mage_ai.data_integrations.sources.constants import SQL_SOURCES_MAPPING
30
42
  from mage_ai.data_preparation.logging.logger import DictLogger
31
43
  from mage_ai.data_preparation.logging.logger_manager_factory import LoggerManagerFactory
@@ -43,6 +55,7 @@ from mage_ai.data_preparation.models.block.dynamic.utils import (
43
55
  uuid_for_output_variables,
44
56
  )
45
57
  from mage_ai.data_preparation.models.block.dynamic.variables import (
58
+ LazyVariableSet,
46
59
  delete_variable_objects_for_dynamic_child,
47
60
  fetch_input_variables_for_dynamic_upstream_blocks,
48
61
  get_outputs_for_dynamic_block,
@@ -51,10 +64,26 @@ from mage_ai.data_preparation.models.block.dynamic.variables import (
51
64
  )
52
65
  from mage_ai.data_preparation.models.block.errors import HasDownstreamDependencies
53
66
  from mage_ai.data_preparation.models.block.extension.utils import handle_run_tests
67
+ from mage_ai.data_preparation.models.block.outputs import (
68
+ format_output_data,
69
+ get_outputs_for_display_async,
70
+ get_outputs_for_display_dynamic_block,
71
+ get_outputs_for_display_sync,
72
+ )
54
73
  from mage_ai.data_preparation.models.block.platform.mixins import (
55
74
  ProjectPlatformAccessible,
56
75
  )
57
76
  from mage_ai.data_preparation.models.block.platform.utils import from_another_project
77
+ from mage_ai.data_preparation.models.block.settings.dynamic.mixins import DynamicMixin
78
+ from mage_ai.data_preparation.models.block.settings.global_data_products.mixins import (
79
+ GlobalDataProductsMixin,
80
+ )
81
+ from mage_ai.data_preparation.models.block.settings.variables.mixins import (
82
+ VariablesMixin,
83
+ )
84
+ from mage_ai.data_preparation.models.block.settings.variables.models import (
85
+ ChunkKeyTypeUnion,
86
+ )
58
87
  from mage_ai.data_preparation.models.block.spark.mixins import SparkBlock
59
88
  from mage_ai.data_preparation.models.block.utils import (
60
89
  clean_name,
@@ -64,13 +93,14 @@ from mage_ai.data_preparation.models.block.utils import (
64
93
  is_valid_print_variable,
65
94
  output_variables,
66
95
  )
67
- from mage_ai.data_preparation.models.constants import (
96
+ from mage_ai.data_preparation.models.constants import ( # PIPELINES_FOLDER,
68
97
  BLOCK_LANGUAGE_TO_FILE_EXTENSION,
69
98
  CALLBACK_STATUSES,
70
99
  CUSTOM_EXECUTION_BLOCK_TYPES,
71
100
  DATAFRAME_ANALYSIS_MAX_COLUMNS,
72
101
  DATAFRAME_ANALYSIS_MAX_ROWS,
73
102
  DATAFRAME_SAMPLE_COUNT_PREVIEW,
103
+ DYNAMIC_CHILD_BLOCK_SAMPLE_COUNT_PREVIEW,
74
104
  FILE_EXTENSION_TO_BLOCK_LANGUAGE,
75
105
  NON_PIPELINE_EXECUTABLE_BLOCK_TYPES,
76
106
  BlockColor,
@@ -82,18 +112,36 @@ from mage_ai.data_preparation.models.constants import (
82
112
  PipelineType,
83
113
  )
84
114
  from mage_ai.data_preparation.models.file import File
85
- from mage_ai.data_preparation.models.variable import VariableType
115
+ from mage_ai.data_preparation.models.utils import is_basic_iterable, warn_for_repo_path
116
+ from mage_ai.data_preparation.models.variable import Variable
117
+ from mage_ai.data_preparation.models.variables.cache import (
118
+ AggregateInformation,
119
+ AggregateInformationData,
120
+ InformationData,
121
+ VariableAggregateCache,
122
+ )
123
+ from mage_ai.data_preparation.models.variables.constants import (
124
+ VariableAggregateDataType,
125
+ VariableAggregateSummaryGroupType,
126
+ VariableType,
127
+ )
128
+ from mage_ai.data_preparation.models.variables.summarizer import (
129
+ aggregate_summary_info_for_all_variables,
130
+ get_aggregate_summary_info,
131
+ )
86
132
  from mage_ai.data_preparation.repo_manager import RepoConfig
87
133
  from mage_ai.data_preparation.shared.stream import StreamToLogger
88
134
  from mage_ai.data_preparation.shared.utils import get_template_vars
89
135
  from mage_ai.data_preparation.templates.data_integrations.utils import get_templates
90
136
  from mage_ai.data_preparation.templates.template import load_template
91
- from mage_ai.server.kernel_output_parser import DataType
137
+ from mage_ai.data_preparation.variable_manager import VariableManager
138
+ from mage_ai.io.base import ExportWritePolicy
92
139
  from mage_ai.services.spark.config import SparkConfig
93
140
  from mage_ai.services.spark.spark import SPARK_ENABLED, get_spark_session
94
141
  from mage_ai.settings.platform.constants import project_platform_activated
95
- from mage_ai.settings.repo import get_repo_path
96
- from mage_ai.shared.array import unique_by
142
+ from mage_ai.settings.repo import base_repo_path_directory_name, get_repo_path
143
+ from mage_ai.settings.server import VARIABLE_DATA_OUTPUT_META_CACHE
144
+ from mage_ai.shared.array import is_iterable, unique_by
97
145
  from mage_ai.shared.constants import ENV_DEV, ENV_TEST
98
146
  from mage_ai.shared.custom_logger import DX_PRINTER
99
147
  from mage_ai.shared.environments import get_env, is_debug
@@ -110,6 +158,10 @@ from mage_ai.shared.strings import format_enum
110
158
  from mage_ai.shared.utils import clean_name as clean_name_orig
111
159
  from mage_ai.shared.utils import is_spark_env
112
160
 
161
+ # from mage_ai.system.memory.manager import MemoryManager
162
+ from mage_ai.system.memory.wrappers import execute_with_memory_tracking
163
+ from mage_ai.system.models import ResourceUsage
164
+
113
165
  PYTHON_COMMAND = 'python3'
114
166
  BLOCK_EXISTS_ERROR = '[ERR_BLOCK_EXISTS]'
115
167
 
@@ -174,8 +226,11 @@ async def run_blocks(
174
226
  while not blocks.empty():
175
227
  block = blocks.get()
176
228
 
177
- if block.type in NON_PIPELINE_EXECUTABLE_BLOCK_TYPES or \
178
- not run_sensors and block.type == BlockType.SENSOR:
229
+ if (
230
+ block.type in NON_PIPELINE_EXECUTABLE_BLOCK_TYPES
231
+ or not run_sensors
232
+ and block.type == BlockType.SENSOR
233
+ ):
179
234
  continue
180
235
 
181
236
  if tries_by_block_uuid.get(block.uuid, None) is None:
@@ -290,7 +345,14 @@ def run_blocks_sync(
290
345
  blocks.put(downstream_block)
291
346
 
292
347
 
293
- class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
348
+ class Block(
349
+ DataIntegrationMixin,
350
+ SparkBlock,
351
+ ProjectPlatformAccessible,
352
+ DynamicMixin,
353
+ GlobalDataProductsMixin,
354
+ VariablesMixin,
355
+ ):
294
356
  def __init__(
295
357
  self,
296
358
  name: str,
@@ -376,6 +438,13 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
376
438
  # Needs to after self._project_platform_activated = None
377
439
  self.configuration = configuration
378
440
 
441
+ self.resource_usage = None
442
+ self._store_variables_in_block_function: Optional[
443
+ Callable[..., Optional[List[VariableType]]]
444
+ ] = None
445
+
446
+ self._variable_aggregate_cache = None
447
+
379
448
  @property
380
449
  def uuid(self) -> str:
381
450
  return self._uuid
@@ -449,6 +518,166 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
449
518
  )
450
519
  return None
451
520
 
521
+ @property
522
+ def variable_manager(self) -> VariableManager:
523
+ if not self.pipeline:
524
+ return
525
+ return self.pipeline.variable_manager
526
+
527
+ @property
528
+ def pipeline_uuid(self) -> str:
529
+ return self.pipeline.uuid if self.pipeline else ''
530
+
531
+ def __load_variable_aggregate_cache(self, variable_uuid: str) -> VariableAggregateCache:
532
+ if not self._variable_aggregate_cache:
533
+ self._variable_aggregate_cache = {variable_uuid: VariableAggregateCache()}
534
+
535
+ if variable_uuid not in self._variable_aggregate_cache:
536
+ self._variable_aggregate_cache[variable_uuid] = VariableAggregateCache()
537
+
538
+ return self._variable_aggregate_cache[variable_uuid]
539
+
540
+ def get_variable_aggregate_cache(
541
+ self,
542
+ variable_uuid: str,
543
+ data_type: VariableAggregateDataType,
544
+ default_group_type: Optional[VariableAggregateSummaryGroupType] = None,
545
+ group_type: Optional[VariableAggregateSummaryGroupType] = None,
546
+ infer_group_type: Optional[bool] = None,
547
+ partition: Optional[str] = None,
548
+ ) -> Optional[Union[AggregateInformationData, InformationData]]:
549
+ if not VARIABLE_DATA_OUTPUT_META_CACHE:
550
+ return
551
+
552
+ cache = self.__load_variable_aggregate_cache(variable_uuid)
553
+ cache = VariableAggregateCache.load(cache)
554
+
555
+ if infer_group_type:
556
+ group_type = (
557
+ VariableAggregateSummaryGroupType.DYNAMIC
558
+ if is_dynamic_block_child(self)
559
+ else VariableAggregateSummaryGroupType.PARTS
560
+ )
561
+
562
+ keys = [v.value for v in [group_type, data_type] if v is not None]
563
+ value = functools.reduce(getattr, keys, cache)
564
+
565
+ if not value:
566
+ cache_new = get_aggregate_summary_info(
567
+ self.variable_manager,
568
+ self.pipeline_uuid,
569
+ self.uuid,
570
+ variable_uuid,
571
+ data_type,
572
+ default_group_type=default_group_type,
573
+ group_type=group_type,
574
+ partition=partition,
575
+ )
576
+ group_value_use = (group_type.value if group_type else None) or (
577
+ default_group_type.value if default_group_type else None
578
+ )
579
+ if group_value_use is not None:
580
+ cache_group = AggregateInformation.load(getattr(cache, group_value_use))
581
+ cache_group_new = AggregateInformation.load(getattr(cache_new, group_value_use))
582
+ if cache_group_new:
583
+ for data in VariableAggregateDataType:
584
+ val = getattr(cache_group, data.value)
585
+ val_new = getattr(cache_group_new, data.value)
586
+ cache_group.update_attributes(**{
587
+ data.value: val_new or val,
588
+ })
589
+ cache.update_attributes(**{
590
+ group_value_use: AggregateInformation.load(cache_group)
591
+ })
592
+
593
+ for data in VariableAggregateDataType:
594
+ val = getattr(cache, data.value)
595
+ val_new = getattr(cache_new, data.value)
596
+ cache.update_attributes(**{
597
+ data.value: val_new or val,
598
+ })
599
+
600
+ cache = VariableAggregateCache.load(cache)
601
+ self._variable_aggregate_cache = merge_dict(
602
+ self._variable_aggregate_cache or {},
603
+ {variable_uuid: cache},
604
+ )
605
+ value = functools.reduce(getattr, keys, cache)
606
+
607
+ return value
608
+
609
+ def get_resource_usage(
610
+ self,
611
+ block_uuid: Optional[str] = None,
612
+ index: Optional[int] = None,
613
+ partition: Optional[str] = None,
614
+ variable_uuid: Optional[str] = None,
615
+ ) -> Optional[ResourceUsage]:
616
+ try:
617
+ if not VARIABLE_DATA_OUTPUT_META_CACHE:
618
+ variable = self.get_variable_object(
619
+ block_uuid or self.uuid, partition=partition, variable_uuid=variable_uuid
620
+ )
621
+ return variable.get_resource_usage(index=index)
622
+
623
+ values = self.get_variable_aggregate_cache(
624
+ variable_uuid,
625
+ VariableAggregateDataType.RESOURCE_USAGE,
626
+ infer_group_type=index is not None,
627
+ partition=partition,
628
+ )
629
+
630
+ if index is not None:
631
+ if values and isinstance(values, list) and len(values) > index:
632
+ values = values[index]
633
+ else:
634
+ values = values
635
+
636
+ if isinstance(values, Iterable) and len(values) >= 1:
637
+ values = values[0]
638
+
639
+ return values
640
+ except Exception as err:
641
+ print(f'[ERROR] Block.get_resource_usage: {err}')
642
+ return ResourceUsage()
643
+
644
+ def get_analysis(
645
+ self,
646
+ block_uuid: Optional[str] = None,
647
+ index: Optional[int] = None,
648
+ partition: Optional[str] = None,
649
+ variable_uuid: Optional[str] = None,
650
+ ) -> Optional[Dict]:
651
+ try:
652
+ if not VARIABLE_DATA_OUTPUT_META_CACHE:
653
+ variable = self.get_variable_object(
654
+ block_uuid or self.uuid, partition=partition, variable_uuid=variable_uuid
655
+ )
656
+ return variable.get_analysis(index=index)
657
+
658
+ values = self.get_variable_aggregate_cache(
659
+ variable_uuid,
660
+ VariableAggregateDataType.STATISTICS,
661
+ infer_group_type=index is not None,
662
+ partition=partition,
663
+ )
664
+
665
+ value = None
666
+ if index is not None:
667
+ if values and isinstance(values, list) and len(values) > index:
668
+ value = values[index]
669
+ else:
670
+ value = values
671
+
672
+ if isinstance(value, Iterable) and len(value) >= 1:
673
+ value = value[0]
674
+
675
+ if value is not None:
676
+ return dict(statistics=value.to_dict() if value else {})
677
+ except Exception as err:
678
+ print(f'[ERROR] Block.get_analysis: {err}')
679
+ return {}
680
+
452
681
  async def content_async(self) -> str:
453
682
  if self.replicated_block and self.replicated_block_object:
454
683
  self._content = await self.replicated_block_object.content_async()
@@ -483,32 +712,42 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
483
712
  upstream_block_uuids = self.upstream_block_uuids
484
713
 
485
714
  if outputs_from_input_vars is None:
486
-
487
- if BlockLanguage.SQL == self.language and any([is_dynamic_block(
488
- upstream_block,
489
- ) or is_dynamic_block_child(
490
- upstream_block,
491
- ) for upstream_block in self.upstream_blocks]):
492
- outputs_from_input_vars, _kwargs_vars, upstream_block_uuids = \
493
- fetch_input_variables_for_dynamic_upstream_blocks(
494
- self,
495
- None,
496
- dynamic_block_index=dynamic_block_index,
497
- dynamic_block_indexes=dynamic_block_indexes,
498
- execution_partition=execution_partition,
499
- from_notebook=from_notebook,
500
- global_vars=variables,
501
- )
715
+ if BlockLanguage.SQL == self.language and any([
716
+ is_dynamic_block(
717
+ upstream_block,
718
+ )
719
+ or is_dynamic_block_child(
720
+ upstream_block,
721
+ )
722
+ for upstream_block in self.upstream_blocks
723
+ ]):
724
+ (
725
+ outputs_from_input_vars,
726
+ _kwargs_vars,
727
+ upstream_block_uuids,
728
+ ) = fetch_input_variables_for_dynamic_upstream_blocks(
729
+ self,
730
+ None,
731
+ dynamic_block_index=dynamic_block_index,
732
+ dynamic_block_indexes=dynamic_block_indexes,
733
+ execution_partition=execution_partition,
734
+ from_notebook=from_notebook,
735
+ global_vars=variables,
736
+ )
502
737
  else:
503
- outputs_from_input_vars, _input_vars, _kwargs_vars, upstream_block_uuids = \
504
- self.__get_outputs_from_input_vars(
505
- dynamic_block_index=dynamic_block_index,
506
- dynamic_block_indexes=dynamic_block_indexes,
507
- dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
508
- execution_partition=execution_partition,
509
- from_notebook=from_notebook,
510
- global_vars=variables,
511
- )
738
+ (
739
+ outputs_from_input_vars,
740
+ _input_vars,
741
+ _kwargs_vars,
742
+ upstream_block_uuids,
743
+ ) = self.__get_outputs_from_input_vars(
744
+ dynamic_block_index=dynamic_block_index,
745
+ dynamic_block_indexes=dynamic_block_indexes,
746
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
747
+ execution_partition=execution_partition,
748
+ from_notebook=from_notebook,
749
+ global_vars=variables,
750
+ )
512
751
 
513
752
  return hydrate_block_outputs(
514
753
  content,
@@ -524,7 +763,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
524
763
  if BlockLanguage.YAML == self.language:
525
764
  content = await self.content_async()
526
765
  if content:
527
- text = self.interpolate_content(content)
766
+ try:
767
+ text = self.interpolate_content(content)
768
+ except Exception:
769
+ traceback.print_exc()
770
+ text = content
528
771
  settings = yaml.safe_load(text)
529
772
  uuid = settings.get('source') or settings.get('destination')
530
773
  mapping = grouped_templates.get(uuid) or {}
@@ -550,11 +793,14 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
550
793
 
551
794
  di_metadata = merge_dict(
552
795
  extract(mapping or {}, ['name']),
553
- ignore_keys(di_settings or {}, [
554
- 'catalog',
555
- 'config',
556
- 'data_integration_uuid',
557
- ]),
796
+ ignore_keys(
797
+ di_settings or {},
798
+ [
799
+ 'catalog',
800
+ 'config',
801
+ 'data_integration_uuid',
802
+ ],
803
+ ),
558
804
  )
559
805
  di_metadata['sql'] = uuid in SQL_SOURCES_MAPPING
560
806
 
@@ -572,9 +818,8 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
572
818
 
573
819
  @property
574
820
  def executable(self) -> bool:
575
- return (
576
- self.type not in NON_PIPELINE_EXECUTABLE_BLOCK_TYPES
577
- and (self.pipeline is None or self.pipeline.type != PipelineType.STREAMING)
821
+ return self.type not in NON_PIPELINE_EXECUTABLE_BLOCK_TYPES and (
822
+ self.pipeline is None or self.pipeline.type != PipelineType.STREAMING
578
823
  )
579
824
 
580
825
  @property
@@ -584,10 +829,15 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
584
829
  self._outputs = self.get_outputs()
585
830
  return self._outputs
586
831
 
587
- async def __outputs_async(self) -> List:
832
+ async def __outputs_async(
833
+ self, exclude_blank_variable_uuids: bool = False, max_results: Optional[int] = None
834
+ ) -> List:
588
835
  if not self._outputs_loaded:
589
836
  if self._outputs is None or len(self._outputs) == 0:
590
- self._outputs = await self.__get_outputs_async()
837
+ self._outputs = await self.__get_outputs_async(
838
+ exclude_blank_variable_uuids=exclude_blank_variable_uuids,
839
+ max_results=max_results,
840
+ )
591
841
  return self._outputs
592
842
 
593
843
  @property
@@ -622,15 +872,20 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
622
872
  block_uuid: str,
623
873
  block_type: BlockType,
624
874
  language: BlockLanguage,
875
+ relative_path: bool = False,
625
876
  ) -> str:
626
877
  file_extension = BLOCK_LANGUAGE_TO_FILE_EXTENSION[language]
627
878
  block_directory = f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
628
879
 
629
- return os.path.join(
630
- repo_path or os.getcwd(),
880
+ parts = []
881
+ if not relative_path:
882
+ parts.append(repo_path or os.getcwd())
883
+ parts += [
631
884
  block_directory,
632
885
  f'{block_uuid}.{file_extension}',
633
- )
886
+ ]
887
+
888
+ return os.path.join(*parts)
634
889
 
635
890
  @property
636
891
  def file_path(self) -> str:
@@ -645,12 +900,59 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
645
900
  return add_absolute_path(file_path)
646
901
 
647
902
  return self.__build_file_path(
648
- self.repo_path or os.getcwd(),
903
+ self.repo_path,
649
904
  self.uuid,
650
905
  self.type,
651
906
  self.language,
652
907
  )
653
908
 
909
+ def build_file_path_directory(
910
+ self,
911
+ block_uuid: Optional[str] = None,
912
+ ) -> Tuple[Optional[str], Optional[str]]:
913
+ file_path = None
914
+ file_path_absolute = None
915
+
916
+ if self.replicated_block and self.replicated_block_object:
917
+ (
918
+ file_path_absolute,
919
+ file_path,
920
+ ) = self.replicated_block_object.build_file_path_directory(
921
+ block_uuid=block_uuid,
922
+ )
923
+
924
+ if not file_path:
925
+ file_path = self.get_file_path_from_source() or self.configuration.get('file_path')
926
+
927
+ if not file_path:
928
+ file_path = self.__build_file_path(
929
+ self.repo_path or os.getcwd(),
930
+ self.uuid,
931
+ self.type,
932
+ self.language,
933
+ relative_path=True,
934
+ )
935
+
936
+ if block_uuid:
937
+ old_file_path = Path(file_path)
938
+ old_file_extension = old_file_path.suffix
939
+ new_file_name = f'{block_uuid}{old_file_extension}'
940
+ file_path = str(old_file_path.with_name(new_file_name))
941
+
942
+ if file_path:
943
+ file_path_absolute = add_absolute_path(file_path)
944
+
945
+ if not file_path_absolute and block_uuid:
946
+ file_path_absolute = self.__build_file_path(
947
+ self.repo_path or os.getcwd(),
948
+ block_uuid,
949
+ self.type,
950
+ self.language,
951
+ relative_path=False,
952
+ )
953
+
954
+ return file_path_absolute, file_path
955
+
654
956
  @property
655
957
  def file(self) -> File:
656
958
  if self.replicated_block and self.replicated_block_object:
@@ -661,15 +963,29 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
661
963
  if file:
662
964
  return file
663
965
 
664
- return File.from_path(self.file_path)
966
+ repo_path = self.pipeline.repo_path if self.pipeline else None
967
+ new_file = File.from_path(self.file_path, repo_path=repo_path)
968
+
969
+ if not new_file.filename or not new_file.dir_path:
970
+ new_file = File.from_path(
971
+ self.__build_file_path(
972
+ new_file.repo_path,
973
+ self.uuid,
974
+ self.type,
975
+ self.language,
976
+ )
977
+ )
978
+
979
+ return new_file
665
980
 
666
981
  @property
667
982
  def table_name(self) -> str:
668
983
  if self.configuration and self.configuration.get('data_provider_table'):
669
984
  return self.configuration['data_provider_table']
670
985
 
671
- table_name = f'{self.pipeline.uuid}_{clean_name_orig(self.uuid)}_'\
672
- f'{self.pipeline.version_name}'
986
+ table_name = (
987
+ f'{self.pipeline_uuid}_{clean_name_orig(self.uuid)}_' f'{self.pipeline.version_name}'
988
+ )
673
989
 
674
990
  env = (self.global_vars or dict()).get('env')
675
991
  if env == ENV_DEV:
@@ -680,31 +996,15 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
680
996
  return table_name
681
997
 
682
998
  @property
683
- def full_table_name(self) -> str:
999
+ def full_table_name(self) -> Optional[str]:
684
1000
  from mage_ai.data_preparation.models.block.sql.utils.shared import (
685
- extract_create_statement_table_name,
686
- extract_insert_statement_table_names,
687
- extract_update_statement_table_names,
1001
+ extract_full_table_name,
688
1002
  )
689
1003
 
690
- if not self.content:
691
- return None
692
-
693
- table_name = extract_create_statement_table_name(self.content)
694
- if table_name:
695
- return table_name
696
-
697
- matches = extract_insert_statement_table_names(self.content)
698
- if len(matches) == 0:
699
- matches = extract_update_statement_table_names(self.content)
700
-
701
- if len(matches) == 0:
702
- return None
703
-
704
- return matches[len(matches) - 1]
1004
+ return extract_full_table_name(self.content)
705
1005
 
706
1006
  @classmethod
707
- def after_create(self, block: 'Block', **kwargs) -> None:
1007
+ def after_create(cls, block: 'Block', **kwargs) -> None:
708
1008
  widget = kwargs.get('widget')
709
1009
  pipeline = kwargs.get('pipeline')
710
1010
  if pipeline is not None:
@@ -766,28 +1066,30 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
766
1066
  # Don’t create a file if it’s from another project.
767
1067
 
768
1068
  file_path_from_source = (
769
- configuration and
770
- configuration.get('file_source') and
771
- (configuration.get('file_source') or {}).get('path')
1069
+ configuration
1070
+ and configuration.get('file_source')
1071
+ and (configuration.get('file_source') or {}).get('path')
772
1072
  )
773
- file_is_from_another_project = (
774
- file_path_from_source and
775
- from_another_project(
776
- file_path=file_path_from_source,
777
- other_file_path=pipeline.dir_path if pipeline else None,
1073
+ file_is_from_another_project = file_path_from_source and from_another_project(
1074
+ file_path=file_path_from_source,
1075
+ other_file_path=pipeline.dir_path if pipeline else None,
1076
+ )
1077
+ absolute_file_path = (
1078
+ add_root_repo_path_to_relative_path(
1079
+ file_path_from_source,
778
1080
  )
1081
+ if file_path_from_source
1082
+ else None
779
1083
  )
780
- absolute_file_path = add_root_repo_path_to_relative_path(
781
- file_path_from_source,
782
- ) if file_path_from_source else None
783
-
784
- if not file_is_from_another_project and \
785
- (not absolute_file_path or not os.path.exists(absolute_file_path)):
786
-
787
- if not replicated_block and \
788
- (BlockType.DBT != block_type or BlockLanguage.YAML == language) and \
789
- BlockType.GLOBAL_DATA_PRODUCT != block_type:
790
1084
 
1085
+ if not file_is_from_another_project and (
1086
+ not absolute_file_path or not os.path.exists(absolute_file_path)
1087
+ ):
1088
+ if (
1089
+ not replicated_block
1090
+ and (BlockType.DBT != block_type or BlockLanguage.YAML == language)
1091
+ and BlockType.GLOBAL_DATA_PRODUCT != block_type
1092
+ ):
791
1093
  block_directory = self.file_directory_name(block_type)
792
1094
  if absolute_file_path:
793
1095
  block_dir_path = os.path.dirname(absolute_file_path)
@@ -803,19 +1105,24 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
803
1105
  file_path = os.path.join(block_dir_path, f'{uuid}.{file_extension}')
804
1106
  if os.path.exists(file_path):
805
1107
  already_exists = True
806
- if (pipeline is not None and pipeline.has_block(
807
- uuid,
808
- block_type=block_type,
809
- extension_uuid=extension_uuid,
810
- )) or require_unique_name:
1108
+ if (
1109
+ pipeline is not None
1110
+ and pipeline.has_block(
1111
+ uuid,
1112
+ block_type=block_type,
1113
+ extension_uuid=extension_uuid,
1114
+ )
1115
+ ) or require_unique_name:
811
1116
  """
812
1117
  The BLOCK_EXISTS_ERROR constant is used on the frontend to identify when
813
1118
  a user is trying to create a new block with an existing block name, and
814
1119
  link them to the existing block file so the user can choose to add the
815
1120
  existing block to their pipeline.
816
1121
  """
817
- raise Exception(f'{BLOCK_EXISTS_ERROR} Block {uuid} already exists. \
818
- Please use a different name.')
1122
+ raise Exception(
1123
+ f'{BLOCK_EXISTS_ERROR} Block {uuid} already exists. \
1124
+ Please use a different name.'
1125
+ )
819
1126
  else:
820
1127
  load_template(
821
1128
  block_type,
@@ -830,8 +1137,9 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
830
1137
  if not configuration.get('file_source'):
831
1138
  configuration['file_source'] = {}
832
1139
  if not configuration['file_source'].get('path'):
1140
+ relative_path = str(Path(repo_path).relative_to(base_repo_path_directory_name()))
833
1141
  configuration['file_source']['path'] = self.__build_file_path(
834
- get_repo_path(absolute_path=False, root_project=False),
1142
+ relative_path,
835
1143
  uuid,
836
1144
  block_type,
837
1145
  language,
@@ -875,8 +1183,13 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
875
1183
  return f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
876
1184
 
877
1185
  @classmethod
878
- def block_type_from_path(self, block_file_absolute_path: str) -> BlockType:
879
- file_path = str(block_file_absolute_path).replace(get_repo_path(), '')
1186
+ def block_type_from_path(
1187
+ self, block_file_absolute_path: str, repo_path: str = None
1188
+ ) -> BlockType:
1189
+ warn_for_repo_path(repo_path)
1190
+ if not repo_path:
1191
+ repo_path = get_repo_path()
1192
+ file_path = str(block_file_absolute_path).replace(repo_path, '')
880
1193
  if file_path.startswith(os.sep):
881
1194
  file_path = file_path[1:]
882
1195
 
@@ -981,9 +1294,13 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
981
1294
  )
982
1295
  # For block_type SCRATCHPAD and MARKDOWN, also delete the file if possible
983
1296
  if self.type in NON_PIPELINE_EXECUTABLE_BLOCK_TYPES:
984
- pipelines = Pipeline.get_pipelines_by_block(self, widget=widget)
1297
+ pipelines = Pipeline.get_pipelines_by_block(
1298
+ self,
1299
+ repo_path=self.pipeline.repo_path,
1300
+ widget=widget,
1301
+ )
985
1302
  pipelines = [
986
- pipeline for pipeline in pipelines if self.pipeline.uuid != pipeline.uuid
1303
+ pipeline for pipeline in pipelines if self.pipeline_uuid != pipeline.uuid
987
1304
  ]
988
1305
  if len(pipelines) == 0:
989
1306
  os.remove(self.file_path)
@@ -991,7 +1308,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
991
1308
 
992
1309
  # TODO (tommy dang): delete this block from all pipelines in all projects
993
1310
  # If pipeline is not specified, delete the block from all pipelines and delete the file.
994
- pipelines = Pipeline.get_pipelines_by_block(self, widget=widget)
1311
+ pipelines = Pipeline.get_pipelines_by_block(self, repo_path=self.repo_path, widget=widget)
995
1312
  if not force:
996
1313
  for p in pipelines:
997
1314
  if not p.block_deletable(self):
@@ -1017,7 +1334,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1017
1334
  logging_tags: Dict = None,
1018
1335
  execution_uuid: str = None,
1019
1336
  from_notebook: bool = False,
1020
- **kwargs
1337
+ **kwargs,
1021
1338
  ) -> Dict:
1022
1339
  """
1023
1340
  This method will execute the block and run the callback functions if they exist
@@ -1051,9 +1368,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1051
1368
  # Print result to block output
1052
1369
  if not result:
1053
1370
  conditional_message += 'This block would not be executed in a trigger run.\n'
1054
- conditional_json = json.dumps(dict(
1055
- message=conditional_message,
1056
- ))
1371
+ conditional_json = json.dumps(
1372
+ dict(
1373
+ message=conditional_message,
1374
+ )
1375
+ )
1057
1376
  print(f'[__internal_test__]{conditional_json}')
1058
1377
 
1059
1378
  callback_arr = []
@@ -1068,7 +1387,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1068
1387
  logger=logger,
1069
1388
  logging_tags=logging_tags,
1070
1389
  from_notebook=from_notebook,
1071
- **kwargs
1390
+ **kwargs,
1072
1391
  )
1073
1392
  except Exception as error:
1074
1393
  for callback_block in callback_arr:
@@ -1122,145 +1441,235 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1122
1441
  data_integration_runtime_settings: Dict = None,
1123
1442
  execution_partition_previous: str = None,
1124
1443
  metadata: Dict = None,
1444
+ override_outputs: bool = True,
1125
1445
  **kwargs,
1126
1446
  ) -> Dict:
1127
- if logging_tags is None:
1128
- logging_tags = dict()
1447
+ def __execute(
1448
+ self=self,
1449
+ analyze_outputs=analyze_outputs,
1450
+ block_run_outputs_cache=block_run_outputs_cache,
1451
+ build_block_output_stdout=build_block_output_stdout,
1452
+ custom_code=custom_code,
1453
+ data_integration_runtime_settings=data_integration_runtime_settings,
1454
+ disable_json_serialization=disable_json_serialization,
1455
+ dynamic_block_index=dynamic_block_index,
1456
+ dynamic_block_indexes=dynamic_block_indexes,
1457
+ dynamic_block_uuid=dynamic_block_uuid,
1458
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
1459
+ execution_partition=execution_partition,
1460
+ execution_partition_previous=execution_partition_previous,
1461
+ from_notebook=from_notebook,
1462
+ global_vars=global_vars,
1463
+ input_from_output=input_from_output,
1464
+ kwargs=kwargs,
1465
+ logger=logger,
1466
+ logging_tags=logging_tags,
1467
+ metadata=metadata,
1468
+ output_messages_to_logs=output_messages_to_logs,
1469
+ override_outputs=override_outputs,
1470
+ run_all_blocks=run_all_blocks,
1471
+ run_settings=run_settings,
1472
+ runtime_arguments=runtime_arguments,
1473
+ store_variables=store_variables,
1474
+ update_status=update_status,
1475
+ verify_output=verify_output,
1476
+ ) -> Dict:
1477
+ if logging_tags is None:
1478
+ logging_tags = dict()
1129
1479
 
1130
- try:
1131
- if not run_all_blocks:
1132
- not_executed_upstream_blocks = list(
1133
- filter(lambda b: b.status == BlockStatus.NOT_EXECUTED, self.upstream_blocks)
1134
- )
1135
- all_upstream_is_dbt = all([BlockType.DBT == b.type
1136
- for b in not_executed_upstream_blocks])
1137
- if not all_upstream_is_dbt and len(not_executed_upstream_blocks) > 0:
1138
- upstream_block_uuids = list(map(lambda b: b.uuid, not_executed_upstream_blocks))
1139
- raise Exception(
1140
- f"Block {self.uuid}'s upstream blocks have not been executed yet. "
1141
- f'Please run upstream blocks {upstream_block_uuids} '
1142
- 'before running the current block.'
1480
+ try:
1481
+ if not run_all_blocks:
1482
+ not_executed_upstream_blocks = list(
1483
+ filter(
1484
+ lambda b: b.status == BlockStatus.NOT_EXECUTED,
1485
+ self.upstream_blocks,
1486
+ )
1143
1487
  )
1144
- global_vars = self.enrich_global_vars(
1145
- global_vars,
1146
- dynamic_block_index=dynamic_block_index,
1147
- )
1148
-
1149
- if output_messages_to_logs and not logger:
1150
- from mage_ai.data_preparation.models.block.constants import (
1151
- LOG_PARTITION_EDIT_PIPELINE,
1488
+ all_upstream_is_dbt = all([
1489
+ BlockType.DBT == b.type for b in not_executed_upstream_blocks
1490
+ ])
1491
+ if not all_upstream_is_dbt and len(not_executed_upstream_blocks) > 0:
1492
+ upstream_block_uuids = list(
1493
+ map(lambda b: b.uuid, not_executed_upstream_blocks)
1494
+ )
1495
+ raise Exception(
1496
+ f"Block {self.uuid}'s upstream blocks have not been executed yet. "
1497
+ f'Please run upstream blocks {upstream_block_uuids} '
1498
+ 'before running the current block.'
1499
+ )
1500
+ global_vars = self.enrich_global_vars(
1501
+ global_vars,
1502
+ dynamic_block_index=dynamic_block_index,
1152
1503
  )
1153
1504
 
1154
- logger_manager = LoggerManagerFactory.get_logger_manager(
1155
- block_uuid=datetime.utcnow().strftime(format='%Y%m%dT%H%M%S'),
1156
- partition=LOG_PARTITION_EDIT_PIPELINE,
1157
- pipeline_uuid=self.pipeline.uuid if self.pipeline else None,
1158
- subpartition=clean_name(self.uuid),
1159
- )
1160
- logger = DictLogger(logger_manager.logger)
1161
- logging_tags = dict(
1162
- block_type=self.type,
1163
- block_uuid=self.uuid,
1164
- pipeline_uuid=self.pipeline.uuid if self.pipeline else None,
1165
- )
1505
+ def __store_variables(
1506
+ variable_mapping: Dict[str, Any],
1507
+ skip_delete: Optional[bool] = None,
1508
+ save_variable_types_only: Optional[bool] = None,
1509
+ clean_variable_uuid: Optional[bool] = None,
1510
+ dynamic_block_index=dynamic_block_index,
1511
+ dynamic_block_uuid=dynamic_block_uuid,
1512
+ execution_partition=execution_partition,
1513
+ global_vars=global_vars,
1514
+ override_outputs=override_outputs,
1515
+ self=self,
1516
+ ) -> Optional[List[Variable]]:
1517
+ return self.store_variables(
1518
+ variable_mapping,
1519
+ clean_variable_uuid=clean_variable_uuid,
1520
+ execution_partition=execution_partition,
1521
+ override_outputs=override_outputs,
1522
+ skip_delete=skip_delete,
1523
+ spark=self.__get_spark_session_from_global_vars(
1524
+ global_vars=global_vars,
1525
+ ),
1526
+ dynamic_block_index=dynamic_block_index,
1527
+ dynamic_block_uuid=dynamic_block_uuid,
1528
+ save_variable_types_only=save_variable_types_only,
1529
+ )
1166
1530
 
1167
- output = self.execute_block(
1168
- block_run_outputs_cache=block_run_outputs_cache,
1169
- build_block_output_stdout=build_block_output_stdout,
1170
- custom_code=custom_code,
1171
- execution_partition=execution_partition,
1172
- from_notebook=from_notebook,
1173
- global_vars=global_vars,
1174
- logger=logger,
1175
- logging_tags=logging_tags,
1176
- input_from_output=input_from_output,
1177
- runtime_arguments=runtime_arguments,
1178
- dynamic_block_index=dynamic_block_index,
1179
- dynamic_block_indexes=dynamic_block_indexes,
1180
- dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
1181
- run_settings=run_settings,
1182
- data_integration_runtime_settings=data_integration_runtime_settings,
1183
- execution_partition_previous=execution_partition_previous,
1184
- metadata=metadata,
1185
- **kwargs,
1186
- )
1531
+ self._store_variables_in_block_function = __store_variables
1187
1532
 
1188
- if self.configuration and self.configuration.get('disable_query_preprocessing'):
1189
- output = dict(output=None)
1190
- else:
1191
- block_output = self.post_process_output(output)
1192
- variable_mapping = dict()
1193
-
1194
- if BlockType.CHART == self.type:
1195
- variable_mapping = block_output
1196
- output = dict(
1197
- output=simplejson.dumps(
1198
- block_output,
1199
- default=encode_complex,
1200
- ignore_nan=True,
1201
- ) if not disable_json_serialization else block_output,
1533
+ if output_messages_to_logs and not logger:
1534
+ from mage_ai.data_preparation.models.block.constants import (
1535
+ LOG_PARTITION_EDIT_PIPELINE,
1202
1536
  )
1203
- else:
1204
- output_count = len(block_output)
1205
- variable_keys = [f'output_{idx}' for idx in range(output_count)]
1206
- variable_mapping = dict(zip(variable_keys, block_output))
1207
1537
 
1208
- if store_variables and \
1209
- self.pipeline and \
1210
- self.pipeline.type != PipelineType.INTEGRATION:
1538
+ logger_manager = LoggerManagerFactory.get_logger_manager(
1539
+ block_uuid=datetime.utcnow().strftime(format='%Y%m%dT%H%M%S'),
1540
+ partition=LOG_PARTITION_EDIT_PIPELINE,
1541
+ pipeline_uuid=self.pipeline_uuid,
1542
+ subpartition=clean_name(self.uuid),
1543
+ )
1544
+ logger = DictLogger(logger_manager.logger)
1545
+ logging_tags = dict(
1546
+ block_type=self.type,
1547
+ block_uuid=self.uuid,
1548
+ pipeline_uuid=self.pipeline_uuid,
1549
+ )
1211
1550
 
1212
- try:
1213
- DX_PRINTER.critical(
1214
- block=self,
1215
- execution_partition=execution_partition,
1216
- override_outputs=True,
1217
- dynamic_block_uuid=dynamic_block_uuid,
1218
- __uuid='store_variables',
1219
- )
1551
+ output = self.execute_block(
1552
+ block_run_outputs_cache=block_run_outputs_cache,
1553
+ build_block_output_stdout=build_block_output_stdout,
1554
+ custom_code=custom_code,
1555
+ execution_partition=execution_partition,
1556
+ from_notebook=from_notebook,
1557
+ global_vars=global_vars,
1558
+ logger=logger,
1559
+ logging_tags=logging_tags,
1560
+ input_from_output=input_from_output,
1561
+ runtime_arguments=runtime_arguments,
1562
+ dynamic_block_index=dynamic_block_index,
1563
+ dynamic_block_indexes=dynamic_block_indexes,
1564
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
1565
+ run_settings=run_settings,
1566
+ data_integration_runtime_settings=data_integration_runtime_settings,
1567
+ execution_partition_previous=execution_partition_previous,
1568
+ metadata=metadata,
1569
+ override_outputs=override_outputs,
1570
+ **kwargs,
1571
+ )
1220
1572
 
1221
- self.store_variables(
1222
- variable_mapping,
1223
- execution_partition=execution_partition,
1224
- override_outputs=True,
1225
- spark=self.__get_spark_session_from_global_vars(
1226
- global_vars=global_vars,
1573
+ if self.configuration and self.configuration.get('disable_query_preprocessing'):
1574
+ output = dict(output=None)
1575
+ else:
1576
+ block_output = self.post_process_output(output)
1577
+ variable_mapping = dict()
1578
+
1579
+ if BlockType.CHART == self.type:
1580
+ variable_mapping = block_output
1581
+ output = dict(
1582
+ output=(
1583
+ simplejson.dumps(
1584
+ block_output,
1585
+ default=encode_complex,
1586
+ ignore_nan=True,
1587
+ )
1588
+ if not disable_json_serialization
1589
+ else block_output
1227
1590
  ),
1228
- dynamic_block_index=dynamic_block_index,
1229
- dynamic_block_uuid=dynamic_block_uuid,
1230
- )
1231
- except ValueError as e:
1232
- if str(e) == 'Circular reference detected':
1233
- raise ValueError(
1234
- 'Please provide dataframe or json serializable data as output.'
1235
- )
1236
- raise e
1237
- # Reset outputs cache
1238
- self._outputs = None
1239
-
1240
- if BlockType.CHART != self.type:
1241
- if analyze_outputs:
1242
- self.analyze_outputs(
1243
- variable_mapping,
1244
- execution_partition=execution_partition,
1245
1591
  )
1246
1592
  else:
1247
- self.analyze_outputs(
1248
- variable_mapping,
1249
- execution_partition=execution_partition,
1250
- shape_only=True,
1251
- )
1593
+ output_count = len(block_output)
1594
+ variable_keys = [f'output_{idx}' for idx in range(output_count)]
1595
+ variable_mapping = dict(zip(variable_keys, block_output))
1596
+
1597
+ if (
1598
+ store_variables
1599
+ and self.pipeline
1600
+ and self.pipeline.type != PipelineType.INTEGRATION
1601
+ ):
1602
+ try:
1603
+ DX_PRINTER.critical(
1604
+ block=self,
1605
+ execution_partition=execution_partition,
1606
+ override_outputs=override_outputs,
1607
+ dynamic_block_uuid=dynamic_block_uuid,
1608
+ __uuid='store_variables',
1609
+ )
1252
1610
 
1253
- if update_status:
1254
- self.status = BlockStatus.EXECUTED
1255
- except Exception as err:
1256
- if update_status:
1257
- self.status = BlockStatus.FAILED
1258
- raise err
1259
- finally:
1260
- if update_status:
1261
- self.__update_pipeline_block(widget=BlockType.CHART == self.type)
1611
+ if self._store_variables_in_block_function and isinstance(
1612
+ variable_mapping, dict
1613
+ ):
1614
+ self._store_variables_in_block_function(variable_mapping)
1615
+
1616
+ except ValueError as e:
1617
+ if str(e) == 'Circular reference detected':
1618
+ raise ValueError(
1619
+ 'Please provide dataframe or json serializable data as output.'
1620
+ )
1621
+ raise e
1622
+
1623
+ if not is_dynamic_block_child(self):
1624
+ # This will be handled in the execute_custom_code file so that it’s only
1625
+ # invoked once.
1626
+ self.aggregate_summary_info()
1627
+
1628
+ # Reset outputs cache
1629
+ self._outputs = None
1630
+
1631
+ if BlockType.CHART != self.type:
1632
+ if analyze_outputs:
1633
+ self.analyze_outputs(
1634
+ variable_mapping,
1635
+ execution_partition=execution_partition,
1636
+ )
1637
+ else:
1638
+ self.analyze_outputs(
1639
+ variable_mapping,
1640
+ execution_partition=execution_partition,
1641
+ shape_only=True,
1642
+ )
1262
1643
 
1263
- return output
1644
+ if update_status:
1645
+ self.status = BlockStatus.EXECUTED
1646
+ except Exception as err:
1647
+ if update_status:
1648
+ self.status = BlockStatus.FAILED
1649
+ raise err
1650
+ finally:
1651
+ if update_status:
1652
+ self.__update_pipeline_block(widget=BlockType.CHART == self.type)
1653
+
1654
+ return output
1655
+
1656
+ # if MEMORY_MANAGER_V2:
1657
+ # metadata = {}
1658
+ # if execution_partition:
1659
+ # metadata['execution_partition'] = execution_partition
1660
+ # if from_notebook:
1661
+ # metadata['origin'] = 'ide'
1662
+ # with MemoryManager(
1663
+ # scope_uuid=os.path.join(
1664
+ # *([PIPELINES_FOLDER, self.pipeline_uuid] if self.pipeline else ['']),
1665
+ # self.uuid,
1666
+ # ),
1667
+ # process_uuid='block.execute_sync',
1668
+ # repo_path=self.repo_path,
1669
+ # metadata=metadata,
1670
+ # ):
1671
+ # return __execute()
1672
+ return __execute()
1264
1673
 
1265
1674
  def post_process_output(self, output: Dict) -> List:
1266
1675
  return output['output'] or []
@@ -1289,7 +1698,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1289
1698
  global_vars=global_vars,
1290
1699
  run_all_blocks=run_all_blocks,
1291
1700
  update_status=update_status,
1292
- )
1701
+ ),
1293
1702
  )
1294
1703
  else:
1295
1704
  self.execute_sync(
@@ -1334,19 +1743,19 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1334
1743
  if num_args > num_inputs:
1335
1744
  if num_upstream < num_args:
1336
1745
  raise Exception(
1337
- f'Block {self.uuid} may be missing upstream dependencies. '
1746
+ f"Block {self.uuid} may be missing upstream dependencies. "
1338
1747
  f'It expected to have {"at least " if has_var_args else ""}{num_args} '
1339
- f'arguments, but only received {num_inputs}. '
1340
- f'Confirm that the @{self.type} method declaration has the correct number '
1341
- 'of arguments.'
1748
+ f"arguments, but only received {num_inputs}. "
1749
+ f"Confirm that the @{self.type} method declaration has the correct number "
1750
+ "of arguments."
1342
1751
  )
1343
1752
  else:
1344
1753
  raise Exception(
1345
- f'Block {self.uuid} is missing input arguments. '
1754
+ f"Block {self.uuid} is missing input arguments. "
1346
1755
  f'It expected to have {"at least " if has_var_args else ""}{num_args} '
1347
- f'arguments, but only received {num_inputs}. '
1348
- f'Double check the @{self.type} method declaration has the correct number '
1349
- 'of arguments and that the upstream blocks have been executed.'
1756
+ f"arguments, but only received {num_inputs}. "
1757
+ f"Double check the @{self.type} method declaration has the correct number "
1758
+ "of arguments and that the upstream blocks have been executed."
1350
1759
  )
1351
1760
  elif num_args < num_inputs and not has_var_args:
1352
1761
  if num_upstream > num_args:
@@ -1378,21 +1787,26 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1378
1787
  input_args: List = None,
1379
1788
  metadata: Dict = None,
1380
1789
  ) -> Tuple[Dict, List, Dict, List[str]]:
1381
- # Only fetch the input variables that the destination block explicitly declares.
1382
- # If all the input variables are fetched, there is a chance that a lot of data from
1383
- # an upstream source block is loaded just to be used as inputs for the block’s
1384
- # decorated functions. Only do this for the notebook because
1790
+ """
1791
+ Only fetch the input variables that the destination block explicitly declares.
1792
+ If all the input variables are fetched, there is a chance that a lot of data from
1793
+ an upstream source block is loaded just to be used as inputs for the block’s
1794
+ decorated functions. Only do this for the notebook because
1795
+ """
1385
1796
  if from_notebook and self.is_data_integration():
1386
- input_vars, kwargs_vars, upstream_block_uuids = \
1387
- self.fetch_input_variables_and_catalog(
1388
- input_args,
1389
- execution_partition,
1390
- global_vars,
1391
- dynamic_block_index=dynamic_block_index,
1392
- dynamic_block_indexes=dynamic_block_indexes,
1393
- dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
1394
- from_notebook=from_notebook,
1395
- )
1797
+ (
1798
+ input_vars,
1799
+ kwargs_vars,
1800
+ upstream_block_uuids,
1801
+ ) = self.fetch_input_variables_and_catalog(
1802
+ input_args,
1803
+ execution_partition,
1804
+ global_vars,
1805
+ dynamic_block_index=dynamic_block_index,
1806
+ dynamic_block_indexes=dynamic_block_indexes,
1807
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
1808
+ from_notebook=from_notebook,
1809
+ )
1396
1810
  else:
1397
1811
  input_vars, kwargs_vars, upstream_block_uuids = self.fetch_input_variables(
1398
1812
  input_args,
@@ -1451,18 +1865,22 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1451
1865
  logging_tags=logging_tags,
1452
1866
  ):
1453
1867
  # Fetch input variables
1454
- outputs_from_input_vars, input_vars, kwargs_vars, upstream_block_uuids = \
1455
- self.__get_outputs_from_input_vars(
1456
- block_run_outputs_cache=block_run_outputs_cache,
1457
- dynamic_block_index=dynamic_block_index,
1458
- dynamic_block_indexes=dynamic_block_indexes,
1459
- dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
1460
- execution_partition=execution_partition,
1461
- from_notebook=from_notebook,
1462
- global_vars=global_vars,
1463
- input_args=input_args,
1464
- metadata=metadata,
1465
- )
1868
+ (
1869
+ outputs_from_input_vars,
1870
+ input_vars,
1871
+ kwargs_vars,
1872
+ upstream_block_uuids,
1873
+ ) = self.__get_outputs_from_input_vars(
1874
+ block_run_outputs_cache=block_run_outputs_cache,
1875
+ dynamic_block_index=dynamic_block_index,
1876
+ dynamic_block_indexes=dynamic_block_indexes,
1877
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
1878
+ execution_partition=execution_partition,
1879
+ from_notebook=from_notebook,
1880
+ global_vars=global_vars,
1881
+ input_args=input_args,
1882
+ metadata=metadata,
1883
+ )
1466
1884
 
1467
1885
  global_vars_copy = global_vars.copy()
1468
1886
  for kwargs_var in kwargs_vars:
@@ -1588,11 +2006,14 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1588
2006
  preprocesser_functions = []
1589
2007
  test_functions = []
1590
2008
 
1591
- results = merge_dict({
1592
- 'preprocesser': self._block_decorator(preprocesser_functions),
1593
- 'test': self._block_decorator(test_functions),
1594
- self.type: self._block_decorator(decorated_functions),
1595
- }, outputs_from_input_vars)
2009
+ results = merge_dict(
2010
+ {
2011
+ 'preprocesser': self._block_decorator(preprocesser_functions),
2012
+ 'test': self._block_decorator(test_functions),
2013
+ self.type: self._block_decorator(decorated_functions),
2014
+ },
2015
+ outputs_from_input_vars,
2016
+ )
1596
2017
 
1597
2018
  if custom_code is not None and custom_code.strip():
1598
2019
  if BlockType.CHART != self.type:
@@ -1610,13 +2031,17 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1610
2031
  self.execute_block_function(
1611
2032
  preprocesser_function,
1612
2033
  input_vars,
2034
+ dynamic_block_index=dynamic_block_index,
2035
+ execution_partition=execution_partition,
1613
2036
  from_notebook=from_notebook,
1614
2037
  global_vars=global_vars,
2038
+ logger=logger,
2039
+ logging_tags=logging_tags,
1615
2040
  )
1616
2041
 
1617
2042
  block_function = self._validate_execution(decorated_functions, input_vars)
1618
2043
  if block_function is not None:
1619
- if logger and 'logger' not in global_vars:
2044
+ if logger:
1620
2045
  global_vars['logger'] = logger
1621
2046
 
1622
2047
  track_spark = from_notebook and self.should_track_spark()
@@ -1629,8 +2054,12 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1629
2054
  outputs = self.execute_block_function(
1630
2055
  block_function,
1631
2056
  input_vars,
2057
+ dynamic_block_index=dynamic_block_index,
2058
+ execution_partition=execution_partition,
1632
2059
  from_notebook=from_notebook,
1633
2060
  global_vars=global_vars,
2061
+ logger=logger,
2062
+ logging_tags=logging_tags,
1634
2063
  )
1635
2064
 
1636
2065
  if track_spark:
@@ -1640,7 +2069,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1640
2069
 
1641
2070
  if outputs is None:
1642
2071
  outputs = []
1643
- if type(outputs) is not list:
2072
+
2073
+ if isinstance(outputs, tuple):
2074
+ outputs = list(outputs)
2075
+
2076
+ if not isinstance(outputs, list):
1644
2077
  outputs = [outputs]
1645
2078
 
1646
2079
  return outputs
@@ -1649,26 +2082,148 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1649
2082
  self,
1650
2083
  block_function: Callable,
1651
2084
  input_vars: List,
2085
+ dynamic_block_index: Optional[int] = None,
2086
+ dynamic_block_uuid: Optional[str] = None,
2087
+ execution_partition: Optional[str] = None,
1652
2088
  from_notebook: bool = False,
1653
- global_vars: Dict = None,
2089
+ global_vars: Optional[Dict] = None,
1654
2090
  initialize_decorator_modules: bool = True,
1655
- ) -> Dict:
1656
- block_function_updated = block_function
2091
+ logger: Optional[Logger] = None,
2092
+ logging_tags: Optional[Dict] = None,
2093
+ ) -> List[Dict[str, Any]]:
2094
+ from mage_ai.settings.server import (
2095
+ MEMORY_MANAGER_V2, # Need here to mock in tests
2096
+ )
2097
+
2098
+ sig = signature(block_function)
2099
+ has_kwargs = any([p.kind == p.VAR_KEYWORD for p in sig.parameters.values()])
1657
2100
 
2101
+ block_function_updated = block_function
1658
2102
  if from_notebook and initialize_decorator_modules:
1659
2103
  block_function_updated = self.__initialize_decorator_modules(
1660
2104
  block_function,
1661
- [self.type],
2105
+ [str(self.type.value) if not isinstance(self.type, str) else str(self.type)]
2106
+ if self.type
2107
+ else [],
1662
2108
  )
1663
2109
 
1664
- sig = signature(block_function)
1665
- has_kwargs = any([p.kind == p.VAR_KEYWORD for p in sig.parameters.values()])
1666
-
1667
- if has_kwargs and global_vars is not None and len(global_vars) != 0:
2110
+ write_policy = (
2111
+ self.write_settings.batch_settings.mode
2112
+ if self.write_settings and self.write_settings.batch_settings
2113
+ else None
2114
+ )
2115
+ append_data = ExportWritePolicy.APPEND == write_policy
2116
+ part_index = None
2117
+ if append_data:
2118
+ block_uuid, changed = uuid_for_output_variables(
2119
+ self,
2120
+ block_uuid=self.uuid,
2121
+ dynamic_block_index=dynamic_block_index,
2122
+ )
2123
+ variable_object = self.get_variable_object(
2124
+ block_uuid=block_uuid,
2125
+ partition=execution_partition,
2126
+ )
2127
+ part_uuids = variable_object.part_uuids
2128
+ if part_uuids is not None:
2129
+ part_index = len(part_uuids)
2130
+ if global_vars:
2131
+ global_vars.update(part_index=part_index)
2132
+
2133
+ if MEMORY_MANAGER_V2:
2134
+ log_message_prefix = self.uuid
2135
+ if self.pipeline:
2136
+ log_message_prefix = f'{self.pipeline_uuid}:{log_message_prefix}'
2137
+ log_message_prefix = f'[{log_message_prefix}:execute_block_function]'
2138
+
2139
+ output, self.resource_usage = execute_with_memory_tracking(
2140
+ block_function_updated,
2141
+ args=input_vars,
2142
+ kwargs=global_vars
2143
+ if has_kwargs and global_vars is not None and len(global_vars) != 0
2144
+ else None,
2145
+ logger=logger,
2146
+ logging_tags=logging_tags,
2147
+ log_message_prefix=log_message_prefix,
2148
+ )
2149
+ elif has_kwargs and global_vars is not None and len(global_vars) != 0:
1668
2150
  output = block_function_updated(*input_vars, **global_vars)
1669
2151
  else:
1670
2152
  output = block_function_updated(*input_vars)
1671
2153
 
2154
+ if MEMORY_MANAGER_V2 and inspect.isgeneratorfunction(block_function_updated):
2155
+ variable_types = []
2156
+ dynamic_child = is_dynamic_block_child(self)
2157
+ output_count = part_index if part_index is not None else 0
2158
+ if output is not None and is_iterable(output):
2159
+ if dynamic_child or self.is_dynamic_child:
2160
+ # Each child will delete its own data
2161
+ # How do we delete everything ahead of time?
2162
+ delete_variable_objects_for_dynamic_child(
2163
+ self,
2164
+ dynamic_block_index=dynamic_block_index,
2165
+ execution_partition=execution_partition,
2166
+ )
2167
+ else:
2168
+ self.delete_variables(
2169
+ dynamic_block_index=dynamic_block_index,
2170
+ dynamic_block_uuid=dynamic_block_uuid,
2171
+ execution_partition=execution_partition,
2172
+ )
2173
+
2174
+ for data in output:
2175
+ if self._store_variables_in_block_function is None:
2176
+ raise Exception(
2177
+ 'Store variables function isn’t defined, '
2178
+ 'don’t proceed or else no data will be persisted'
2179
+ )
2180
+
2181
+ store_options = {}
2182
+ if output_count >= 1:
2183
+ store_options['override_outputs'] = False
2184
+
2185
+ variable_mapping = {}
2186
+
2187
+ def __output_key(order: int, output_count=output_count):
2188
+ return os.path.join(f'output_{order}', str(output_count))
2189
+
2190
+ if is_basic_iterable(data):
2191
+ if data is None or len(data) == 1:
2192
+ variable_mapping[__output_key(0)] = data
2193
+ elif len(data) == 2 and isinstance(data[1], dict):
2194
+ variable_mapping[__output_key(0)] = data[0]
2195
+ variable_mapping[__output_key(1)] = data[1]
2196
+ else:
2197
+ for idx, item in enumerate(data):
2198
+ variable_mapping[__output_key(idx)] = item
2199
+ else:
2200
+ variable_mapping[__output_key(0)] = data
2201
+
2202
+ variables = self._store_variables_in_block_function(
2203
+ variable_mapping=variable_mapping,
2204
+ clean_variable_uuid=False,
2205
+ skip_delete=True,
2206
+ **store_options,
2207
+ )
2208
+ if variables is not None and isinstance(variables, list):
2209
+ variable_types += [
2210
+ variable.variable_type
2211
+ for variable in variables
2212
+ if isinstance(variable, Variable)
2213
+ ]
2214
+
2215
+ output_count += 1
2216
+
2217
+ if len(variable_types) >= 1 and self._store_variables_in_block_function:
2218
+ self._store_variables_in_block_function(
2219
+ {'output_0': variable_types},
2220
+ save_variable_types_only=True,
2221
+ )
2222
+
2223
+ self._store_variables_in_block_function = None
2224
+
2225
+ if output is None:
2226
+ return []
1672
2227
  return output
1673
2228
 
1674
2229
  def __initialize_decorator_modules(
@@ -1747,11 +2302,18 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1747
2302
  upstream block UUIDs.
1748
2303
  """
1749
2304
 
1750
- if any([is_dynamic_block(
1751
- upstream_block,
1752
- ) or is_dynamic_block_child(
1753
- upstream_block,
1754
- ) for upstream_block in self.upstream_blocks]):
2305
+ if (self.is_dynamic_v2 and self.is_dynamic_child) or (
2306
+ not self.is_dynamic_v2
2307
+ and any([
2308
+ is_dynamic_block(
2309
+ upstream_block,
2310
+ )
2311
+ or is_dynamic_block_child(
2312
+ upstream_block,
2313
+ )
2314
+ for upstream_block in self.upstream_blocks
2315
+ ])
2316
+ ):
1755
2317
  return fetch_input_variables_for_dynamic_upstream_blocks(
1756
2318
  self,
1757
2319
  input_args,
@@ -1760,7 +2322,6 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1760
2322
  execution_partition=execution_partition,
1761
2323
  from_notebook=from_notebook,
1762
2324
  global_vars=global_vars,
1763
- # For non-dynamic upstream blocks
1764
2325
  block_run_outputs_cache=block_run_outputs_cache,
1765
2326
  data_integration_settings_mapping=data_integration_settings_mapping,
1766
2327
  upstream_block_uuids_override=upstream_block_uuids_override,
@@ -1780,6 +2341,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1780
2341
  global_vars=global_vars,
1781
2342
  metadata=metadata,
1782
2343
  upstream_block_uuids_override=upstream_block_uuids_override,
2344
+ current_block=self,
1783
2345
  )
1784
2346
 
1785
2347
  return variables
@@ -1794,8 +2356,8 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1794
2356
  for v in output_variable_objects:
1795
2357
  if v.variable_type != VariableType.DATAFRAME:
1796
2358
  continue
1797
- data = self.pipeline.variable_manager.get_variable(
1798
- self.pipeline.uuid,
2359
+ data = self.variable_manager.get_variable(
2360
+ self.pipeline_uuid,
1799
2361
  self.uuid,
1800
2362
  v.uuid,
1801
2363
  variable_type=VariableType.DATAFRAME_ANALYSIS,
@@ -1807,12 +2369,12 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1807
2369
  def get_variables_by_block(
1808
2370
  self,
1809
2371
  block_uuid: str,
1810
- dynamic_block_index: int = None,
1811
- dynamic_block_uuid: str = None,
1812
- partition: str = None,
1813
- ):
1814
- variable_manager = self.pipeline.variable_manager
1815
-
2372
+ clean_block_uuid: bool = True,
2373
+ dynamic_block_index: Optional[int] = None,
2374
+ dynamic_block_uuid: Optional[str] = None,
2375
+ max_results: Optional[int] = None,
2376
+ partition: Optional[str] = None,
2377
+ ) -> List[str]:
1816
2378
  block_uuid_use, changed = uuid_for_output_variables(
1817
2379
  self,
1818
2380
  block_uuid=block_uuid,
@@ -1820,10 +2382,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1820
2382
  dynamic_block_uuid=dynamic_block_uuid,
1821
2383
  )
1822
2384
 
1823
- res = variable_manager.get_variables_by_block(
1824
- self.pipeline.uuid,
2385
+ res = self.variable_manager.get_variables_by_block(
2386
+ self.pipeline_uuid,
1825
2387
  block_uuid=block_uuid_use,
1826
- clean_block_uuid=not changed,
2388
+ clean_block_uuid=not changed and clean_block_uuid,
2389
+ max_results=max_results,
1827
2390
  partition=partition,
1828
2391
  )
1829
2392
 
@@ -1847,9 +2410,12 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1847
2410
  partition: str = None,
1848
2411
  raise_exception: bool = False,
1849
2412
  spark=None,
2413
+ input_data_types: Optional[List[InputDataType]] = None,
2414
+ read_batch_settings: Optional[BatchSettings] = None,
2415
+ read_chunks: Optional[List[ChunkKeyTypeUnion]] = None,
2416
+ write_batch_settings: Optional[BatchSettings] = None,
2417
+ write_chunks: Optional[List[ChunkKeyTypeUnion]] = None,
1850
2418
  ):
1851
- variable_manager = self.pipeline.variable_manager
1852
-
1853
2419
  block_uuid_use, changed = uuid_for_output_variables(
1854
2420
  self,
1855
2421
  block_uuid=block_uuid,
@@ -1857,40 +2423,151 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1857
2423
  dynamic_block_uuid=dynamic_block_uuid,
1858
2424
  )
1859
2425
 
1860
- value = variable_manager.get_variable(
1861
- self.pipeline.uuid,
2426
+ value = self.variable_manager.get_variable(
2427
+ self.pipeline_uuid,
1862
2428
  block_uuid=block_uuid_use,
1863
2429
  clean_block_uuid=not changed,
1864
2430
  partition=partition,
1865
2431
  raise_exception=raise_exception,
1866
2432
  spark=spark,
1867
2433
  variable_uuid=variable_uuid,
2434
+ input_data_types=input_data_types,
2435
+ read_batch_settings=read_batch_settings,
2436
+ read_chunks=read_chunks,
2437
+ write_batch_settings=write_batch_settings,
2438
+ write_chunks=write_chunks,
1868
2439
  )
1869
2440
 
1870
2441
  return value
1871
2442
 
1872
- def get_variable_object(
2443
+ def read_partial_data(
1873
2444
  self,
1874
- block_uuid: str,
1875
- variable_uuid: str = None,
1876
- dynamic_block_index: int = None,
1877
- partition: str = None,
2445
+ variable_uuid: str,
2446
+ batch_settings: Optional[BatchSettings] = None,
2447
+ chunks: Optional[List[ChunkKeyTypeUnion]] = None,
2448
+ input_data_types: Optional[List[InputDataType]] = None,
2449
+ part_uuid: Optional[Union[int, str]] = None,
2450
+ partition: Optional[str] = None,
1878
2451
  ):
1879
- variable_manager = self.pipeline.variable_manager
2452
+ return self.get_variable_object(
2453
+ self.uuid,
2454
+ variable_uuid,
2455
+ partition=partition,
2456
+ ).read_partial_data(
2457
+ batch_settings=batch_settings,
2458
+ chunks=chunks,
2459
+ input_data_types=input_data_types,
2460
+ part_uuid=part_uuid,
2461
+ )
1880
2462
 
2463
+ def get_variable_object(
2464
+ self,
2465
+ block_uuid: Optional[str] = None,
2466
+ variable_uuid: Optional[str] = None,
2467
+ clean_block_uuid: bool = True,
2468
+ dynamic_block_index: Optional[int] = None,
2469
+ input_data_types: Optional[List[InputDataType]] = None,
2470
+ ordinal_position: Optional[int] = None, # Used to get cached variable information
2471
+ partition: Optional[str] = None,
2472
+ read_batch_settings: Optional[BatchSettings] = None,
2473
+ read_chunks: Optional[List[ChunkKeyTypeUnion]] = None,
2474
+ skip_check_variable_type: Optional[bool] = None,
2475
+ write_batch_settings: Optional[BatchSettings] = None,
2476
+ write_chunks: Optional[List[ChunkKeyTypeUnion]] = None,
2477
+ ) -> Variable:
2478
+ block_uuid = block_uuid or self.uuid
1881
2479
  block_uuid, changed = uuid_for_output_variables(
1882
2480
  self,
1883
2481
  block_uuid=block_uuid,
1884
2482
  dynamic_block_index=dynamic_block_index,
1885
2483
  )
1886
2484
 
1887
- return variable_manager.get_variable_object(
1888
- self.pipeline.uuid,
2485
+ variable_type_information = None
2486
+ variable_types_information = None
2487
+ skip_check_variable_type = skip_check_variable_type or False
2488
+ if VARIABLE_DATA_OUTPUT_META_CACHE:
2489
+ dynamic_child = is_dynamic_block_child(self)
2490
+ group_type = (
2491
+ VariableAggregateSummaryGroupType.DYNAMIC
2492
+ if dynamic_child
2493
+ else VariableAggregateSummaryGroupType.PARTS
2494
+ )
2495
+ variable_type_information = self.get_variable_aggregate_cache(
2496
+ variable_uuid, VariableAggregateDataType.TYPE, default_group_type=group_type
2497
+ )
2498
+ variable_types_information = self.get_variable_aggregate_cache(
2499
+ variable_uuid,
2500
+ VariableAggregateDataType.TYPE,
2501
+ group_type=group_type,
2502
+ partition=partition,
2503
+ )
2504
+
2505
+ if (
2506
+ variable_type_information
2507
+ and variable_type_information.type == VariableType.ITERABLE
2508
+ and not variable_types_information
2509
+ ):
2510
+ # If the dynamic parent block is an interable with no types information,
2511
+ # then the data from the parent block won’t have any type information.
2512
+ # Skip variable type check when instantiating a variable object for the children.
2513
+ skip_check_variable_type = True
2514
+ elif (
2515
+ dynamic_child
2516
+ and variable_types_information is not None
2517
+ and isinstance(variable_types_information, Iterable)
2518
+ and (
2519
+ (
2520
+ dynamic_block_index is not None
2521
+ and int(dynamic_block_index) < len(variable_types_information)
2522
+ )
2523
+ or (
2524
+ ordinal_position is not None
2525
+ and int(ordinal_position) < len(variable_types_information)
2526
+ )
2527
+ )
2528
+ ):
2529
+ position = (
2530
+ int(ordinal_position)
2531
+ if ordinal_position is not None
2532
+ else int(dynamic_block_index)
2533
+ if dynamic_block_index is not None
2534
+ else None
2535
+ )
2536
+ if position is not None and isinstance(variable_types_information, Iterable):
2537
+ variable_type_information = variable_types_information[position]
2538
+ if (
2539
+ isinstance(variable_type_information, Iterable)
2540
+ and len(variable_type_information) >= 1
2541
+ ):
2542
+ variable_type_information = variable_type_information[0]
2543
+
2544
+ variable_types_information = None
2545
+
2546
+ variable_types = []
2547
+ if isinstance(variable_types_information, Iterable):
2548
+ for v in variable_types_information:
2549
+ if isinstance(v, list):
2550
+ variable_types += [vv.type for vv in v]
2551
+ else:
2552
+ variable_types.append(v.type)
2553
+
2554
+ return self.variable_manager.get_variable_object(
2555
+ self.pipeline_uuid,
1889
2556
  block_uuid=block_uuid,
1890
- clean_block_uuid=not changed,
2557
+ clean_block_uuid=not changed and clean_block_uuid,
1891
2558
  partition=partition,
1892
2559
  spark=self.get_spark_session(),
2560
+ skip_check_variable_type=skip_check_variable_type,
2561
+ variable_type=variable_type_information.type
2562
+ if variable_type_information is not None
2563
+ else None,
2564
+ variable_types=variable_types,
1893
2565
  variable_uuid=variable_uuid,
2566
+ input_data_types=input_data_types,
2567
+ read_batch_settings=read_batch_settings,
2568
+ read_chunks=read_chunks,
2569
+ write_batch_settings=write_batch_settings,
2570
+ write_chunks=write_chunks,
1894
2571
  )
1895
2572
 
1896
2573
  def get_raw_outputs(
@@ -1901,6 +2578,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1901
2578
  global_vars: Dict = None,
1902
2579
  dynamic_block_index: int = None,
1903
2580
  dynamic_block_uuid: str = None,
2581
+ input_data_types: Optional[List[InputDataType]] = None,
2582
+ read_batch_settings: Optional[BatchSettings] = None,
2583
+ read_chunks: Optional[List[ChunkKeyTypeUnion]] = None,
2584
+ write_batch_settings: Optional[BatchSettings] = None,
2585
+ write_chunks: Optional[List[ChunkKeyTypeUnion]] = None,
1904
2586
  ) -> List[Any]:
1905
2587
  all_variables = self.get_variables_by_block(
1906
2588
  block_uuid=block_uuid,
@@ -1912,6 +2594,9 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1912
2594
  outputs = []
1913
2595
 
1914
2596
  for variable_uuid in all_variables:
2597
+ if not is_output_variable(variable_uuid):
2598
+ continue
2599
+
1915
2600
  variable = self.pipeline.get_block_variable(
1916
2601
  block_uuid,
1917
2602
  variable_uuid,
@@ -1922,6 +2607,11 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1922
2607
  spark=self.__get_spark_session_from_global_vars(global_vars),
1923
2608
  dynamic_block_index=dynamic_block_index,
1924
2609
  dynamic_block_uuid=dynamic_block_uuid,
2610
+ input_data_types=input_data_types,
2611
+ read_batch_settings=read_batch_settings,
2612
+ read_chunks=read_chunks,
2613
+ write_batch_settings=write_batch_settings,
2614
+ write_chunks=write_chunks,
1925
2615
  )
1926
2616
  outputs.append(variable)
1927
2617
 
@@ -1929,411 +2619,181 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
1929
2619
 
1930
2620
  def get_outputs(
1931
2621
  self,
1932
- execution_partition: str = None,
1933
- include_print_outputs: bool = True,
2622
+ block_uuid: Optional[str] = None,
1934
2623
  csv_lines_only: bool = False,
2624
+ dynamic_block_index: Optional[int] = None,
2625
+ exclude_blank_variable_uuids: bool = False,
2626
+ execution_partition: Optional[str] = None,
2627
+ include_print_outputs: bool = True,
2628
+ metadata: Optional[Dict] = None,
1935
2629
  sample: bool = True,
1936
- sample_count: int = DATAFRAME_SAMPLE_COUNT_PREVIEW,
1937
- variable_type: VariableType = None,
1938
- block_uuid: str = None,
1939
- selected_variables: List[str] = None,
1940
- metadata: Dict = None,
1941
- dynamic_block_index: int = None,
1942
- ) -> List[Dict]:
1943
- data_products = []
1944
- outputs = []
1945
-
2630
+ sample_count: Optional[int] = None,
2631
+ selected_variables: Optional[List[str]] = None,
2632
+ variable_type: Optional[VariableType] = None,
2633
+ ) -> List[Dict[str, Any]]:
1946
2634
  is_dynamic_child = is_dynamic_block_child(self)
1947
2635
  is_dynamic = is_dynamic_block(self)
1948
2636
 
1949
- if is_dynamic_child or is_dynamic:
1950
- pairs = []
1951
-
1952
- if is_dynamic_child:
1953
- lazy_variable_controller = get_outputs_for_dynamic_child(
1954
- self,
1955
- execution_partition=execution_partition,
1956
- sample=sample,
1957
- sample_count=sample_count,
1958
- )
1959
- pairs = lazy_variable_controller.render(
1960
- dynamic_block_index=dynamic_block_index,
1961
- )
1962
- elif is_dynamic:
1963
- tup = get_outputs_for_dynamic_block(
1964
- self,
1965
- execution_partition=execution_partition,
1966
- sample=sample,
1967
- sample_count=sample_count,
1968
- )
1969
- pairs.append(tup)
1970
-
1971
- for pair in pairs:
1972
- child_data = None
1973
- metadata = None
1974
- if len(pair) >= 1:
1975
- child_data = pair[0]
1976
- if len(pair) >= 2:
1977
- metadata = pair[1]
1978
-
1979
- for output, variable_uuid in [
1980
- (child_data, 'child_data'),
1981
- (metadata, 'metadata'),
1982
- ]:
1983
- if output is None:
1984
- continue
1985
-
1986
- data, is_data_product = self.__format_output_data(
1987
- output,
1988
- variable_uuid,
1989
- block_uuid=self.uuid,
1990
- csv_lines_only=csv_lines_only,
1991
- execution_partition=execution_partition,
1992
- )
1993
-
1994
- outputs_below_limit = not sample or not sample_count
1995
- if is_data_product:
1996
- outputs_below_limit = outputs_below_limit or \
1997
- (sample_count is not None and len(data_products) < sample_count)
1998
- else:
1999
- outputs_below_limit = outputs_below_limit or \
2000
- (sample_count is not None and len(outputs) < sample_count)
2001
-
2002
- if outputs_below_limit:
2003
- if is_data_product:
2004
- data_products.append(data)
2005
- else:
2006
- outputs.append(data)
2007
- else:
2008
- if self.pipeline is None:
2009
- return
2010
-
2011
- if not block_uuid:
2012
- block_uuid = self.uuid
2013
-
2014
- # The block_run’s block_uuid for replicated blocks will be in this format:
2015
- # [block_uuid]:[replicated_block_uuid]
2016
- # We need to use the original block_uuid to get the proper output.
2017
-
2018
- # Block runs for dynamic child blocks will have the following block UUID:
2019
- # [block.uuid]:[index]
2020
- # Don’t use the original UUID even if the block is a replica because it will get rid of
2021
- # the dynamic child block index.
2022
-
2023
- data_products = []
2024
- outputs = []
2025
-
2026
- all_variables = self.get_variables_by_block(
2637
+ if not is_dynamic and not is_dynamic_child:
2638
+ return get_outputs_for_display_sync(
2639
+ self,
2027
2640
  block_uuid=block_uuid,
2028
- partition=execution_partition,
2641
+ csv_lines_only=csv_lines_only,
2642
+ exclude_blank_variable_uuids=exclude_blank_variable_uuids,
2643
+ execution_partition=execution_partition,
2644
+ include_print_outputs=include_print_outputs,
2645
+ sample=sample,
2646
+ sample_count=sample_count or DATAFRAME_SAMPLE_COUNT_PREVIEW,
2647
+ selected_variables=selected_variables,
2648
+ variable_type=variable_type,
2029
2649
  )
2030
2650
 
2031
- if not include_print_outputs:
2032
- all_variables = self.output_variables(
2033
- execution_partition=execution_partition,
2034
- )
2035
-
2036
- for v in all_variables:
2037
- if selected_variables and v not in selected_variables:
2038
- continue
2039
-
2040
- variable_object = self.get_variable_object(
2041
- block_uuid=block_uuid,
2042
- partition=execution_partition,
2043
- variable_uuid=v,
2044
- )
2651
+ sample_count_use = sample_count or DYNAMIC_CHILD_BLOCK_SAMPLE_COUNT_PREVIEW
2652
+ output_sets = []
2653
+ variable_sets = []
2045
2654
 
2046
- if variable_type is not None and variable_object.variable_type != variable_type:
2047
- continue
2048
-
2049
- data = variable_object.read_data(
2050
- sample=sample,
2051
- sample_count=sample_count,
2052
- spark=self.get_spark_session(),
2053
- )
2054
- data, is_data_product = self.__format_output_data(
2055
- data,
2056
- v,
2057
- block_uuid=block_uuid,
2058
- csv_lines_only=csv_lines_only,
2059
- execution_partition=execution_partition,
2060
- )
2061
- if is_data_product:
2062
- data_products.append(data)
2063
- else:
2064
- outputs.append(data)
2065
-
2066
- return outputs + data_products
2067
-
2068
- async def __get_outputs_async(
2069
- self,
2070
- csv_lines_only: bool = False,
2071
- execution_partition: str = None,
2072
- include_print_outputs: bool = True,
2073
- sample: bool = True,
2074
- sample_count: int = DATAFRAME_SAMPLE_COUNT_PREVIEW,
2075
- variable_type: VariableType = None,
2076
- block_uuid: str = None,
2077
- dynamic_block_index: int = None,
2078
- ) -> List[Dict]:
2079
- data_products = []
2080
- outputs = []
2081
-
2082
- is_dynamic_child = is_dynamic_block_child(self)
2083
- is_dynamic = is_dynamic_block(self)
2084
-
2085
- if is_dynamic_child or is_dynamic:
2086
- pairs = []
2087
-
2088
- if is_dynamic_child:
2089
- lazy_variable_controller = get_outputs_for_dynamic_child(
2090
- self,
2091
- execution_partition=execution_partition,
2092
- sample=sample,
2093
- sample_count=sample_count,
2094
- )
2095
- pairs = await lazy_variable_controller.render_async(
2096
- dynamic_block_index=dynamic_block_index,
2097
- )
2098
- elif is_dynamic:
2099
- tup = await get_outputs_for_dynamic_block_async(
2100
- self,
2101
- execution_partition=execution_partition,
2102
- sample=sample,
2103
- sample_count=sample_count,
2104
- )
2105
- pairs.append(tup)
2106
-
2107
- if len(pairs) > 10:
2108
- # Limit the number of dynamic block children we display output for in the UI
2109
- pairs = pairs[:DATAFRAME_SAMPLE_COUNT_PREVIEW]
2110
- for pair in pairs:
2111
- child_data = None
2112
- metadata = None
2113
- if len(pair) >= 1:
2114
- child_data = pair[0]
2115
- if len(pair) >= 2:
2116
- metadata = pair[1]
2117
-
2118
- for output, variable_uuid in [
2119
- (child_data, 'child_data'),
2120
- (metadata, 'metadata'),
2121
- ]:
2122
- if output is None:
2123
- continue
2124
-
2125
- data, is_data_product = self.__format_output_data(
2126
- output,
2127
- variable_uuid,
2128
- block_uuid=self.uuid,
2129
- csv_lines_only=csv_lines_only,
2130
- execution_partition=execution_partition,
2131
- )
2132
-
2133
- if is_data_product:
2134
- data_products.append(data)
2135
- else:
2136
- outputs.append(data)
2137
- else:
2138
- if self.pipeline is None:
2139
- return
2140
-
2141
- if not block_uuid:
2142
- block_uuid = self.uuid
2143
-
2144
- variable_manager = self.pipeline.variable_manager
2145
-
2146
- all_variables = variable_manager.get_variables_by_block(
2147
- self.pipeline.uuid,
2148
- block_uuid,
2149
- partition=execution_partition,
2150
- max_results=DATAFRAME_SAMPLE_COUNT_PREVIEW,
2655
+ if is_dynamic_child:
2656
+ lazy_variable_controller = get_outputs_for_dynamic_child(
2657
+ self,
2658
+ execution_partition=execution_partition,
2659
+ sample=sample,
2660
+ sample_count=sample_count_use,
2151
2661
  )
2662
+ variable_sets: List[
2663
+ Union[
2664
+ Tuple[Optional[Any], Dict],
2665
+ List[LazyVariableSet],
2666
+ ],
2667
+ ] = lazy_variable_controller.render(
2668
+ dynamic_block_index=dynamic_block_index,
2669
+ lazy_load=True,
2670
+ )
2671
+ elif is_dynamic:
2672
+ output_pair: List[
2673
+ Optional[
2674
+ Union[
2675
+ Any,
2676
+ Dict,
2677
+ int,
2678
+ pd.DataFrame,
2679
+ str,
2680
+ ]
2681
+ ]
2682
+ ] = get_outputs_for_dynamic_block(
2683
+ self,
2684
+ execution_partition=execution_partition,
2685
+ sample=sample,
2686
+ sample_count=sample_count_use,
2687
+ )
2688
+ output_sets.append(output_pair)
2152
2689
 
2153
- if not include_print_outputs:
2154
- all_variables = self.output_variables(execution_partition=execution_partition)
2155
-
2156
- for v in all_variables:
2157
- variable_object = variable_manager.get_variable_object(
2158
- self.pipeline.uuid,
2159
- block_uuid,
2160
- v,
2161
- partition=execution_partition,
2162
- spark=self.get_spark_session(),
2163
- )
2164
-
2165
- if variable_type is not None and variable_object.variable_type != variable_type:
2166
- continue
2167
-
2168
- data = await variable_object.read_data_async(
2169
- sample=True,
2170
- sample_count=sample_count,
2171
- spark=self.get_spark_session(),
2172
- )
2173
- data, is_data_product = self.__format_output_data(
2174
- data,
2175
- v,
2176
- block_uuid=block_uuid,
2177
- csv_lines_only=csv_lines_only,
2178
- execution_partition=execution_partition,
2179
- )
2180
- if is_data_product:
2181
- data_products.append(data)
2182
- else:
2183
- outputs.append(data)
2184
-
2185
- return outputs + data_products
2690
+ output_sets = output_sets[:DATAFRAME_SAMPLE_COUNT_PREVIEW]
2691
+ variable_sets = variable_sets[:DATAFRAME_SAMPLE_COUNT_PREVIEW]
2692
+ child_data_sets = [lazy_variable_set.read_data() for lazy_variable_set in variable_sets]
2186
2693
 
2187
- def __format_output_data(
2694
+ return get_outputs_for_display_dynamic_block(
2695
+ self,
2696
+ output_sets,
2697
+ child_data_sets,
2698
+ block_uuid=block_uuid,
2699
+ csv_lines_only=csv_lines_only,
2700
+ dynamic_block_index=dynamic_block_index,
2701
+ exclude_blank_variable_uuids=exclude_blank_variable_uuids,
2702
+ execution_partition=execution_partition,
2703
+ metadata=metadata,
2704
+ sample=sample,
2705
+ sample_count=sample_count_use,
2706
+ )
2707
+
2708
+ async def __get_outputs_async(
2188
2709
  self,
2189
- data: Any,
2190
- variable_uuid: str,
2191
- block_uuid: str = None,
2710
+ execution_partition: Optional[str] = None,
2711
+ include_print_outputs: bool = True,
2192
2712
  csv_lines_only: bool = False,
2193
- execution_partition: str = None,
2194
- skip_dynamic_block: bool = False,
2195
- ) -> Tuple[Dict, bool]:
2196
- """
2197
- Takes variable data and formats it to return to the frontend.
2198
-
2199
- Returns:
2200
- Tuple[Dict, bool]: Tuple of the formatted data and is_data_product boolean. Data product
2201
- outputs and non data product outputs are handled slightly differently.
2202
- """
2203
- variable_manager = self.pipeline.variable_manager
2204
-
2713
+ sample: bool = True,
2714
+ sample_count: Optional[int] = None,
2715
+ variable_type: Optional[VariableType] = None,
2716
+ block_uuid: Optional[str] = None,
2717
+ selected_variables: Optional[List[str]] = None,
2718
+ metadata: Optional[Dict] = None,
2719
+ dynamic_block_index: Optional[int] = None,
2720
+ exclude_blank_variable_uuids: bool = False,
2721
+ max_results: Optional[int] = None,
2722
+ ) -> List[Dict[str, Any]]:
2205
2723
  is_dynamic_child = is_dynamic_block_child(self)
2206
2724
  is_dynamic = is_dynamic_block(self)
2207
2725
 
2208
- if (is_dynamic_child or is_dynamic) and not skip_dynamic_block:
2209
- from mage_ai.data_preparation.models.block.dynamic.utils import (
2210
- coerce_into_dataframe,
2726
+ if not is_dynamic and not is_dynamic_child:
2727
+ return await get_outputs_for_display_async(
2728
+ self,
2729
+ block_uuid=block_uuid,
2730
+ csv_lines_only=csv_lines_only,
2731
+ exclude_blank_variable_uuids=exclude_blank_variable_uuids,
2732
+ execution_partition=execution_partition,
2733
+ include_print_outputs=include_print_outputs,
2734
+ sample=sample,
2735
+ sample_count=sample_count or DATAFRAME_SAMPLE_COUNT_PREVIEW,
2736
+ selected_variables=selected_variables,
2737
+ variable_type=variable_type,
2738
+ max_results=max_results,
2211
2739
  )
2212
2740
 
2213
- data, is_data_product = self.__format_output_data(
2214
- coerce_into_dataframe(data),
2215
- variable_uuid=variable_uuid,
2216
- skip_dynamic_block=True,
2741
+ sample_count_use = sample_count or DYNAMIC_CHILD_BLOCK_SAMPLE_COUNT_PREVIEW
2742
+ output_sets = []
2743
+ variable_sets = []
2744
+
2745
+ if is_dynamic_child:
2746
+ lazy_variable_controller = get_outputs_for_dynamic_child(
2747
+ self,
2748
+ execution_partition=execution_partition,
2749
+ limit_parts=max_results,
2750
+ sample=sample,
2751
+ sample_count=sample_count_use,
2752
+ )
2753
+ variable_sets: List[
2754
+ Union[
2755
+ Tuple[Optional[Any], Dict],
2756
+ List[LazyVariableSet],
2757
+ ],
2758
+ ] = await lazy_variable_controller.render_async(
2759
+ dynamic_block_index=dynamic_block_index,
2760
+ lazy_load=True,
2217
2761
  )
2218
2762
 
2219
- return merge_dict(data, dict(multi_output=True)), is_data_product
2220
- elif isinstance(data, pd.DataFrame):
2221
- if csv_lines_only:
2222
- data = dict(
2223
- table=data.to_csv(header=True, index=False).strip('\n').split('\n')
2224
- )
2225
- else:
2226
- try:
2227
- analysis = variable_manager.get_variable(
2228
- self.pipeline.uuid,
2229
- block_uuid,
2230
- variable_uuid,
2231
- dataframe_analysis_keys=['metadata', 'statistics'],
2232
- partition=execution_partition,
2233
- variable_type=VariableType.DATAFRAME_ANALYSIS,
2234
- )
2235
- except Exception:
2236
- analysis = None
2237
- if analysis is not None and \
2238
- (analysis.get('statistics') or analysis.get('metadata')):
2239
-
2240
- stats = analysis.get('statistics', {})
2241
- column_types = (analysis.get('metadata') or {}).get('column_types', {})
2242
- row_count = stats.get('original_row_count', stats.get('count'))
2243
- column_count = stats.get('original_column_count', len(column_types))
2244
- else:
2245
- row_count, column_count = data.shape
2246
-
2247
- columns_to_display = data.columns.tolist()[:DATAFRAME_ANALYSIS_MAX_COLUMNS]
2248
- data = dict(
2249
- sample_data=dict(
2250
- columns=columns_to_display,
2251
- rows=json.loads(
2252
- data[columns_to_display].to_json(orient='split', date_format='iso'),
2253
- )['data']
2254
- ),
2255
- shape=[row_count, column_count],
2256
- type=DataType.TABLE,
2257
- variable_uuid=variable_uuid,
2258
- )
2259
- return data, True
2260
- elif isinstance(data, pl.DataFrame):
2261
- try:
2262
- analysis = variable_manager.get_variable(
2263
- self.pipeline.uuid,
2264
- block_uuid,
2265
- variable_uuid,
2266
- dataframe_analysis_keys=['statistics'],
2267
- partition=execution_partition,
2268
- variable_type=VariableType.DATAFRAME_ANALYSIS,
2269
- )
2270
- except Exception:
2271
- analysis = None
2272
- if analysis is not None:
2273
- stats = analysis.get('statistics', {})
2274
- row_count = stats.get('original_row_count')
2275
- column_count = stats.get('original_column_count')
2276
- else:
2277
- row_count, column_count = data.shape
2278
- columns_to_display = data.columns[:DATAFRAME_ANALYSIS_MAX_COLUMNS]
2279
- data = dict(
2280
- sample_data=dict(
2281
- columns=columns_to_display,
2282
- rows=[
2283
- list(row.values()) for row in json.loads(
2284
- data[columns_to_display].write_json(row_oriented=True)
2285
- )
2286
- ]
2287
- ),
2288
- shape=[row_count, column_count],
2289
- type=DataType.TABLE,
2290
- variable_uuid=variable_uuid,
2291
- )
2292
- return data, True
2293
- elif is_geo_dataframe(data):
2294
- data = dict(
2295
- text_data=f'''Use the code in a scratchpad to get the output of the block:
2296
-
2297
- from mage_ai.data_preparation.variable_manager import get_variable
2298
- df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2299
- ''',
2300
- type=DataType.TEXT,
2301
- variable_uuid=variable_uuid,
2302
- )
2303
- return data, False
2304
- elif type(data) is str:
2305
- data = dict(
2306
- text_data=data,
2307
- type=DataType.TEXT,
2308
- variable_uuid=variable_uuid,
2309
- )
2310
- return data, False
2311
- elif type(data) is dict or type(data) is list:
2312
- data = dict(
2313
- text_data=simplejson.dumps(
2314
- data,
2315
- default=encode_complex,
2316
- ignore_nan=True,
2317
- ),
2318
- type=DataType.TEXT,
2319
- variable_uuid=variable_uuid,
2320
- )
2321
- return data, False
2322
- elif is_spark_dataframe(data):
2323
- df = data.toPandas()
2324
- columns_to_display = df.columns.tolist()[:DATAFRAME_ANALYSIS_MAX_COLUMNS]
2325
- data = dict(
2326
- sample_data=dict(
2327
- columns=columns_to_display,
2328
- rows=json.loads(
2329
- df[columns_to_display].to_json(orient='split', date_format='iso'),
2330
- )['data']
2331
- ),
2332
- type=DataType.TABLE,
2333
- variable_uuid=variable_uuid,
2763
+ elif is_dynamic:
2764
+ output_pair: List[
2765
+ Optional[Union[Dict, int, str, pd.DataFrame, Any]],
2766
+ ] = await get_outputs_for_dynamic_block_async(
2767
+ self,
2768
+ execution_partition=execution_partition,
2769
+ limit_parts=max_results,
2770
+ sample=sample,
2771
+ sample_count=sample_count_use,
2334
2772
  )
2335
- return data, True
2336
- return data, False
2773
+ output_sets.append(output_pair)
2774
+
2775
+ # Limit the number of dynamic block children we display output for in the UI
2776
+ output_sets = output_sets[:DATAFRAME_SAMPLE_COUNT_PREVIEW]
2777
+ variable_sets = variable_sets[:DATAFRAME_SAMPLE_COUNT_PREVIEW]
2778
+ child_data_sets = await asyncio.gather(*[
2779
+ lazy_variable_set.read_data_async() for lazy_variable_set in variable_sets
2780
+ ])
2781
+
2782
+ return get_outputs_for_display_dynamic_block(
2783
+ self,
2784
+ output_sets,
2785
+ child_data_sets,
2786
+ block_uuid=block_uuid,
2787
+ csv_lines_only=csv_lines_only,
2788
+ exclude_blank_variable_uuids=exclude_blank_variable_uuids,
2789
+ execution_partition=execution_partition,
2790
+ metadata=metadata,
2791
+ sample=sample,
2792
+ sample_count=sample_count_use,
2793
+ )
2794
+
2795
+ def __format_output_data(self, *args, **kwargs) -> Tuple[Dict, bool]:
2796
+ return format_output_data(self, *args, **kwargs)
2337
2797
 
2338
2798
  def __save_outputs_prepare(self, outputs, override_output_variable: bool = False) -> Dict:
2339
2799
  variable_mapping = dict()
@@ -2344,10 +2804,11 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2344
2804
  if not isinstance(o, dict):
2345
2805
  continue
2346
2806
 
2347
- if all(k in o for k in ['variable_uuid', 'text_data']) and \
2348
- (not is_output_variable(o['variable_uuid']) or
2349
- BlockType.SCRATCHPAD == self.type or
2350
- override_output_variable):
2807
+ if all(k in o for k in ['variable_uuid', 'text_data']) and (
2808
+ not is_output_variable(o['variable_uuid'])
2809
+ or BlockType.SCRATCHPAD == self.type
2810
+ or override_output_variable
2811
+ ):
2351
2812
  variable_mapping[o['variable_uuid']] = o['text_data']
2352
2813
 
2353
2814
  self._outputs = outputs
@@ -2375,9 +2836,26 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2375
2836
  self,
2376
2837
  outputs,
2377
2838
  override: bool = False,
2839
+ override_conditionally: bool = False,
2378
2840
  override_outputs: bool = False,
2379
2841
  ) -> None:
2380
2842
  variable_mapping = self.__save_outputs_prepare(outputs)
2843
+
2844
+ if override_conditionally:
2845
+ for variable_uuid, _ in variable_mapping.items():
2846
+ variable = self.get_variable_object(
2847
+ self.uuid,
2848
+ variable_uuid,
2849
+ )
2850
+ if not variable or not variable.variable_type:
2851
+ continue
2852
+
2853
+ # if VariableType
2854
+ # variable = self.get_variable_object(variable_uuid=variable_uuid)
2855
+ # if variable.exists():
2856
+ # variable_mapping.pop(variable_uuid)
2857
+ pass
2858
+
2381
2859
  await self.store_variables_async(
2382
2860
  variable_mapping,
2383
2861
  override=override,
@@ -2386,18 +2864,29 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2386
2864
 
2387
2865
  def get_executor_type(self) -> str:
2388
2866
  if self.executor_type:
2389
- return Template(self.executor_type).render(**get_template_vars())
2390
- return self.executor_type
2867
+ block_executor_type = Template(self.executor_type).render(**get_template_vars())
2868
+ else:
2869
+ block_executor_type = None
2870
+ if not block_executor_type or block_executor_type == ExecutorType.LOCAL_PYTHON:
2871
+ # If block executor_type is not set, fall back to pipeline level executor_type
2872
+ if self.pipeline:
2873
+ pipeline_executor_type = self.pipeline.get_executor_type()
2874
+ else:
2875
+ pipeline_executor_type = None
2876
+ block_executor_type = pipeline_executor_type or block_executor_type
2877
+ return block_executor_type
2391
2878
 
2392
2879
  def get_pipelines_from_cache(self, block_cache: BlockCache = None) -> List[Dict]:
2393
2880
  if block_cache is None:
2394
2881
  block_cache = BlockCache()
2395
- arr = block_cache.get_pipelines(self)
2882
+ arr = block_cache.get_pipelines(self, self.repo_path)
2396
2883
 
2397
2884
  return unique_by(
2398
2885
  arr,
2399
- lambda x: (f"{(x.get('pipeline') or {}).get('uuid')}_"
2400
- f"{(x.get('pipeline') or {}).get('repo_path')}"),
2886
+ lambda x: (
2887
+ f"{(x.get('pipeline') or {}).get('uuid')}_"
2888
+ f"{(x.get('pipeline') or {}).get('repo_path')}"
2889
+ ),
2401
2890
  )
2402
2891
 
2403
2892
  def to_dict_base(
@@ -2417,7 +2906,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2417
2906
  configuration=self.configuration or {},
2418
2907
  downstream_blocks=self.downstream_block_uuids,
2419
2908
  executor_config=self.executor_config,
2420
- executor_type=format_enum(self.executor_type) if self.executor_type else None,
2909
+ executor_type=(format_enum(self.executor_type) if self.executor_type else None),
2421
2910
  has_callback=self.has_callback,
2422
2911
  name=self.name,
2423
2912
  language=language,
@@ -2474,10 +2963,11 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2474
2963
  if include_outputs_use:
2475
2964
  data['outputs'] = self.outputs
2476
2965
 
2477
- if check_if_file_exists and not \
2478
- self.replicated_block and \
2479
- BlockType.GLOBAL_DATA_PRODUCT != self.type:
2480
-
2966
+ if (
2967
+ check_if_file_exists
2968
+ and not self.replicated_block
2969
+ and BlockType.GLOBAL_DATA_PRODUCT != self.type
2970
+ ):
2481
2971
  file_path = self.file_path
2482
2972
  if not os.path.isfile(file_path):
2483
2973
  data['error'] = dict(
@@ -2501,9 +2991,12 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2501
2991
  include_content: bool = False,
2502
2992
  include_outputs: bool = False,
2503
2993
  include_outputs_spark: bool = False,
2504
- sample_count: int = None,
2505
- block_cache: BlockCache = None,
2994
+ sample_count: Optional[int] = None,
2995
+ block_cache: Optional[BlockCache] = None,
2506
2996
  check_if_file_exists: bool = False,
2997
+ disable_output_preview: bool = False,
2998
+ exclude_blank_variable_uuids: bool = False,
2999
+ max_results: Optional[int] = None,
2507
3000
  **kwargs,
2508
3001
  ) -> Dict:
2509
3002
  data = self.to_dict_base(
@@ -2521,25 +3014,44 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2521
3014
 
2522
3015
  if include_outputs:
2523
3016
  include_outputs_use = include_outputs
2524
- if self.is_using_spark() and self.compute_management_enabled():
3017
+ if (
3018
+ self.is_using_spark()
3019
+ and self.compute_management_enabled()
3020
+ and self.pipeline
3021
+ and self.pipeline.type == PipelineType.PYSPARK
3022
+ ):
2525
3023
  include_outputs_use = include_outputs_use and include_outputs_spark
2526
3024
 
2527
3025
  if include_outputs_use:
2528
- data['outputs'] = await self.__outputs_async()
2529
-
2530
- if check_if_file_exists and not \
2531
- self.replicated_block and \
2532
- BlockType.GLOBAL_DATA_PRODUCT != self.type:
3026
+ if (
3027
+ disable_output_preview
3028
+ and self.configuration
3029
+ and self.configuration.get('disable_output_preview', False)
3030
+ ):
3031
+ data['outputs'] = [
3032
+ 'Output preview is disabled for this block. '
3033
+ 'To enable it, go to block settings.',
3034
+ ]
3035
+ else:
3036
+ data['outputs'] = await self.__outputs_async(
3037
+ exclude_blank_variable_uuids=exclude_blank_variable_uuids,
3038
+ max_results=max_results,
3039
+ )
2533
3040
 
2534
- file_path = self.file.file_path
2535
- if not os.path.isfile(file_path):
2536
- data['error'] = dict(
2537
- error='No such file or directory',
2538
- message='You may have moved it or changed its filename. '
2539
- 'Delete the current block to remove it from the pipeline '
2540
- 'or write code and save the pipeline to create a new file at '
2541
- f'{file_path}.',
2542
- )
3041
+ if (
3042
+ check_if_file_exists
3043
+ and not self.replicated_block
3044
+ and BlockType.GLOBAL_DATA_PRODUCT != self.type
3045
+ ):
3046
+ file_path = self.file.file_path
3047
+ if not os.path.isfile(file_path):
3048
+ data['error'] = dict(
3049
+ error='No such file or directory',
3050
+ message='You may have moved it or changed its filename. '
3051
+ 'Delete the current block to remove it from the pipeline '
3052
+ 'or write code and save the pipeline to create a new file at '
3053
+ f'{file_path}.',
3054
+ )
2543
3055
 
2544
3056
  if include_block_metadata:
2545
3057
  data['metadata'] = await self.metadata_async()
@@ -2570,9 +3082,8 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2570
3082
 
2571
3083
  check_upstream_block_order = kwargs.get('check_upstream_block_order', False)
2572
3084
  if 'upstream_blocks' in data and (
2573
- (check_upstream_block_order and
2574
- data['upstream_blocks'] != self.upstream_block_uuids) or
2575
- set(data['upstream_blocks']) != set(self.upstream_block_uuids)
3085
+ (check_upstream_block_order and data['upstream_blocks'] != self.upstream_block_uuids)
3086
+ or set(data['upstream_blocks']) != set(self.upstream_block_uuids)
2576
3087
  ):
2577
3088
  self.__update_upstream_blocks(
2578
3089
  data['upstream_blocks'],
@@ -2607,7 +3118,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2607
3118
  if 'has_callback' in data and data['has_callback'] != self.has_callback:
2608
3119
  self.has_callback = data['has_callback']
2609
3120
  if self.has_callback:
2610
- CallbackBlock.create(self.uuid)
3121
+ CallbackBlock.create(self.uuid, self.repo_path)
2611
3122
  self.__update_pipeline_block()
2612
3123
 
2613
3124
  if 'retry_config' in data and data['retry_config'] != self.retry_config:
@@ -2704,10 +3215,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2704
3215
  return list(visited)
2705
3216
 
2706
3217
  def run_upstream_blocks(
2707
- self,
2708
- from_notebook: bool = False,
2709
- incomplete_only: bool = False,
2710
- **kwargs
3218
+ self, from_notebook: bool = False, incomplete_only: bool = False, **kwargs
2711
3219
  ) -> None:
2712
3220
  def process_upstream_block(
2713
3221
  block: 'Block',
@@ -2717,10 +3225,12 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2717
3225
  root_blocks.append(block)
2718
3226
  return block.uuid
2719
3227
 
2720
- upstream_blocks = list(filter(
2721
- lambda x: not incomplete_only or BlockStatus.EXECUTED != x.status,
2722
- self.get_all_upstream_blocks(),
2723
- ))
3228
+ upstream_blocks = list(
3229
+ filter(
3230
+ lambda x: not incomplete_only or BlockStatus.EXECUTED != x.status,
3231
+ self.get_all_upstream_blocks(),
3232
+ )
3233
+ )
2724
3234
  root_blocks = []
2725
3235
  upstream_block_uuids = list(
2726
3236
  map(lambda x: process_upstream_block(x, root_blocks), upstream_blocks),
@@ -2758,10 +3268,11 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2758
3268
 
2759
3269
  self.dynamic_block_uuid = dynamic_block_uuid
2760
3270
 
2761
- if self.pipeline \
2762
- and PipelineType.INTEGRATION == self.pipeline.type \
2763
- and self.type in [BlockType.DATA_LOADER, BlockType.DATA_EXPORTER]:
2764
-
3271
+ if (
3272
+ self.pipeline
3273
+ and PipelineType.INTEGRATION == self.pipeline.type
3274
+ and self.type in [BlockType.DATA_LOADER, BlockType.DATA_EXPORTER]
3275
+ ):
2765
3276
  return
2766
3277
 
2767
3278
  test_functions = []
@@ -2804,7 +3315,9 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2804
3315
  test_function = getattr(self.module, func.__name__)
2805
3316
  try:
2806
3317
  sig = signature(test_function)
2807
- has_kwargs = any([p.kind == p.VAR_KEYWORD for p in sig.parameters.values()])
3318
+ has_kwargs = any([
3319
+ p.kind == p.VAR_KEYWORD for p in sig.parameters.values()
3320
+ ])
2808
3321
  if has_kwargs and global_vars is not None and len(global_vars) != 0:
2809
3322
  test_function(*outputs, **global_vars)
2810
3323
  else:
@@ -2815,11 +3328,13 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2815
3328
  stacktrace = traceback.format_exc()
2816
3329
 
2817
3330
  if from_notebook:
2818
- error_json = json.dumps(dict(
2819
- error=str(err),
2820
- message=error_message,
2821
- stacktrace=stacktrace.split('\n'),
2822
- ))
3331
+ error_json = json.dumps(
3332
+ dict(
3333
+ error=str(err),
3334
+ message=error_message,
3335
+ stacktrace=stacktrace.split('\n'),
3336
+ )
3337
+ )
2823
3338
  print(f'[__internal_test__]{error_json}')
2824
3339
  else:
2825
3340
  print('==============================================================')
@@ -2830,9 +3345,11 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2830
3345
  message = f'{tests_passed}/{len(test_functions)} tests passed.'
2831
3346
  if from_notebook:
2832
3347
  if len(test_functions) >= 1:
2833
- success_json = json.dumps(dict(
2834
- message=message,
2835
- ))
3348
+ success_json = json.dumps(
3349
+ dict(
3350
+ message=message,
3351
+ )
3352
+ )
2836
3353
  print(f'[__internal_test__]{success_json}')
2837
3354
  else:
2838
3355
  print('--------------------------------------------------------------')
@@ -2859,11 +3376,12 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2859
3376
  ) -> None:
2860
3377
  if self.pipeline is None:
2861
3378
  return
3379
+
2862
3380
  for uuid, data in variable_mapping.items():
2863
- if type(data) is pd.DataFrame:
3381
+ if isinstance(data, pd.DataFrame):
2864
3382
  if data.shape[1] > DATAFRAME_ANALYSIS_MAX_COLUMNS or shape_only:
2865
- self.pipeline.variable_manager.add_variable(
2866
- self.pipeline.uuid,
3383
+ self.variable_manager.add_variable(
3384
+ self.pipeline_uuid,
2867
3385
  self.uuid,
2868
3386
  uuid,
2869
3387
  dict(
@@ -2874,6 +3392,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2874
3392
  ),
2875
3393
  partition=execution_partition,
2876
3394
  variable_type=VariableType.DATAFRAME_ANALYSIS,
3395
+ disable_variable_type_inference=True,
2877
3396
  )
2878
3397
  continue
2879
3398
  if data.shape[0] > DATAFRAME_ANALYSIS_MAX_ROWS:
@@ -2884,14 +3403,15 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2884
3403
  data_for_analysis = data.reset_index(drop=True)
2885
3404
  try:
2886
3405
  from mage_ai.data_cleaner.data_cleaner import clean as clean_data
3406
+
2887
3407
  analysis = clean_data(
2888
3408
  data_for_analysis,
2889
3409
  df_original=data,
2890
3410
  transform=False,
2891
3411
  verbose=False,
2892
3412
  )
2893
- self.pipeline.variable_manager.add_variable(
2894
- self.pipeline.uuid,
3413
+ self.variable_manager.add_variable(
3414
+ self.pipeline_uuid,
2895
3415
  self.uuid,
2896
3416
  uuid,
2897
3417
  dict(
@@ -2903,14 +3423,15 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2903
3423
  partition=execution_partition,
2904
3424
  variable_type=VariableType.DATAFRAME_ANALYSIS,
2905
3425
  )
2906
- except Exception:
2907
- pass
3426
+ except Exception as err:
3427
+ if is_debug():
3428
+ raise err
2908
3429
  # TODO: we use to silently fail, but it looks bad when using BigQuery
2909
3430
  # print('\nFailed to analyze dataframe:')
2910
3431
  # print(traceback.format_exc())
2911
- elif type(data) is pl.DataFrame:
2912
- self.pipeline.variable_manager.add_variable(
2913
- self.pipeline.uuid,
3432
+ elif isinstance(data, pl.DataFrame):
3433
+ self.variable_manager.add_variable(
3434
+ self.pipeline_uuid,
2914
3435
  self.uuid,
2915
3436
  uuid,
2916
3437
  dict(
@@ -2921,6 +3442,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2921
3442
  ),
2922
3443
  partition=execution_partition,
2923
3444
  variable_type=VariableType.DATAFRAME_ANALYSIS,
3445
+ disable_variable_type_inference=True,
2924
3446
  )
2925
3447
 
2926
3448
  def set_global_vars(self, global_vars: Dict) -> None:
@@ -2930,9 +3452,15 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2930
3452
 
2931
3453
  def __consolidate_variables(self, variable_mapping: Dict) -> Dict:
2932
3454
  # Consolidate print variables
2933
- output_variables = {k: v for k, v in variable_mapping.items() if is_output_variable(k)}
2934
- print_variables = {k: v for k, v in variable_mapping.items()
2935
- if is_valid_print_variable(k, v, self.uuid)}
3455
+ if BlockType.SCRATCHPAD == self.type:
3456
+ output_variables = {}
3457
+ print_variables = variable_mapping.copy()
3458
+ else:
3459
+ output_variables = {k: v for k, v in variable_mapping.items() if is_output_variable(k)}
3460
+ print_variables = {
3461
+ k: v for k, v in variable_mapping.items()
3462
+ if is_valid_print_variable(k, v, self.uuid)
3463
+ }
2936
3464
 
2937
3465
  print_variables_keys = sorted(print_variables.keys(), key=lambda k: int(k.split('_')[-1]))
2938
3466
 
@@ -2967,6 +3495,10 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2967
3495
  save_variable_and_reset_state()
2968
3496
  continue
2969
3497
 
3498
+ if json_value.get('msg_type') == 'status':
3499
+ # Do not save status messages
3500
+ continue
3501
+
2970
3502
  if state['msg_key'] is not None and json_value['msg_type'] != state['msg_type']:
2971
3503
  save_variable_and_reset_state()
2972
3504
 
@@ -2994,6 +3526,8 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
2994
3526
  global_vars: Dict = None,
2995
3527
  dynamic_block_index: int = None,
2996
3528
  ) -> Dict:
3529
+ from mage_ai.data_preparation.models.block.remote.models import RemoteBlock
3530
+
2997
3531
  """
2998
3532
  Enriches the provided global variables dictionary with additional context, Spark session,
2999
3533
  environment, configuration, and an empty context dictionary.
@@ -3021,8 +3555,9 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3021
3555
  """
3022
3556
  if global_vars is None:
3023
3557
  global_vars = dict()
3024
- if ((self.pipeline is not None and self.pipeline.type == PipelineType.DATABRICKS) or
3025
- is_spark_env()):
3558
+ if (
3559
+ self.pipeline is not None and self.pipeline.type == PipelineType.DATABRICKS
3560
+ ) or is_spark_env():
3026
3561
  if not global_vars.get('spark'):
3027
3562
  spark = self.get_spark_session()
3028
3563
  if spark is not None:
@@ -3034,12 +3569,21 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3034
3569
  global_vars['context'] = dict()
3035
3570
 
3036
3571
  # Add pipeline uuid and block uuid to global_vars
3037
- global_vars['pipeline_uuid'] = self.pipeline.uuid if self.pipeline else None
3572
+ global_vars['pipeline_uuid'] = self.pipeline_uuid
3038
3573
  global_vars['block_uuid'] = self.uuid
3039
3574
 
3040
3575
  if dynamic_block_index is not None:
3041
3576
  global_vars['dynamic_block_index'] = dynamic_block_index
3042
3577
 
3578
+ # Remote blocks
3579
+ if global_vars.get('remote_blocks'):
3580
+ global_vars['remote_blocks'] = [
3581
+ RemoteBlock.load(
3582
+ **remote_block_dict,
3583
+ ).get_outputs()
3584
+ for remote_block_dict in global_vars['remote_blocks']
3585
+ ]
3586
+
3043
3587
  self.global_vars = global_vars
3044
3588
 
3045
3589
  return global_vars
@@ -3047,18 +3591,15 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3047
3591
  def get_spark_session(self):
3048
3592
  if not SPARK_ENABLED:
3049
3593
  return None
3050
- if self.spark_init and (not self.pipeline or
3051
- not self.pipeline.spark_config):
3594
+ if self.spark_init and (not self.pipeline or not self.pipeline.spark_config):
3052
3595
  return self.spark
3053
3596
 
3054
3597
  try:
3055
3598
  if self.pipeline and self.pipeline.spark_config:
3056
- spark_config = SparkConfig.load(
3057
- config=self.pipeline.spark_config)
3599
+ spark_config = SparkConfig.load(config=self.pipeline.spark_config)
3058
3600
  else:
3059
3601
  repo_config = RepoConfig(repo_path=self.repo_path)
3060
- spark_config = SparkConfig.load(
3061
- config=repo_config.spark_config)
3602
+ spark_config = SparkConfig.load(config=repo_config.spark_config)
3062
3603
  self.spark = get_spark_session(spark_config)
3063
3604
  except Exception:
3064
3605
  self.spark = None
@@ -3083,18 +3624,19 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3083
3624
  spark_config = SparkConfig.load(config=spark_config)
3084
3625
  if spark_config.use_custom_session:
3085
3626
  return global_vars.get('context', dict()).get(
3086
- spark_config.custom_session_var_name, spark)
3627
+ spark_config.custom_session_var_name, spark
3628
+ )
3087
3629
  return spark
3088
3630
 
3089
- def __store_variables_prepare(
3631
+ def __get_variable_uuids(
3090
3632
  self,
3091
- variable_mapping: Dict,
3092
- execution_partition: str = None,
3093
- override: bool = False,
3094
- override_outputs: bool = False,
3095
- dynamic_block_index: int = None,
3096
- dynamic_block_uuid: str = None,
3097
- ) -> Dict:
3633
+ dynamic_block_index: Optional[int] = None,
3634
+ dynamic_block_uuid: Optional[str] = None,
3635
+ execution_partition: Optional[str] = None,
3636
+ ) -> Tuple[List[str], str, bool]:
3637
+ if self.pipeline is None:
3638
+ return []
3639
+
3098
3640
  self.dynamic_block_uuid = dynamic_block_uuid
3099
3641
 
3100
3642
  block_uuid, changed = uuid_for_output_variables(
@@ -3104,21 +3646,37 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3104
3646
  dynamic_block_uuid=dynamic_block_uuid,
3105
3647
  )
3106
3648
 
3107
- if self.pipeline is None:
3108
- return
3649
+ return (
3650
+ self.variable_manager.get_variables_by_block(
3651
+ self.pipeline_uuid,
3652
+ block_uuid=block_uuid,
3653
+ partition=execution_partition,
3654
+ clean_block_uuid=not changed,
3655
+ ),
3656
+ block_uuid,
3657
+ changed,
3658
+ )
3109
3659
 
3110
- all_variables = self.pipeline.variable_manager.get_variables_by_block(
3111
- self.pipeline.uuid,
3112
- block_uuid=block_uuid,
3113
- partition=execution_partition,
3114
- clean_block_uuid=not changed,
3660
+ def __store_variables_prepare(
3661
+ self,
3662
+ variable_mapping: Dict,
3663
+ execution_partition: Optional[str] = None,
3664
+ override: bool = False,
3665
+ override_outputs: bool = False,
3666
+ dynamic_block_index: Optional[int] = None,
3667
+ dynamic_block_uuid: Optional[str] = None,
3668
+ ) -> Dict:
3669
+ variable_uuids, _block_uuid, _changed = self.__get_variable_uuids(
3670
+ dynamic_block_index=dynamic_block_index,
3671
+ dynamic_block_uuid=dynamic_block_uuid,
3672
+ execution_partition=execution_partition,
3115
3673
  )
3116
3674
 
3117
3675
  variable_mapping = self.__consolidate_variables(variable_mapping)
3118
3676
 
3119
3677
  variable_names = [clean_name_orig(v) for v in variable_mapping]
3120
3678
  removed_variables = []
3121
- for v in all_variables:
3679
+ for v in variable_uuids:
3122
3680
  if v in variable_names:
3123
3681
  continue
3124
3682
 
@@ -3131,16 +3689,97 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3131
3689
  variable_mapping=variable_mapping,
3132
3690
  )
3133
3691
 
3692
+ def delete_variables(
3693
+ self,
3694
+ dynamic_block_index: Optional[int] = None,
3695
+ dynamic_block_uuid: Optional[str] = None,
3696
+ execution_partition: Optional[str] = None,
3697
+ variable_uuids: Optional[List[str]] = None,
3698
+ ) -> None:
3699
+ if self.pipeline is None:
3700
+ return
3701
+
3702
+ variable_uuids_all, block_uuid, _changed = self.__get_variable_uuids(
3703
+ dynamic_block_index=dynamic_block_index,
3704
+ dynamic_block_uuid=dynamic_block_uuid,
3705
+ execution_partition=execution_partition,
3706
+ )
3707
+
3708
+ for variable_uuid in variable_uuids or variable_uuids_all:
3709
+ variable_object = self.variable_manager.get_variable_object(
3710
+ self.pipeline_uuid,
3711
+ block_uuid,
3712
+ variable_uuid,
3713
+ partition=execution_partition,
3714
+ )
3715
+ write_policy = (
3716
+ self.write_settings.batch_settings.mode
3717
+ if self.write_settings and self.write_settings.batch_settings
3718
+ else None
3719
+ )
3720
+ if write_policy and variable_object.data_exists():
3721
+ if ExportWritePolicy.FAIL == write_policy:
3722
+ raise Exception(f'Write policy for block {self.uuid} is {write_policy}.')
3723
+ elif ExportWritePolicy.APPEND == write_policy:
3724
+ return
3725
+
3726
+ variable_object.delete()
3727
+
3728
+ def aggregate_summary_info(self, execution_partition: Optional[str] = None):
3729
+ """
3730
+ Run this only after executing blocks in a notebook so that reading pipelines
3731
+ don’t take forever to load while waiting for all the nested variable folders
3732
+ to be read.
3733
+ """
3734
+ if not VARIABLE_DATA_OUTPUT_META_CACHE or not self.variable_manager:
3735
+ return
3736
+
3737
+ aggregate_summary_info_for_all_variables(
3738
+ self.variable_manager,
3739
+ self.pipeline_uuid,
3740
+ self.uuid,
3741
+ partition=execution_partition,
3742
+ )
3743
+
3134
3744
  def store_variables(
3135
3745
  self,
3136
3746
  variable_mapping: Dict,
3137
- execution_partition: str = None,
3747
+ clean_variable_uuid: Optional[bool] = True,
3748
+ dynamic_block_index: Optional[int] = None,
3749
+ dynamic_block_uuid: Optional[str] = None,
3750
+ execution_partition: Optional[str] = None,
3138
3751
  override: bool = False,
3139
3752
  override_outputs: bool = False,
3140
- spark=None,
3141
- dynamic_block_index: int = None,
3142
- dynamic_block_uuid: str = None,
3143
- ) -> None:
3753
+ skip_delete: Optional[bool] = None,
3754
+ spark: Optional[Any] = None,
3755
+ save_variable_types_only: Optional[Any] = None,
3756
+ ) -> Optional[List[Variable]]:
3757
+ block_uuid, changed = uuid_for_output_variables(
3758
+ self,
3759
+ block_uuid=self.uuid,
3760
+ dynamic_block_index=dynamic_block_index,
3761
+ )
3762
+
3763
+ is_dynamic = is_dynamic_block(self)
3764
+ is_dynamic_child = is_dynamic_block_child(self)
3765
+
3766
+ shared_args = dict(
3767
+ clean_block_uuid=not changed,
3768
+ clean_variable_uuid=not is_dynamic and not is_dynamic_child and clean_variable_uuid,
3769
+ partition=execution_partition,
3770
+ )
3771
+
3772
+ if save_variable_types_only:
3773
+ for variable_uuid, variable_types in variable_mapping.items():
3774
+ self.variable_manager.add_variable_types(
3775
+ self.pipeline_uuid,
3776
+ block_uuid,
3777
+ variable_uuid,
3778
+ variable_types,
3779
+ **shared_args,
3780
+ )
3781
+ return []
3782
+
3144
3783
  variables_data = self.__store_variables_prepare(
3145
3784
  variable_mapping,
3146
3785
  execution_partition,
@@ -3149,40 +3788,45 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3149
3788
  dynamic_block_index=dynamic_block_index,
3150
3789
  )
3151
3790
 
3152
- block_uuid, changed = uuid_for_output_variables(
3153
- self,
3154
- block_uuid=self.uuid,
3155
- dynamic_block_index=dynamic_block_index,
3156
- )
3157
-
3158
- is_dynamic_child = is_dynamic_block_child(self)
3159
- if is_dynamic_child:
3791
+ if not skip_delete and is_dynamic_child:
3792
+ # execute_block_function will take care of this if decorated function is a generator
3160
3793
  delete_variable_objects_for_dynamic_child(
3161
3794
  self,
3162
3795
  dynamic_block_index=dynamic_block_index,
3163
3796
  execution_partition=execution_partition,
3164
3797
  )
3165
3798
 
3799
+ variables = []
3166
3800
  for uuid, data in variables_data['variable_mapping'].items():
3167
- if spark is not None and self.pipeline.type == PipelineType.PYSPARK \
3168
- and type(data) is pd.DataFrame:
3801
+ if (
3802
+ spark is not None
3803
+ and self.pipeline is not None
3804
+ and self.pipeline.type == PipelineType.PYSPARK
3805
+ and isinstance(data, pd.DataFrame)
3806
+ ):
3169
3807
  data = spark.createDataFrame(data)
3170
- self.pipeline.variable_manager.add_variable(
3171
- self.pipeline.uuid,
3172
- block_uuid,
3173
- uuid,
3174
- data,
3175
- partition=execution_partition,
3176
- clean_block_uuid=not changed,
3177
- )
3178
3808
 
3179
- if not is_dynamic_child:
3180
- for uuid in variables_data['removed_variables']:
3181
- self.pipeline.variable_manager.delete_variable(
3182
- self.pipeline.uuid,
3809
+ variables.append(
3810
+ self.variable_manager.add_variable(
3811
+ self.pipeline_uuid,
3183
3812
  block_uuid,
3184
3813
  uuid,
3814
+ data,
3815
+ resource_usage=self.resource_usage,
3816
+ write_batch_settings=self.write_batch_settings,
3817
+ write_chunks=self.write_chunks,
3818
+ **shared_args,
3185
3819
  )
3820
+ )
3821
+ if not skip_delete and not is_dynamic_child and variables_data.get('removed_variables'):
3822
+ self.delete_variables(
3823
+ dynamic_block_index=dynamic_block_index,
3824
+ dynamic_block_uuid=dynamic_block_uuid,
3825
+ execution_partition=execution_partition,
3826
+ variable_uuids=variables_data['removed_variables'],
3827
+ )
3828
+
3829
+ return variables
3186
3830
 
3187
3831
  async def store_variables_async(
3188
3832
  self,
@@ -3220,19 +3864,22 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3220
3864
  if spark is not None and type(data) is pd.DataFrame:
3221
3865
  data = spark.createDataFrame(data)
3222
3866
 
3223
- await self.pipeline.variable_manager.add_variable_async(
3224
- self.pipeline.uuid,
3867
+ await self.variable_manager.add_variable_async(
3868
+ self.pipeline_uuid,
3225
3869
  block_uuid,
3226
3870
  uuid,
3227
3871
  data,
3228
3872
  partition=execution_partition,
3229
3873
  clean_block_uuid=not changed,
3874
+ write_batch_settings=self.write_batch_settings,
3875
+ write_chunks=self.write_chunks,
3876
+ resource_usage=self.resource_usage,
3230
3877
  )
3231
3878
 
3232
3879
  if not is_dynamic_child:
3233
3880
  for uuid in variables_data['removed_variables']:
3234
- self.pipeline.variable_manager.delete_variable(
3235
- self.pipeline.uuid,
3881
+ self.variable_manager.delete_variable(
3882
+ self.pipeline_uuid,
3236
3883
  block_uuid,
3237
3884
  uuid,
3238
3885
  clean_block_uuid=not changed,
@@ -3261,8 +3908,8 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3261
3908
  for b in self.upstream_blocks:
3262
3909
  for v in b.output_variables(execution_partition=execution_partition):
3263
3910
  objs.append(
3264
- self.pipeline.variable_manager.get_variable_object(
3265
- self.pipeline.uuid,
3911
+ self.get_variable_object(
3912
+ self.pipeline_uuid,
3266
3913
  b.uuid,
3267
3914
  v,
3268
3915
  partition=execution_partition,
@@ -3278,6 +3925,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3278
3925
  from_notebook: bool = False,
3279
3926
  global_vars: Dict = None,
3280
3927
  input_args: List[Any] = None,
3928
+ max_results: Optional[int] = None,
3281
3929
  block_uuid: str = None,
3282
3930
  ) -> List[str]:
3283
3931
  return output_variables(
@@ -3289,6 +3937,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3289
3937
  from_notebook=from_notebook,
3290
3938
  global_vars=global_vars,
3291
3939
  input_args=input_args,
3940
+ max_results=max_results,
3292
3941
  )
3293
3942
 
3294
3943
  def output_variable_objects(
@@ -3312,12 +3961,15 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3312
3961
  if len(output_variables) == 0:
3313
3962
  return []
3314
3963
 
3315
- variable_objects = [self.pipeline.variable_manager.get_variable_object(
3316
- self.pipeline.uuid,
3317
- self.uuid,
3318
- v,
3319
- partition=execution_partition,
3320
- ) for v in output_variables]
3964
+ variable_objects = [
3965
+ self.variable_manager.get_variable_object(
3966
+ self.pipeline_uuid,
3967
+ self.uuid,
3968
+ v,
3969
+ partition=execution_partition,
3970
+ )
3971
+ for v in output_variables
3972
+ ]
3321
3973
  if variable_type is not None:
3322
3974
  variable_objects = [v for v in variable_objects if v.variable_type == variable_type]
3323
3975
  return variable_objects
@@ -3353,8 +4005,8 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3353
4005
  ) -> Any:
3354
4006
  if self.pipeline is None:
3355
4007
  return []
3356
- return self.pipeline.variable_manager.get_variable_object(
3357
- self.pipeline.uuid,
4008
+ return self.variable_manager.get_variable_object(
4009
+ self.pipeline_uuid,
3358
4010
  self.uuid,
3359
4011
  variable_uuid,
3360
4012
  partition=execution_partition,
@@ -3391,32 +4043,41 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3391
4043
  self.uuid = new_uuid
3392
4044
 
3393
4045
  # This file has a path in its file_source that must be updated.
3394
- if project_platform_activated() and \
3395
- self.file_source_path() and \
3396
- add_absolute_path(self.file_source_path()) == self.file_path:
3397
-
3398
- # /home/src/data-vault/perftools/mage/data_loaders/team/illusory_glitter
4046
+ if (
4047
+ project_platform_activated()
4048
+ and self.file_source_path()
4049
+ and add_absolute_path(self.file_source_path()) == self.file_path
4050
+ ):
4051
+ # /home/src/data-vault/perftools/mage/data_loaders/team/illusory_glitter
3399
4052
  old_file_path_without_extension = str(Path(old_file_path).with_suffix(''))
3400
4053
  # /home/src/data-vault/perftools/mage/data_loaders/team
3401
- old_file_path_without_uuid = str(Path(old_file_path_without_extension.replace(
3402
- str(Path(old_uuid)),
3403
- '',
3404
- )))
4054
+ old_file_path_without_uuid = str(
4055
+ Path(
4056
+ old_file_path_without_extension.replace(
4057
+ str(Path(old_uuid)),
4058
+ '',
4059
+ )
4060
+ )
4061
+ )
3405
4062
 
3406
- # perftools/mage/data_loaders/team
4063
+ # perftools/mage/data_loaders/team
3407
4064
  old_file_path_without_repo_path = remove_base_repo_path(old_file_path_without_uuid)
3408
- # perftools/mage
4065
+ # perftools/mage
3409
4066
  path_without_block_directory = str(old_file_path_without_repo_path).split(
3410
4067
  directory_name,
3411
4068
  )[0]
3412
4069
 
3413
4070
  file_extension_new = Path(self.uuid).suffix or file_extension
3414
4071
  # perftools/mage/data_loaders/load_titanic.py
3415
- new_path = str(Path(os.path.join(
3416
- path_without_block_directory,
3417
- directory_name,
3418
- str(Path(self.uuid).with_suffix('')),
3419
- )).with_suffix(file_extension_new))
4072
+ new_path = str(
4073
+ Path(
4074
+ os.path.join(
4075
+ path_without_block_directory,
4076
+ directory_name,
4077
+ str(Path(self.uuid).with_suffix('')),
4078
+ )
4079
+ ).with_suffix(file_extension_new)
4080
+ )
3420
4081
 
3421
4082
  configuration = self.configuration or {}
3422
4083
  if not configuration.get('file_source'):
@@ -3425,7 +4086,16 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3425
4086
  self.configuration = configuration
3426
4087
 
3427
4088
  # This has to be here
3428
- new_file_path = self.file_path
4089
+ new_file_path, new_file_path_relative = self.build_file_path_directory(
4090
+ block_uuid=new_uuid,
4091
+ )
4092
+
4093
+ configuration = self.configuration or {}
4094
+ if not configuration.get('file_source'):
4095
+ configuration['file_source'] = {}
4096
+ configuration['file_path'] = new_file_path_relative
4097
+ configuration['file_source']['path'] = new_file_path_relative
4098
+ self.configuration = configuration
3429
4099
 
3430
4100
  if self.pipeline is not None:
3431
4101
  DX_PRINTER.critical(
@@ -3437,7 +4107,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3437
4107
  name=self.name,
3438
4108
  uuid=self.uuid,
3439
4109
  file_path=new_file_path,
3440
- pipeline=self.pipeline.uuid,
4110
+ pipeline=self.pipeline_uuid,
3441
4111
  repo_path=self.pipeline.repo_path,
3442
4112
  configuration=self.configuration,
3443
4113
  __uuid='__update_name',
@@ -3453,20 +4123,23 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3453
4123
  )
3454
4124
 
3455
4125
  if not self.replicated_block and BlockType.GLOBAL_DATA_PRODUCT != self.type:
3456
- if os.path.exists(new_file_path):
3457
- raise Exception(f'Block {new_uuid} already exists. Please use a different name.')
4126
+ if new_file_path and os.path.exists(new_file_path):
4127
+ raise Exception(
4128
+ f'Block {new_uuid} already exists at {new_file_path}. '
4129
+ 'Please use a different name.'
4130
+ )
3458
4131
 
3459
4132
  parent_dir = os.path.dirname(new_file_path)
3460
4133
  os.makedirs(parent_dir, exist_ok=True)
3461
4134
 
3462
4135
  if detach:
3463
- """"
4136
+ """ "
3464
4137
  Detaching a block creates a copy of the block file while keeping the existing block
3465
4138
  file the same. Without detaching a block, the existing block file is simply renamed.
3466
4139
  """
3467
4140
  with open(new_file_path, 'w') as f:
3468
4141
  f.write(block_content)
3469
- else:
4142
+ elif os.path.exists(old_file_path):
3470
4143
  os.rename(old_file_path, new_file_path)
3471
4144
 
3472
4145
  if self.pipeline is not None:
@@ -3477,11 +4150,12 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3477
4150
  from mage_ai.data_preparation.models.block.block_factory import (
3478
4151
  BlockFactory,
3479
4152
  )
4153
+
3480
4154
  """"
3481
4155
  New block added to pipeline, so it must be added to the block cache.
3482
4156
  Old block no longer in pipeline, so it must be removed from block cache.
3483
4157
  """
3484
- cache.add_pipeline(self, self.pipeline)
4158
+ cache.add_pipeline(self, self.pipeline, self.pipeline.repo_path)
3485
4159
  old_block = BlockFactory.get_block(
3486
4160
  old_uuid,
3487
4161
  old_uuid,
@@ -3489,12 +4163,16 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
3489
4163
  language=self.language,
3490
4164
  pipeline=self.pipeline,
3491
4165
  )
3492
- cache.remove_pipeline(old_block, self.pipeline.uuid, self.pipeline.repo_path)
4166
+ cache.remove_pipeline(old_block, self.pipeline_uuid, self.pipeline.repo_path)
3493
4167
  else:
3494
- cache.move_pipelines(self, dict(
3495
- type=self.type,
3496
- uuid=old_uuid,
3497
- ))
4168
+ cache.move_pipelines(
4169
+ self,
4170
+ dict(
4171
+ type=self.type,
4172
+ uuid=old_uuid,
4173
+ ),
4174
+ self.pipeline.repo_path,
4175
+ )
3498
4176
 
3499
4177
  def __update_pipeline_block(self, widget=False) -> None:
3500
4178
  if self.pipeline is None:
@@ -3568,7 +4246,9 @@ class SensorBlock(Block):
3568
4246
  block_function: Callable,
3569
4247
  input_vars: List,
3570
4248
  from_notebook: bool = False,
3571
- global_vars: Dict = None,
4249
+ global_vars: Optional[Dict] = None,
4250
+ *args,
4251
+ **kwargs,
3572
4252
  ) -> List:
3573
4253
  if from_notebook:
3574
4254
  return super().execute_block_function(
@@ -3584,8 +4264,9 @@ class SensorBlock(Block):
3584
4264
  use_global_vars = has_kwargs and global_vars is not None and len(global_vars) != 0
3585
4265
  args = input_vars if has_args else []
3586
4266
  while True:
3587
- condition = block_function(*args, **global_vars) \
3588
- if use_global_vars else block_function()
4267
+ condition = (
4268
+ block_function(*args, **global_vars) if use_global_vars else block_function()
4269
+ )
3589
4270
  if condition:
3590
4271
  break
3591
4272
  print('Sensor sleeping for 1 minute...')
@@ -3605,7 +4286,7 @@ class AddonBlock(Block):
3605
4286
  global_vars = merge_dict(
3606
4287
  global_vars or dict(),
3607
4288
  dict(
3608
- pipeline_uuid=self.pipeline.uuid,
4289
+ pipeline_uuid=self.pipeline_uuid,
3609
4290
  block_uuid=self.uuid,
3610
4291
  pipeline_run=pipeline_run,
3611
4292
  ),
@@ -3624,9 +4305,7 @@ class AddonBlock(Block):
3624
4305
  global_vars = merge_dict(parent_block.global_vars, global_vars)
3625
4306
  global_vars['parent_block_uuid'] = parent_block.uuid
3626
4307
 
3627
- if parent_block.pipeline and \
3628
- PipelineType.INTEGRATION == parent_block.pipeline.type:
3629
-
4308
+ if parent_block.pipeline and PipelineType.INTEGRATION == parent_block.pipeline.type:
3630
4309
  template_runtime_configuration = parent_block.template_runtime_configuration
3631
4310
  index = template_runtime_configuration.get('index', None)
3632
4311
  is_last_block_run = template_runtime_configuration.get('is_last_block_run', False)
@@ -3652,7 +4331,7 @@ class ConditionalBlock(AddonBlock):
3652
4331
  global_vars: Dict = None,
3653
4332
  logger: Logger = None,
3654
4333
  logging_tags: Dict = None,
3655
- **kwargs
4334
+ **kwargs,
3656
4335
  ) -> bool:
3657
4336
  with self._redirect_streams(
3658
4337
  logger=logger,
@@ -3696,11 +4375,11 @@ class ConditionalBlock(AddonBlock):
3696
4375
 
3697
4376
  class CallbackBlock(AddonBlock):
3698
4377
  @classmethod
3699
- def create(cls, orig_block_name) -> 'CallbackBlock':
4378
+ def create(cls, orig_block_name, repo_path: str) -> 'CallbackBlock':
3700
4379
  return Block.create(
3701
4380
  f'{clean_name_orig(orig_block_name)}_callback',
3702
4381
  BlockType.CALLBACK,
3703
- get_repo_path(),
4382
+ repo_path,
3704
4383
  language=BlockLanguage.PYTHON,
3705
4384
  )
3706
4385
 
@@ -3719,7 +4398,7 @@ class CallbackBlock(AddonBlock):
3719
4398
  from_notebook: bool = False,
3720
4399
  metadata: Dict = None,
3721
4400
  upstream_block_uuids_override: List[str] = None,
3722
- **kwargs
4401
+ **kwargs,
3723
4402
  ) -> None:
3724
4403
  with self._redirect_streams(
3725
4404
  logger=logger,
@@ -3731,7 +4410,7 @@ class CallbackBlock(AddonBlock):
3731
4410
  global_vars,
3732
4411
  parent_block,
3733
4412
  dynamic_block_index=dynamic_block_index,
3734
- **kwargs
4413
+ **kwargs,
3735
4414
  )
3736
4415
 
3737
4416
  callback_functions = []
@@ -3813,12 +4492,10 @@ class CallbackBlock(AddonBlock):
3813
4492
  # As of version 0.8.81, callback functions have access to the parent block’s
3814
4493
  # data output. If the callback function has any positional arguments, we will
3815
4494
  # pass in the input variables as positional arguments.
3816
- if not input_vars or any(
3817
- [
3818
- p.kind not in [p.KEYWORD_ONLY, p.VAR_KEYWORD]
3819
- for p in sig.parameters.values()
3820
- ]
3821
- ):
4495
+ if not input_vars or any([
4496
+ p.kind not in [p.KEYWORD_ONLY, p.VAR_KEYWORD]
4497
+ for p in sig.parameters.values()
4498
+ ]):
3822
4499
  inner_function(
3823
4500
  *input_vars,
3824
4501
  **callback_kwargs,
@@ -3827,9 +4504,12 @@ class CallbackBlock(AddonBlock):
3827
4504
  # a user has already written callback functions with only keyword arguments.
3828
4505
  else:
3829
4506
  inner_function(
3830
- **merge_dict(callback_kwargs, dict(
3831
- __input=outputs_from_input_vars,
3832
- )),
4507
+ **merge_dict(
4508
+ callback_kwargs,
4509
+ dict(
4510
+ __input=outputs_from_input_vars,
4511
+ ),
4512
+ ),
3833
4513
  )
3834
4514
 
3835
4515
  def update_content(self, content, widget=False) -> 'CallbackBlock':
@@ -3852,9 +4532,11 @@ class CallbackBlock(AddonBlock):
3852
4532
  def custom_code(callback_status: CallbackStatus = CallbackStatus.SUCCESS, *args, **kwargs):
3853
4533
  # If the decorator is just @callback with no arguments, default to success callback
3854
4534
  if isfunction(callback_status):
4535
+
3855
4536
  def func(callback_status_inner, *args, **kwargs):
3856
4537
  if callback_status_inner == CallbackStatus.SUCCESS:
3857
4538
  return callback_status(*args, **kwargs)
4539
+
3858
4540
  decorated_functions.append(func)
3859
4541
  return func
3860
4542
 
@@ -3868,6 +4550,7 @@ class CallbackBlock(AddonBlock):
3868
4550
  def func(callback_status_inner):
3869
4551
  if callback_status_inner == callback_status:
3870
4552
  return function
4553
+
3871
4554
  decorated_functions.append(func)
3872
4555
 
3873
4556
  return inner