mage-ai 0.8.91__py3-none-any.whl → 0.8.93__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 (103) hide show
  1. mage_ai/api/policies/BlockPolicy.py +1 -0
  2. mage_ai/api/policies/PipelinePolicy.py +6 -0
  3. mage_ai/api/presenters/BlockPresenter.py +1 -0
  4. mage_ai/api/presenters/PipelinePresenter.py +2 -0
  5. mage_ai/data_preparation/decorators.py +4 -0
  6. mage_ai/data_preparation/executors/block_executor.py +68 -3
  7. mage_ai/data_preparation/models/block/__init__.py +212 -67
  8. mage_ai/data_preparation/models/block/constants.py +3 -1
  9. mage_ai/data_preparation/models/constants.py +8 -8
  10. mage_ai/data_preparation/models/pipeline.py +83 -5
  11. mage_ai/data_preparation/repo_manager.py +5 -2
  12. mage_ai/data_preparation/shared/constants.py +2 -1
  13. mage_ai/data_preparation/templates/conditionals/base.jinja +11 -0
  14. mage_ai/data_preparation/templates/constants.py +7 -0
  15. mage_ai/io/mssql.py +11 -1
  16. mage_ai/io/sql.py +8 -1
  17. mage_ai/orchestration/db/migrations/versions/dfe49d040487_add_condition_failed_status_to_block_.py +39 -0
  18. mage_ai/orchestration/db/models/schedules.py +5 -1
  19. mage_ai/orchestration/pipeline_scheduler.py +27 -17
  20. mage_ai/server/api/downloads.py +64 -0
  21. mage_ai/server/constants.py +1 -1
  22. mage_ai/server/execution_manager.py +3 -2
  23. mage_ai/server/frontend_dist/404.html +2 -2
  24. mage_ai/server/frontend_dist/404.html.html +2 -2
  25. mage_ai/server/frontend_dist/_next/static/chunks/1424-321c8c08a2b05c19.js +1 -0
  26. mage_ai/server/frontend_dist/_next/static/chunks/2786-2b3ad2cf216fae42.js +1 -0
  27. mage_ai/server/frontend_dist/_next/static/chunks/3714-3bd2a8c979d6d820.js +1 -0
  28. mage_ai/server/frontend_dist/_next/static/chunks/3752-8f15fe0ca9c23cf4.js +1 -0
  29. mage_ai/server/frontend_dist/_next/static/chunks/{4476-cdae7a65db573bb7.js → 4476-c1a62e69cd8e14d5.js} +1 -1
  30. mage_ai/server/frontend_dist/_next/static/chunks/pages/_app-8aaee96edc252aa3.js +1 -0
  31. mage_ai/server/frontend_dist/_next/static/chunks/pages/manage-88c03376d807012e.js +1 -0
  32. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipeline-runs-915825c19bf42fa1.js +1 -0
  33. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/backfills-bf2d83dabe1bd25a.js +1 -0
  34. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-f0940870ff5a17f6.js +1 -0
  35. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/logs-8ee12ce8362ed576.js +1 -0
  36. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runs-a64f7a0aba0f481d.js +1 -0
  37. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runtime-3a3a115ab1a86e2f.js +1 -0
  38. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/runs/[run]-160881dab5ef66d8.js +1 -0
  39. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/runs-69d63c14abf8cf68.js +1 -0
  40. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/syncs-6092226e191dd720.js +1 -0
  41. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines-549f4708f2912a7a.js +1 -0
  42. mage_ai/server/frontend_dist/_next/static/zlCBBK90aKYZtPlYLj9_T/_buildManifest.js +1 -0
  43. mage_ai/server/frontend_dist/files.html +2 -2
  44. mage_ai/server/frontend_dist/index.html +2 -2
  45. mage_ai/server/frontend_dist/manage/users/[user].html +2 -2
  46. mage_ai/server/frontend_dist/manage/users.html +2 -2
  47. mage_ai/server/frontend_dist/manage.html +2 -2
  48. mage_ai/server/frontend_dist/pipeline-runs.html +2 -2
  49. mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills/[...slug].html +2 -2
  50. mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills.html +2 -2
  51. mage_ai/server/frontend_dist/pipelines/[pipeline]/edit.html +2 -2
  52. mage_ai/server/frontend_dist/pipelines/[pipeline]/logs.html +2 -2
  53. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runs.html +2 -2
  54. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runtime.html +2 -2
  55. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors.html +2 -2
  56. mage_ai/server/frontend_dist/pipelines/[pipeline]/runs/[run].html +2 -2
  57. mage_ai/server/frontend_dist/pipelines/[pipeline]/runs.html +2 -2
  58. mage_ai/server/frontend_dist/pipelines/[pipeline]/settings.html +2 -2
  59. mage_ai/server/frontend_dist/pipelines/[pipeline]/syncs.html +2 -2
  60. mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers/[...slug].html +2 -2
  61. mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers.html +2 -2
  62. mage_ai/server/frontend_dist/pipelines/[pipeline].html +2 -2
  63. mage_ai/server/frontend_dist/pipelines.html +2 -2
  64. mage_ai/server/frontend_dist/settings/account/profile.html +2 -2
  65. mage_ai/server/frontend_dist/settings/workspace/preferences.html +2 -2
  66. mage_ai/server/frontend_dist/settings/workspace/sync-data.html +2 -2
  67. mage_ai/server/frontend_dist/settings/workspace/users.html +2 -2
  68. mage_ai/server/frontend_dist/settings.html +2 -2
  69. mage_ai/server/frontend_dist/sign-in.html +2 -2
  70. mage_ai/server/frontend_dist/terminal.html +2 -2
  71. mage_ai/server/frontend_dist/test.html +2 -2
  72. mage_ai/server/frontend_dist/triggers.html +2 -2
  73. mage_ai/server/server.py +8 -0
  74. mage_ai/server/websocket_server.py +3 -2
  75. mage_ai/services/spark/config.py +8 -2
  76. mage_ai/services/spark/spark.py +64 -22
  77. mage_ai/shared/environments.py +4 -8
  78. mage_ai/tests/api/operations/test_syncs.py +1 -1
  79. mage_ai/tests/data_preparation/models/test_pipeline.py +11 -1
  80. {mage_ai-0.8.91.dist-info → mage_ai-0.8.93.dist-info}/METADATA +1 -1
  81. {mage_ai-0.8.91.dist-info → mage_ai-0.8.93.dist-info}/RECORD +87 -83
  82. mage_ai/server/frontend_dist/_next/static/chunks/1424-c6b0d89ffb4a10b9.js +0 -1
  83. mage_ai/server/frontend_dist/_next/static/chunks/3714-c70e815b08e3d9be.js +0 -1
  84. mage_ai/server/frontend_dist/_next/static/chunks/3752-bd78037feb0a755f.js +0 -1
  85. mage_ai/server/frontend_dist/_next/static/chunks/pages/_app-aa11738683e2250f.js +0 -1
  86. mage_ai/server/frontend_dist/_next/static/chunks/pages/manage-06aa8a8f1ca2e8d8.js +0 -1
  87. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipeline-runs-3260a2dac8df672e.js +0 -1
  88. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/backfills-f08b51d9dc56eab5.js +0 -1
  89. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-43e71712d3fc0299.js +0 -1
  90. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/logs-264439be4f197741.js +0 -1
  91. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runs-91ba61b9030eff1f.js +0 -1
  92. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/monitors/block-runtime-0bbae5456b0e6e82.js +0 -1
  93. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/runs/[run]-86d1477c6671ea30.js +0 -1
  94. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/runs-4b0c098074dd3e6d.js +0 -1
  95. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/syncs-891c3d3f7a2b634b.js +0 -1
  96. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines-f4d470fe28b74de7.js +0 -1
  97. mage_ai/server/frontend_dist/_next/static/j-J6532RA0pcVgjHCeKSz/_buildManifest.js +0 -1
  98. /mage_ai/server/frontend_dist/_next/static/{j-J6532RA0pcVgjHCeKSz → zlCBBK90aKYZtPlYLj9_T}/_middlewareManifest.js +0 -0
  99. /mage_ai/server/frontend_dist/_next/static/{j-J6532RA0pcVgjHCeKSz → zlCBBK90aKYZtPlYLj9_T}/_ssgManifest.js +0 -0
  100. {mage_ai-0.8.91.dist-info → mage_ai-0.8.93.dist-info}/LICENSE +0 -0
  101. {mage_ai-0.8.91.dist-info → mage_ai-0.8.93.dist-info}/WHEEL +0 -0
  102. {mage_ai-0.8.91.dist-info → mage_ai-0.8.93.dist-info}/entry_points.txt +0 -0
  103. {mage_ai-0.8.91.dist-info → mage_ai-0.8.93.dist-info}/top_level.txt +0 -0
@@ -81,6 +81,7 @@ BlockPolicy.allow_write([
81
81
  'bookmark_values',
82
82
  'callback_blocks',
83
83
  'color',
84
+ 'conditional_blocks',
84
85
  'configuration',
85
86
  'content',
86
87
  'destination_table',
@@ -37,6 +37,7 @@ PipelinePolicy.allow_actions([
37
37
 
38
38
  PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
39
39
  'callbacks',
40
+ 'conditionals',
40
41
  'extensions',
41
42
  ], scopes=[
42
43
  OauthScope.CLIENT_PRIVATE,
@@ -46,6 +47,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
46
47
 
47
48
  PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
48
49
  'callbacks',
50
+ 'conditionals',
49
51
  'extensions',
50
52
  ], scopes=[
51
53
  OauthScope.CLIENT_PRIVATE,
@@ -56,6 +58,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
56
58
 
57
59
  PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
58
60
  'callbacks',
61
+ 'conditionals',
59
62
  'extensions',
60
63
  ], scopes=[
61
64
  OauthScope.CLIENT_PRIVATE,
@@ -65,6 +68,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
65
68
 
66
69
  PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
67
70
  'callbacks',
71
+ 'conditionals',
68
72
  'extensions',
69
73
  'schedules',
70
74
  ], scopes=[
@@ -75,6 +79,7 @@ PipelinePolicy.allow_read(PipelinePresenter.default_attributes + [
75
79
 
76
80
  PipelinePolicy.allow_write([
77
81
  'callbacks',
82
+ 'conditionals',
78
83
  'clone_pipeline_uuid',
79
84
  'extensions',
80
85
  'name',
@@ -88,6 +93,7 @@ PipelinePolicy.allow_write([
88
93
  PipelinePolicy.allow_write([
89
94
  'add_upstream_for_block_uuid',
90
95
  'callbacks',
96
+ 'conditionals',
91
97
  'extensions',
92
98
  'schedules',
93
99
  ] + PipelinePresenter.default_attributes, scopes=[
@@ -11,6 +11,7 @@ class BlockPresenter(BasePresenter):
11
11
  'all_upstream_blocks_executed',
12
12
  'callback_blocks',
13
13
  'color',
14
+ 'conditional_blocks',
14
15
  'configuration',
15
16
  'downstream_blocks',
16
17
  'executor_config',
@@ -12,6 +12,7 @@ class PipelinePresenter(BasePresenter):
12
12
  'executor_count',
13
13
  'executor_type',
14
14
  'name',
15
+ 'spark_config',
15
16
  'type',
16
17
  'updated_at',
17
18
  'uuid',
@@ -44,6 +45,7 @@ class PipelinePresenter(BasePresenter):
44
45
  include_block_metadata=include_block_metadata,
45
46
  include_block_tags=True,
46
47
  include_callback_blocks=True,
48
+ include_conditional_blocks=True,
47
49
  include_content=include_content,
48
50
  include_extensions=include_extensions,
49
51
  include_outputs=include_outputs,
@@ -9,6 +9,10 @@ def callback(function):
9
9
  return function
10
10
 
11
11
 
12
+ def condition(function):
13
+ return function
14
+
15
+
12
16
  def custom(function):
13
17
  return function
14
18
 
@@ -64,6 +64,31 @@ class BlockExecutor:
64
64
  on_start(self.block_uuid)
65
65
  pipeline_run = PipelineRun.query.get(kwargs['pipeline_run_id']) \
66
66
  if 'pipeline_run_id' in kwargs else None
67
+
68
+ conditional_result = self._execute_conditional(
69
+ dynamic_block_index=dynamic_block_index,
70
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
71
+ global_vars=global_vars,
72
+ logging_tags=tags,
73
+ pipeline_run=pipeline_run,
74
+ )
75
+ if not conditional_result:
76
+ self.logger.info(
77
+ f'Conditional block(s) returned false for {self.block.uuid}. '
78
+ 'This block run and downstream blocks will set as CONDITION_FAILED.',
79
+ **merge_dict(tags, dict(
80
+ block_type=self.block.type,
81
+ block_uuid=self.block.uuid,
82
+ )),
83
+ )
84
+ self.__update_block_run_status(
85
+ 'condition_failed',
86
+ block_run_id=kwargs.get('block_run_id'),
87
+ callback_url=callback_url,
88
+ tags=tags,
89
+ )
90
+ return dict(output=[])
91
+
67
92
  try:
68
93
  from mage_ai.shared.retry import retry
69
94
 
@@ -207,12 +232,52 @@ class BlockExecutor:
207
232
 
208
233
  return result
209
234
 
235
+ def _execute_conditional(
236
+ self,
237
+ global_vars: Dict,
238
+ logging_tags: Dict,
239
+ pipeline_run: PipelineRun,
240
+ dynamic_block_index: Union[int, None] = None,
241
+ dynamic_upstream_block_uuids: Union[List[str], None] = None,
242
+ ):
243
+ result = True
244
+ for conditional_block in self.block.conditional_blocks:
245
+ try:
246
+ block_result = conditional_block.execute_conditional(
247
+ self.block,
248
+ dynamic_block_index=dynamic_block_index,
249
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
250
+ execution_partition=self.execution_partition,
251
+ global_vars=global_vars,
252
+ logger=self.logger,
253
+ logging_tags=logging_tags,
254
+ pipeline_run=pipeline_run,
255
+ )
256
+ if not block_result:
257
+ self.logger.info(
258
+ f'Conditional block {conditional_block.uuid} evaluated as False ' +
259
+ f'for block {self.block.uuid}',
260
+ logging_tags,
261
+ )
262
+ result = result and block_result
263
+ except Exception as conditional_err:
264
+ self.logger.exception(
265
+ f'Failed to execute conditional block {conditional_block.uuid} ' +
266
+ f'for block {self.block.uuid}.',
267
+ **merge_dict(logging_tags, dict(
268
+ error=conditional_err,
269
+ )),
270
+ )
271
+ result = False
272
+
273
+ return result
274
+
210
275
  def _execute_callback(
211
276
  self,
212
277
  callback: str,
213
- global_vars,
214
- logging_tags,
215
- pipeline_run,
278
+ global_vars: Dict,
279
+ logging_tags: Dict,
280
+ pipeline_run: PipelineRun,
216
281
  dynamic_block_index: Union[int, None] = None,
217
282
  dynamic_upstream_block_uuids: Union[List[str], None] = None,
218
283
  ):
@@ -47,7 +47,7 @@ from mage_ai.data_preparation.models.constants import (
47
47
  )
48
48
  from mage_ai.data_preparation.models.file import File
49
49
  from mage_ai.data_preparation.models.variable import VariableType
50
- from mage_ai.data_preparation.repo_manager import get_repo_path
50
+ from mage_ai.data_preparation.repo_manager import RepoConfig, get_repo_path
51
51
  from mage_ai.data_preparation.shared.stream import StreamToLogger
52
52
  from mage_ai.data_preparation.templates.template import load_template
53
53
  from mage_ai.server.kernel_output_parser import DataType
@@ -253,6 +253,7 @@ class Block:
253
253
 
254
254
  self._outputs = None
255
255
  self._outputs_loaded = False
256
+ self.conditional_blocks = []
256
257
  self.callback_blocks = []
257
258
  self.upstream_blocks = []
258
259
  self.downstream_blocks = []
@@ -348,6 +349,10 @@ class Block:
348
349
  def callback_block_uuids(self) -> List[str]:
349
350
  return [b.uuid for b in self.callback_blocks]
350
351
 
352
+ @property
353
+ def conditional_block_uuids(self) -> List[str]:
354
+ return [b.uuid for b in self.conditional_blocks]
355
+
351
356
  @property
352
357
  def upstream_block_uuids(self) -> List[str]:
353
358
  return [b.uuid for b in self.upstream_blocks]
@@ -456,6 +461,8 @@ class Block:
456
461
  elif pipeline and PipelineType.INTEGRATION == pipeline.type:
457
462
  if BlockType.CALLBACK == block_type:
458
463
  return CallbackBlock
464
+ elif BlockType.CONDITIONAL == block_type:
465
+ return ConditionalBlock
459
466
  elif BlockType.DATA_LOADER == block_type:
460
467
  return SourceBlock
461
468
  elif BlockType.DATA_EXPORTER == block_type:
@@ -670,11 +677,33 @@ class Block:
670
677
  if logging_tags is None:
671
678
  logging_tags = dict()
672
679
 
673
- arr = []
680
+ if self.conditional_blocks and len(self.conditional_blocks) > 0:
681
+ conditional_message = ''
682
+ result = True
683
+ for conditional_block in self.conditional_blocks:
684
+ block_result = conditional_block.execute_conditional(
685
+ self,
686
+ global_vars=global_vars,
687
+ logger=logger,
688
+ logging_tags=logging_tags,
689
+ )
690
+ conditional_message += \
691
+ f'Conditional block {conditional_block.uuid} evaluated to {block_result}.\n'
692
+ result = result and block_result
693
+
694
+ # Print result to block output
695
+ if not result:
696
+ conditional_message += 'This block would not be executed in a trigger run.\n'
697
+ conditional_json = json.dumps(dict(
698
+ message=conditional_message,
699
+ ))
700
+ print(f'[__internal_test__]{conditional_json}')
701
+
702
+ callback_arr = []
674
703
  if self.callback_block:
675
- arr.append(self.callback_block)
704
+ callback_arr.append(self.callback_block)
676
705
  if self.callback_blocks:
677
- arr += self.callback_blocks
706
+ callback_arr += self.callback_blocks
678
707
 
679
708
  try:
680
709
  output = self.execute_sync(
@@ -684,7 +713,7 @@ class Block:
684
713
  **kwargs
685
714
  )
686
715
  except Exception as e:
687
- for callback_block in arr:
716
+ for callback_block in callback_arr:
688
717
  callback_block.execute_callback(
689
718
  'on_failure',
690
719
  global_vars=global_vars,
@@ -694,7 +723,7 @@ class Block:
694
723
  )
695
724
  raise e
696
725
 
697
- for callback_block in arr:
726
+ for callback_block in callback_arr:
698
727
  callback_block.execute_callback(
699
728
  'on_success',
700
729
  global_vars=global_vars,
@@ -1131,6 +1160,8 @@ class Block:
1131
1160
  self,
1132
1161
  execution_partition: str = None,
1133
1162
  include_print_outputs: bool = True,
1163
+ csv_lines_only: bool = False,
1164
+ sample: bool = True,
1134
1165
  sample_count: int = DATAFRAME_SAMPLE_COUNT_PREVIEW,
1135
1166
  variable_type: VariableType = None,
1136
1167
  block_uuid: str = None,
@@ -1166,40 +1197,47 @@ class Block:
1166
1197
  continue
1167
1198
 
1168
1199
  data = variable_object.read_data(
1169
- sample=True,
1200
+ sample=sample,
1170
1201
  sample_count=sample_count,
1171
1202
  spark=self.__get_spark_session(),
1172
1203
  )
1173
1204
  if type(data) is pd.DataFrame:
1174
- try:
1175
- analysis = variable_manager.get_variable(
1176
- self.pipeline.uuid,
1177
- block_uuid,
1178
- v,
1179
- dataframe_analysis_keys=['metadata', 'statistics'],
1180
- partition=execution_partition,
1181
- variable_type=VariableType.DATAFRAME_ANALYSIS,
1205
+ if csv_lines_only:
1206
+ data = dict(
1207
+ table=data.to_csv(header=True, index=False).strip('\n').split('\n')
1182
1208
  )
1183
- except Exception:
1184
- analysis = None
1185
- if analysis is not None:
1186
- stats = analysis.get('statistics', {})
1187
- column_types = (analysis.get('metadata') or {}).get('column_types', {})
1188
- row_count = stats.get('original_row_count', stats.get('count'))
1189
- column_count = stats.get('original_column_count', len(column_types))
1190
1209
  else:
1191
- row_count, column_count = data.shape
1192
-
1193
- columns_to_display = data.columns.tolist()[:DATAFRAME_ANALYSIS_MAX_COLUMNS]
1194
- data = dict(
1195
- sample_data=dict(
1196
- columns=columns_to_display,
1197
- rows=json.loads(data[columns_to_display].to_json(orient='split'))['data']
1198
- ),
1199
- shape=[row_count, column_count],
1200
- type=DataType.TABLE,
1201
- variable_uuid=v,
1202
- )
1210
+ try:
1211
+ analysis = variable_manager.get_variable(
1212
+ self.pipeline.uuid,
1213
+ block_uuid,
1214
+ v,
1215
+ dataframe_analysis_keys=['metadata', 'statistics'],
1216
+ partition=execution_partition,
1217
+ variable_type=VariableType.DATAFRAME_ANALYSIS,
1218
+ )
1219
+ except Exception:
1220
+ analysis = None
1221
+ if analysis is not None:
1222
+ stats = analysis.get('statistics', {})
1223
+ column_types = (analysis.get('metadata') or {}).get('column_types', {})
1224
+ row_count = stats.get('original_row_count', stats.get('count'))
1225
+ column_count = stats.get('original_column_count', len(column_types))
1226
+ else:
1227
+ row_count, column_count = data.shape
1228
+
1229
+ columns_to_display = data.columns.tolist()[:DATAFRAME_ANALYSIS_MAX_COLUMNS]
1230
+ data = dict(
1231
+ sample_data=dict(
1232
+ columns=columns_to_display,
1233
+ rows=json.loads(
1234
+ data[columns_to_display].to_json(orient='split')
1235
+ )['data']
1236
+ ),
1237
+ shape=[row_count, column_count],
1238
+ type=DataType.TABLE,
1239
+ variable_uuid=v,
1240
+ )
1203
1241
  data_products.append(data)
1204
1242
  continue
1205
1243
  elif is_geo_dataframe(data):
@@ -1383,7 +1421,11 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
1383
1421
  variable_mapping = self.__save_outputs_prepare(outputs)
1384
1422
  await self.store_variables_async(variable_mapping, override=override)
1385
1423
 
1386
- def to_dict_base(self, include_callback_blocks: bool = False) -> Dict:
1424
+ def to_dict_base(
1425
+ self,
1426
+ include_callback_blocks: bool = False,
1427
+ include_conditional_blocks: bool = False,
1428
+ ) -> Dict:
1387
1429
  language = self.language
1388
1430
  if language and type(self.language) is not str:
1389
1431
  language = self.language.value
@@ -1410,6 +1452,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
1410
1452
  if include_callback_blocks:
1411
1453
  data['callback_blocks'] = self.callback_block_uuids
1412
1454
 
1455
+ if include_conditional_blocks:
1456
+ data['conditional_blocks'] = self.conditional_block_uuids
1457
+
1413
1458
  if self.replicated_block:
1414
1459
  data['replicated_block'] = self.replicated_block
1415
1460
 
@@ -1451,12 +1496,16 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
1451
1496
  include_block_metadata: bool = False,
1452
1497
  include_block_tags: bool = False,
1453
1498
  include_callback_blocks: bool = False,
1499
+ include_conditional_blocks: bool = False,
1454
1500
  include_content: bool = False,
1455
1501
  include_outputs: bool = False,
1456
1502
  sample_count: int = None,
1457
1503
  check_if_file_exists: bool = False,
1458
1504
  ) -> Dict:
1459
- data = self.to_dict_base(include_callback_blocks=include_callback_blocks)
1505
+ data = self.to_dict_base(
1506
+ include_callback_blocks=include_callback_blocks,
1507
+ include_conditional_blocks=include_conditional_blocks,
1508
+ )
1460
1509
 
1461
1510
  if include_content:
1462
1511
  data['content'] = await self.content_async()
@@ -1509,6 +1558,11 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
1509
1558
  ):
1510
1559
  self.__update_callback_blocks(data['callback_blocks'])
1511
1560
 
1561
+ if 'conditional_blocks' in data and set(data['conditional_blocks']) != set(
1562
+ self.conditional_block_uuids
1563
+ ):
1564
+ self.__update_conditional_blocks(data['conditional_blocks'])
1565
+
1512
1566
  if 'executor_type' in data and data['executor_type'] != self.executor_type:
1513
1567
  self.executor_type = data['executor_type']
1514
1568
  self.__update_pipeline_block()
@@ -1524,6 +1578,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
1524
1578
  def update_callback_blocks(self, callback_blocks: List[Any]) -> None:
1525
1579
  self.callback_blocks = callback_blocks
1526
1580
 
1581
+ def update_conditional_blocks(self, conditional_blocks: List[Any]) -> None:
1582
+ self.conditional_blocks = conditional_blocks
1583
+
1527
1584
  def update_upstream_blocks(self, upstream_blocks: List[Any]) -> None:
1528
1585
  self.upstream_blocks = upstream_blocks
1529
1586
 
@@ -1839,11 +1896,17 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
1839
1896
  return global_vars
1840
1897
 
1841
1898
  def __get_spark_session(self):
1842
- if self.spark_init:
1899
+ if self.spark_init and (not self.pipeline or
1900
+ not self.pipeline.spark_config):
1843
1901
  return self.spark
1844
1902
  try:
1845
- spark_config = SparkConfig.load(
1846
- config={'repo_path': self.repo_path})
1903
+ if self.pipeline and self.pipeline.spark_config:
1904
+ spark_config = SparkConfig.load(
1905
+ config=self.pipeline.spark_config)
1906
+ else:
1907
+ repo_config = RepoConfig(repo_path=self.repo_path)
1908
+ spark_config = SparkConfig.load(
1909
+ config=repo_config.spark_config)
1847
1910
  self.spark = get_spark_session(spark_config)
1848
1911
  except Exception:
1849
1912
  self.spark = None
@@ -2026,6 +2089,7 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
2026
2089
 
2027
2090
  def tags(self) -> List[str]:
2028
2091
  from mage_ai.data_preparation.models.block.constants import (
2092
+ TAG_CONDITION,
2029
2093
  TAG_DYNAMIC,
2030
2094
  TAG_DYNAMIC_CHILD,
2031
2095
  TAG_REDUCE_OUTPUT,
@@ -2046,6 +2110,9 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
2046
2110
  if self.replicated_block:
2047
2111
  arr.append(TAG_REPLICA)
2048
2112
 
2113
+ if len(self.conditional_blocks) > 0:
2114
+ arr.append(TAG_CONDITION)
2115
+
2049
2116
  return arr
2050
2117
 
2051
2118
  def variable_object(
@@ -2155,6 +2222,16 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
2155
2222
  widget=BlockType.CHART == self.type,
2156
2223
  )
2157
2224
 
2225
+ def __update_conditional_blocks(self, block_uuids: List[str]) -> None:
2226
+ if self.pipeline is None:
2227
+ return
2228
+
2229
+ self.pipeline.update_block(
2230
+ self,
2231
+ conditional_block_uuids=block_uuids,
2232
+ widget=BlockType.CHART == self.type,
2233
+ )
2234
+
2158
2235
 
2159
2236
  class SensorBlock(Block):
2160
2237
  def execute_block_function(
@@ -2187,7 +2264,98 @@ class SensorBlock(Block):
2187
2264
  return []
2188
2265
 
2189
2266
 
2190
- class CallbackBlock(Block):
2267
+ class AddonBlock(Block):
2268
+ def _create_global_vars(
2269
+ self,
2270
+ global_vars: Dict,
2271
+ parent_block: Block,
2272
+ **kwargs,
2273
+ ) -> Dict:
2274
+ pipeline_run = kwargs.get('pipeline_run')
2275
+ global_vars = merge_dict(
2276
+ global_vars or dict(),
2277
+ dict(
2278
+ pipeline_uuid=self.pipeline.uuid,
2279
+ block_uuid=self.uuid,
2280
+ pipeline_run=pipeline_run,
2281
+ ),
2282
+ )
2283
+ if parent_block:
2284
+ global_vars['parent_block_uuid'] = parent_block.uuid
2285
+
2286
+ if parent_block and \
2287
+ parent_block.pipeline and \
2288
+ PipelineType.INTEGRATION == parent_block.pipeline.type:
2289
+
2290
+ template_runtime_configuration = parent_block.template_runtime_configuration
2291
+ index = template_runtime_configuration.get('index', None)
2292
+ is_last_block_run = template_runtime_configuration.get('is_last_block_run', False)
2293
+ selected_streams = template_runtime_configuration.get('selected_streams', [])
2294
+ stream = selected_streams[0] if len(selected_streams) >= 1 else None
2295
+ destination_table = template_runtime_configuration.get('destination_table', stream)
2296
+
2297
+ global_vars['index'] = index
2298
+ global_vars['is_last_block_run'] = is_last_block_run
2299
+ global_vars['stream'] = stream
2300
+ global_vars['destination_table'] = destination_table
2301
+
2302
+ return global_vars
2303
+
2304
+
2305
+ class ConditionalBlock(AddonBlock):
2306
+ def execute_conditional(
2307
+ self,
2308
+ parent_block: Block,
2309
+ dynamic_block_index: Union[int, None] = None,
2310
+ dynamic_upstream_block_uuids: Union[List[str], None] = None,
2311
+ execution_partition: str = None,
2312
+ global_vars: Dict = None,
2313
+ logger: Logger = None,
2314
+ logging_tags: Dict = None,
2315
+ **kwargs
2316
+ ) -> bool:
2317
+ if logger is not None:
2318
+ stdout = StreamToLogger(logger, logging_tags=logging_tags)
2319
+ else:
2320
+ stdout = sys.stdout
2321
+
2322
+ with redirect_stdout(stdout):
2323
+ global_vars = self._create_global_vars(
2324
+ global_vars,
2325
+ parent_block,
2326
+ **kwargs,
2327
+ )
2328
+
2329
+ condition_functions = []
2330
+
2331
+ results = dict(
2332
+ condition=self._block_decorator(condition_functions),
2333
+ )
2334
+ exec(self.content, results)
2335
+
2336
+ global_vars_copy = global_vars.copy()
2337
+ input_vars = []
2338
+ # Fetch input variables for parent block
2339
+ if parent_block is not None:
2340
+ input_vars, kwargs_vars, _ = parent_block.fetch_input_variables(
2341
+ None,
2342
+ execution_partition,
2343
+ global_vars,
2344
+ dynamic_block_index=dynamic_block_index,
2345
+ dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
2346
+ )
2347
+
2348
+ for kwargs_var in kwargs_vars:
2349
+ global_vars_copy.update(kwargs_var)
2350
+
2351
+ result = True
2352
+ for condition_function in condition_functions:
2353
+ result = condition_function(*input_vars, **global_vars_copy) and result
2354
+
2355
+ return result
2356
+
2357
+
2358
+ class CallbackBlock(AddonBlock):
2191
2359
  @classmethod
2192
2360
  def create(cls, orig_block_name) -> 'CallbackBlock':
2193
2361
  return Block.create(
@@ -2209,40 +2377,17 @@ class CallbackBlock(Block):
2209
2377
  parent_block: Block = None,
2210
2378
  **kwargs
2211
2379
  ) -> None:
2212
- pipeline_run = kwargs.get('pipeline_run')
2213
-
2214
2380
  if logger is not None:
2215
2381
  stdout = StreamToLogger(logger, logging_tags=logging_tags)
2216
2382
  else:
2217
2383
  stdout = sys.stdout
2218
2384
 
2219
2385
  with redirect_stdout(stdout):
2220
- global_vars = merge_dict(
2221
- global_vars or dict(),
2222
- dict(
2223
- pipeline_uuid=self.pipeline.uuid,
2224
- block_uuid=self.uuid,
2225
- pipeline_run=pipeline_run,
2226
- ),
2386
+ global_vars = self._create_global_vars(
2387
+ global_vars,
2388
+ parent_block,
2389
+ **kwargs
2227
2390
  )
2228
- if parent_block:
2229
- global_vars['parent_block_uuid'] = parent_block.uuid
2230
-
2231
- if parent_block and \
2232
- parent_block.pipeline and \
2233
- PipelineType.INTEGRATION == parent_block.pipeline.type:
2234
-
2235
- template_runtime_configuration = parent_block.template_runtime_configuration
2236
- index = template_runtime_configuration.get('index', None)
2237
- is_last_block_run = template_runtime_configuration.get('is_last_block_run', False)
2238
- selected_streams = template_runtime_configuration.get('selected_streams', [])
2239
- stream = selected_streams[0] if len(selected_streams) >= 1 else None
2240
- destination_table = template_runtime_configuration.get('destination_table', stream)
2241
-
2242
- global_vars['index'] = index
2243
- global_vars['is_last_block_run'] = is_last_block_run
2244
- global_vars['stream'] = stream
2245
- global_vars['destination_table'] = destination_table
2246
2391
 
2247
2392
  callback_functions = []
2248
2393
  failure_functions = []
@@ -1,15 +1,16 @@
1
1
  from mage_ai.data_preparation.models.block import (
2
2
  Block,
3
3
  CallbackBlock,
4
+ ConditionalBlock,
4
5
  SensorBlock,
5
6
  )
6
7
  from mage_ai.data_preparation.models.block.dbt import DBTBlock
7
8
  from mage_ai.data_preparation.models.block.extension.block import ExtensionBlock
8
9
  from mage_ai.data_preparation.models.constants import BlockType
9
10
 
10
-
11
11
  BLOCK_TYPE_TO_CLASS = {
12
12
  BlockType.CALLBACK: CallbackBlock,
13
+ BlockType.CONDITIONAL: ConditionalBlock,
13
14
  BlockType.CUSTOM: Block,
14
15
  BlockType.DATA_EXPORTER: Block,
15
16
  BlockType.DATA_LOADER: Block,
@@ -21,6 +22,7 @@ BLOCK_TYPE_TO_CLASS = {
21
22
  BlockType.SENSOR: SensorBlock,
22
23
  }
23
24
 
25
+ TAG_CONDITION = 'condition'
24
26
  TAG_DYNAMIC = 'dynamic'
25
27
  TAG_DYNAMIC_CHILD = 'dynamic_child'
26
28
  TAG_REDUCE_OUTPUT = 'reduce_output'
@@ -1,12 +1,7 @@
1
1
  import os
2
2
  from enum import Enum
3
3
 
4
- PIPELINES_FOLDER = 'pipelines'
5
- PIPELINE_CONFIG_FILE = 'metadata.yaml'
6
- PREFERENCES_FILE = '.preferences.yaml'
7
4
  DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json'
8
-
9
-
10
5
  DATAFRAME_ANALYSIS_KEYS = frozenset(
11
6
  [
12
7
  'metadata',
@@ -15,14 +10,18 @@ DATAFRAME_ANALYSIS_KEYS = frozenset(
15
10
  'suggestions',
16
11
  ]
17
12
  )
18
- DATAFRAME_ANALYSIS_MAX_ROWS = 100_000
19
13
  DATAFRAME_ANALYSIS_MAX_COLUMNS = 100
20
- DATAFRAME_SAMPLE_COUNT_PREVIEW = 10
14
+ DATAFRAME_ANALYSIS_MAX_ROWS = 100_000
21
15
  DATAFRAME_SAMPLE_COUNT = 1000
16
+ DATAFRAME_SAMPLE_COUNT_PREVIEW = 10
22
17
  DATAFRAME_SAMPLE_MAX_COLUMNS = 1000
18
+ LOGS_DIR = '.logs'
23
19
  MAX_PRINT_OUTPUT_LINES = int(os.getenv('MAX_PRINT_OUTPUT_LINES', 1000) or 1000)
20
+ PIPELINE_CONFIG_FILE = 'metadata.yaml'
21
+ PIPELINES_FOLDER = 'pipelines'
22
+ PREFERENCES_FILE = '.preferences.yaml'
23
+ REPO_CONFIG_FILE = 'metadata.yaml'
24
24
  VARIABLE_DIR = '.variables'
25
- LOGS_DIR = '.logs'
26
25
 
27
26
 
28
27
  class BlockLanguage(str, Enum):
@@ -42,6 +41,7 @@ class BlockStatus(str, Enum):
42
41
 
43
42
  class BlockType(str, Enum):
44
43
  CALLBACK = 'callback'
44
+ CONDITIONAL = 'conditional'
45
45
  CHART = 'chart'
46
46
  CUSTOM = 'custom'
47
47
  DATA_EXPORTER = 'data_exporter'