mage-ai 0.8.37__py3-none-any.whl → 0.8.39__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 (91) hide show
  1. mage_ai/api/logging.py +2 -2
  2. mage_ai/api/resources/PipelineRunResource.py +6 -3
  3. mage_ai/data_integrations/sources/constants.py +1 -0
  4. mage_ai/data_preparation/models/block/dbt/__init__.py +20 -0
  5. mage_ai/data_preparation/models/block/dbt/utils/__init__.py +4 -1
  6. mage_ai/data_preparation/models/pipeline.py +10 -1
  7. mage_ai/io/base.py +2 -2
  8. mage_ai/io/constants.py +3 -0
  9. mage_ai/io/export_utils.py +14 -1
  10. mage_ai/io/mssql.py +5 -2
  11. mage_ai/io/mysql.py +7 -3
  12. mage_ai/io/postgres.py +64 -21
  13. mage_ai/io/sql.py +35 -6
  14. mage_ai/io/trino.py +7 -3
  15. mage_ai/orchestration/pipeline_scheduler.py +37 -27
  16. mage_ai/server/active_kernel.py +6 -3
  17. mage_ai/server/api/clusters.py +4 -1
  18. mage_ai/server/api/integration_sources.py +5 -2
  19. mage_ai/server/client/mage.py +2 -2
  20. mage_ai/server/constants.py +1 -1
  21. mage_ai/server/data/base.py +2 -2
  22. mage_ai/server/data/models.py +2 -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/{2626-905774aafeb2c600.js → 2626-e7fa4f83f8214c97.js} +1 -1
  26. mage_ai/server/frontend_dist/_next/static/chunks/4178-9103014b7dae3c49.js +1 -0
  27. mage_ai/server/frontend_dist/_next/static/chunks/{4538-8a3c3e47be976ede.js → 4538-347283088b83c6bf.js} +1 -1
  28. mage_ai/server/frontend_dist/_next/static/chunks/5141-ddf4ba0a362d6f34.js +1 -0
  29. mage_ai/server/frontend_dist/_next/static/chunks/{5477-793cd2120261d023.js → 5477-b439f211b6146a11.js} +1 -1
  30. mage_ai/server/frontend_dist/_next/static/chunks/{5872-103815a4a043489b.js → 5872-1767c45ee6690ae5.js} +1 -1
  31. mage_ai/server/frontend_dist/_next/static/chunks/{5896-f84e336fb8877027.js → 5896-10a676bcc86978cc.js} +1 -1
  32. mage_ai/server/frontend_dist/_next/static/chunks/{7400-365cb7888b6db7d9.js → 7400-f4db9b5d41f67f75.js} +1 -1
  33. mage_ai/server/frontend_dist/_next/static/chunks/{9386-fb899ca8ecc2a350.js → 9386-d4cc11bab74eec8d.js} +1 -1
  34. mage_ai/server/frontend_dist/_next/static/chunks/{9832-c8b8970bb522f302.js → 9832-f97919376d52e3bf.js} +1 -1
  35. mage_ai/server/frontend_dist/_next/static/chunks/pages/{manage-9e5f315db570ac77.js → manage-3046bc53d24917c7.js} +1 -1
  36. mage_ai/server/frontend_dist/_next/static/chunks/pages/{pipeline-runs-79e10a783afec3df.js → pipeline-runs-e64ba4e8b2bfe73c.js} +1 -1
  37. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-017f6a2de38f8658.js +1 -0
  38. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/{syncs-1767a2f57f887ef7.js → syncs-e1271453ed0c8d6e.js} +1 -1
  39. mage_ai/server/frontend_dist/_next/static/chunks/pages/{pipelines-52c3ee3817e5554b.js → pipelines-7446a70bdd8381a5.js} +1 -1
  40. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/{preferences-a23b61bab04a16f3.js → preferences-997acba85f777259.js} +1 -1
  41. mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/{sync-data-de28e502102defde.js → sync-data-8c903140c99e487c.js} +1 -1
  42. mage_ai/server/frontend_dist/_next/static/chunks/pages/terminal-6c0ea500b3bc6b61.js +1 -0
  43. mage_ai/server/frontend_dist/_next/static/chunks/pages/{triggers-8a2169d30b643ae7.js → triggers-783b9526167f1249.js} +1 -1
  44. mage_ai/server/frontend_dist/_next/static/{23nXkA2GRjJCZDwY3kuKd → ngCg1gM6urg83y2nV5KHK}/_buildManifest.js +1 -1
  45. mage_ai/server/frontend_dist/index.html +2 -2
  46. mage_ai/server/frontend_dist/manage.html +2 -2
  47. mage_ai/server/frontend_dist/pipeline-runs.html +2 -2
  48. mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills/[...slug].html +2 -2
  49. mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills.html +2 -2
  50. mage_ai/server/frontend_dist/pipelines/[pipeline]/edit.html +2 -2
  51. mage_ai/server/frontend_dist/pipelines/[pipeline]/logs.html +2 -2
  52. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runs.html +2 -2
  53. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runtime.html +2 -2
  54. mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors.html +2 -2
  55. mage_ai/server/frontend_dist/pipelines/[pipeline]/runs/[run].html +2 -2
  56. mage_ai/server/frontend_dist/pipelines/[pipeline]/runs.html +2 -2
  57. mage_ai/server/frontend_dist/pipelines/[pipeline]/syncs.html +2 -2
  58. mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers/[...slug].html +2 -2
  59. mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers.html +2 -2
  60. mage_ai/server/frontend_dist/pipelines/[pipeline].html +2 -2
  61. mage_ai/server/frontend_dist/pipelines.html +2 -2
  62. mage_ai/server/frontend_dist/settings/account/profile.html +2 -2
  63. mage_ai/server/frontend_dist/settings/workspace/preferences.html +2 -2
  64. mage_ai/server/frontend_dist/settings/workspace/sync-data.html +2 -2
  65. mage_ai/server/frontend_dist/settings/workspace/users.html +2 -2
  66. mage_ai/server/frontend_dist/settings.html +2 -2
  67. mage_ai/server/frontend_dist/sign-in.html +2 -2
  68. mage_ai/server/frontend_dist/terminal.html +2 -2
  69. mage_ai/server/frontend_dist/test.html +2 -2
  70. mage_ai/server/frontend_dist/triggers.html +2 -2
  71. mage_ai/server/logger.py +16 -0
  72. mage_ai/server/scheduler_manager.py +8 -4
  73. mage_ai/server/server.py +7 -14
  74. mage_ai/server/subscriber.py +6 -3
  75. mage_ai/server/terminal_server.py +65 -0
  76. mage_ai/server/utils/frontend_renderer.py +3 -2
  77. mage_ai/server/websocket_server.py +4 -1
  78. mage_ai/settings/__init__.py +2 -0
  79. mage_ai/shared/logger.py +4 -0
  80. {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/METADATA +1 -1
  81. {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/RECORD +87 -85
  82. mage_ai/server/frontend_dist/_next/static/chunks/4178-e17f37d21253b832.js +0 -1
  83. mage_ai/server/frontend_dist/_next/static/chunks/5141-2ae9eae00ec2cdfa.js +0 -1
  84. mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-6a52671f1853e1a5.js +0 -1
  85. mage_ai/server/frontend_dist/_next/static/chunks/pages/terminal-9c21edae8f1b6737.js +0 -1
  86. /mage_ai/server/frontend_dist/_next/static/{23nXkA2GRjJCZDwY3kuKd → ngCg1gM6urg83y2nV5KHK}/_middlewareManifest.js +0 -0
  87. /mage_ai/server/frontend_dist/_next/static/{23nXkA2GRjJCZDwY3kuKd → ngCg1gM6urg83y2nV5KHK}/_ssgManifest.js +0 -0
  88. {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/LICENSE +0 -0
  89. {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/WHEEL +0 -0
  90. {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/entry_points.txt +0 -0
  91. {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/top_level.txt +0 -0
mage_ai/api/logging.py CHANGED
@@ -1,7 +1,7 @@
1
1
  from datetime import datetime
2
- import logging
2
+ from mage_ai.server.logger import Logger
3
3
 
4
- LOGGER = logging.getLogger(__name__)
4
+ LOGGER = Logger().new_server_logger(__name__)
5
5
 
6
6
 
7
7
  def debug(text):
@@ -6,7 +6,7 @@ from mage_ai.data_preparation.models.constants import PipelineType
6
6
  from mage_ai.data_preparation.models.pipeline import Pipeline
7
7
  from mage_ai.orchestration.db import safe_db_query
8
8
  from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun
9
- from mage_ai.orchestration.pipeline_scheduler import get_variables
9
+ from mage_ai.orchestration.pipeline_scheduler import get_variables, stop_pipeline_run
10
10
  from sqlalchemy.orm import selectinload
11
11
 
12
12
 
@@ -219,8 +219,11 @@ class PipelineRunResource(DatabaseResource):
219
219
 
220
220
  return super().update(dict(status=PipelineRun.PipelineRunStatus.RUNNING))
221
221
  elif PipelineRun.PipelineRunStatus.CANCELLED == payload.get('status'):
222
- from mage_ai.orchestration.pipeline_scheduler import PipelineScheduler
222
+ pipeline = Pipeline.get(
223
+ self.model.pipeline_uuid,
224
+ check_if_exists=True,
225
+ )
223
226
 
224
- PipelineScheduler(self.model).stop()
227
+ stop_pipeline_run(self.model, pipeline)
225
228
 
226
229
  return self
@@ -23,6 +23,7 @@ SOURCES = sorted([
23
23
  dict(name='Facebook Ads'),
24
24
  dict(name='Freshdesk'),
25
25
  dict(name='Front'),
26
+ dict(name='Google Ads'),
26
27
  dict(name='Google Analytics'),
27
28
  dict(name='Google Search Console'),
28
29
  dict(name='Google Sheets'),
@@ -4,6 +4,7 @@ from mage_ai.data_preparation.models.block.dbt.utils import (
4
4
  create_upstream_tables,
5
5
  fetch_model_data,
6
6
  load_profiles_async,
7
+ load_profiles_file,
7
8
  parse_attributes,
8
9
  query_from_compiled_sql,
9
10
  run_dbt_tests,
@@ -15,7 +16,9 @@ from mage_ai.shared.hash import merge_dict
15
16
  from typing import Any, Dict, List
16
17
  import json
17
18
  import os
19
+ import shutil
18
20
  import subprocess
21
+ import yaml
19
22
 
20
23
 
21
24
  class DBTBlock(Block):
@@ -120,8 +123,20 @@ class DBTBlock(Block):
120
123
  test_execution=test_execution,
121
124
  )
122
125
  project_full_path = command_line_dict['project_full_path']
126
+ profiles_dir = command_line_dict['profiles_dir']
123
127
  dbt_profile_target = command_line_dict['profile_target']
124
128
 
129
+ # Create a temporary profiles file with variables and secrets interpolated
130
+ attributes_dict = parse_attributes(self)
131
+ profiles_full_path = attributes_dict['profiles_full_path']
132
+ profile = load_profiles_file(profiles_full_path)
133
+
134
+ temp_profile_full_path = f'{profiles_dir}/profiles.yml'
135
+ os.makedirs(os.path.dirname(temp_profile_full_path), exist_ok=True)
136
+
137
+ with open(temp_profile_full_path, 'w') as f:
138
+ yaml.safe_dump(profile, f)
139
+
125
140
  is_sql = BlockLanguage.SQL == self.language
126
141
  if is_sql:
127
142
  create_upstream_tables(
@@ -214,4 +229,9 @@ class DBTBlock(Block):
214
229
  )
215
230
  outputs = [df]
216
231
 
232
+ try:
233
+ shutil.rmtree(profiles_dir)
234
+ except Exception as err:
235
+ print(f'Error removing temporary profile at {temp_profile_full_path}: {err}')
236
+
217
237
  return outputs
@@ -912,11 +912,13 @@ def build_command_line_arguments(
912
912
  project_full_path = f'{get_repo_path()}/dbt/{project_name}'
913
913
  args += block.content.split(' ')
914
914
 
915
+ profiles_dir = f'{project_full_path}/.mage_temp_profiles'
916
+
915
917
  args += [
916
918
  '--project-dir',
917
919
  project_full_path,
918
920
  '--profiles-dir',
919
- project_full_path,
921
+ profiles_dir,
920
922
  ]
921
923
 
922
924
  dbt_profile_target = block.configuration.get('dbt_profile_target') \
@@ -934,6 +936,7 @@ def build_command_line_arguments(
934
936
 
935
937
  return dbt_command, args, dict(
936
938
  profile_target=dbt_profile_target,
939
+ profiles_dir=profiles_dir,
937
940
  project_full_path=project_full_path,
938
941
  )
939
942
 
@@ -180,10 +180,19 @@ class Pipeline:
180
180
  )
181
181
 
182
182
  @classmethod
183
- def get(self, uuid, repo_path: str = None):
183
+ def get(self, uuid, repo_path: str = None, check_if_exists: bool = False):
184
184
  from mage_ai.data_preparation.models.pipelines.integration_pipeline \
185
185
  import IntegrationPipeline
186
186
 
187
+ if check_if_exists and not os.path.exists(
188
+ os.path.join(
189
+ repo_path or get_repo_path(),
190
+ PIPELINES_FOLDER,
191
+ uuid,
192
+ ),
193
+ ):
194
+ return None
195
+
187
196
  pipeline = self(uuid, repo_path=repo_path)
188
197
  if PipelineType.INTEGRATION == pipeline.type:
189
198
  pipeline = IntegrationPipeline(uuid, repo_path=repo_path)
mage_ai/io/base.py CHANGED
@@ -275,9 +275,9 @@ class BaseSQLDatabase(BaseIO):
275
275
  """
276
276
  return query_string.strip(' \n\t')
277
277
 
278
- def _clean_column_name(self, column_name: str) -> str:
278
+ def _clean_column_name(self, column_name: str, allow_reserved_words: bool = False) -> str:
279
279
  col_new = re.sub(r'\W', '_', column_name.lower())
280
- if col_new.upper() in SQL_RESERVED_WORDS:
280
+ if not allow_reserved_words and col_new.upper() in SQL_RESERVED_WORDS:
281
281
  col_new = f'_{col_new}'
282
282
  return col_new
283
283
 
mage_ai/io/constants.py CHANGED
@@ -841,3 +841,6 @@ SQL_RESERVED_WORDS_SUBSET = {
841
841
  'USER',
842
842
  'WHERE',
843
843
  }
844
+
845
+ UNIQUE_CONFLICT_METHOD_IGNORE = 'IGNORE'
846
+ UNIQUE_CONFLICT_METHOD_UPDATE = 'UPDATE'
@@ -2,7 +2,7 @@ from enum import Enum
2
2
  from mage_ai.shared.utils import clean_name
3
3
  from pandas import DataFrame, Series
4
4
  from pandas.api.types import infer_dtype
5
- from typing import Callable, Dict, Mapping
5
+ from typing import Callable, Dict, List, Mapping
6
6
 
7
7
  """
8
8
  Utilities for exporting Python data frames to external databases.
@@ -31,6 +31,7 @@ class PandasTypes(str, Enum):
31
31
  DATETIME64 = 'datetime64'
32
32
  DECIMAL = 'decimal'
33
33
  INTEGER = 'integer'
34
+ INT64 = 'int64'
34
35
  EMPTY = 'empty'
35
36
  FLOATING = 'floating'
36
37
  MIXED = 'mixed'
@@ -101,6 +102,7 @@ def gen_table_creation_query(
101
102
  dtypes: Mapping[str, str],
102
103
  schema_name: str,
103
104
  table_name: str,
105
+ unique_constraints: List[str] = [],
104
106
  ) -> str:
105
107
  """
106
108
  Generates a database table creation query from a data frame.
@@ -123,4 +125,15 @@ def gen_table_creation_query(
123
125
  else:
124
126
  full_table_name = table_name
125
127
 
128
+ if unique_constraints:
129
+ unique_constraints_clean = [clean_name(col) for col in unique_constraints]
130
+ unique_constraints_escaped = [f'"{col}"'
131
+ for col in unique_constraints_clean]
132
+ index_name = '_'.join([
133
+ clean_name(full_table_name),
134
+ ] + unique_constraints_clean)
135
+ index_name = f'unique{index_name}'[:64]
136
+ query.append(
137
+ f"CONSTRAINT {index_name} UNIQUE ({', '.join(unique_constraints_escaped)})",
138
+ )
126
139
  return f'CREATE TABLE {full_table_name} (' + ','.join(query) + ');'
mage_ai/io/mssql.py CHANGED
@@ -3,7 +3,7 @@ from mage_ai.io.export_utils import PandasTypes
3
3
  from mage_ai.io.base import QUERY_ROW_LIMIT
4
4
  from mage_ai.io.sql import BaseSQL
5
5
  from pandas import DataFrame, Series
6
- from typing import Any, IO, Union
6
+ from typing import Any, IO, List, Union
7
7
  import json
8
8
  import numpy as np
9
9
  import pyodbc
@@ -87,8 +87,11 @@ class MSSQL(BaseSQL):
87
87
  self,
88
88
  cursor: Any,
89
89
  df: DataFrame,
90
+ db_dtypes: List[str],
91
+ dtypes: List[str],
90
92
  full_table_name: str,
91
- buffer: Union[IO, None] = None
93
+ buffer: Union[IO, None] = None,
94
+ **kwargs,
92
95
  ) -> None:
93
96
  values_placeholder = ', '.join(["?" for i in range(len(df.columns))])
94
97
  values = []
mage_ai/io/mysql.py CHANGED
@@ -5,7 +5,7 @@ from mage_ai.shared.utils import clean_name
5
5
  from mysql.connector import connect
6
6
  from mysql.connector.cursor import MySQLCursor
7
7
  from pandas import DataFrame, Series
8
- from typing import IO, Mapping, Union
8
+ from typing import IO, List, Mapping, Union
9
9
  import numpy as np
10
10
 
11
11
 
@@ -47,7 +47,8 @@ class MySQL(BaseSQL):
47
47
  self,
48
48
  dtypes: Mapping[str, str],
49
49
  schema_name: str,
50
- table_name: str
50
+ table_name: str,
51
+ unique_constraints: List[str] = [],
51
52
  ) -> str:
52
53
  query = []
53
54
  for cname in dtypes:
@@ -73,8 +74,11 @@ class MySQL(BaseSQL):
73
74
  self,
74
75
  cursor: MySQLCursor,
75
76
  df: DataFrame,
77
+ db_dtypes: List[str],
78
+ dtypes: List[str],
76
79
  full_table_name: str,
77
- buffer: Union[IO, None] = None
80
+ buffer: Union[IO, None] = None,
81
+ **kwargs,
78
82
  ) -> None:
79
83
  values_placeholder = ', '.join(["%s" for i in range(len(df.columns))])
80
84
  values = []
mage_ai/io/postgres.py CHANGED
@@ -1,4 +1,5 @@
1
1
  from mage_ai.io.config import BaseConfigLoader, ConfigKey
2
+ from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
2
3
  from mage_ai.io.export_utils import BadConversionError, PandasTypes
3
4
  from mage_ai.io.sql import BaseSQL
4
5
  from mage_ai.shared.parsers import encode_complex
@@ -6,12 +7,21 @@ from mage_ai.shared.utils import is_port_in_use
6
7
  from pandas import DataFrame, Series
7
8
  from psycopg2 import connect, _psycopg
8
9
  from sshtunnel import SSHTunnelForwarder
9
- from typing import Union, IO
10
+ from typing import IO, List, Union
10
11
  import numpy as np
11
12
  import pandas as pd
12
13
  import simplejson
13
14
 
14
15
 
16
+ JSON_SERIALIZABLE_TYPES = frozenset([
17
+ PandasTypes.DATE,
18
+ PandasTypes.DATETIME,
19
+ PandasTypes.DATETIME64,
20
+ PandasTypes.OBJECT,
21
+ PandasTypes.TIME,
22
+ ])
23
+
24
+
15
25
  class Postgres(BaseSQL):
16
26
  """
17
27
  Handles data transfer between a PostgreSQL database and the Mage app.
@@ -159,9 +169,9 @@ class Postgres(BaseSQL):
159
169
  column_type = None
160
170
 
161
171
  if len(values) >= 1:
162
- value = values[0]
163
172
  column_type = 'JSONB'
164
173
 
174
+ value = values[0]
165
175
  if type(value) is list:
166
176
  if len(value) >= 1:
167
177
  item = value[0]
@@ -204,7 +214,7 @@ class Postgres(BaseSQL):
204
214
  return 'bytea'
205
215
  elif dtype in (PandasTypes.FLOATING, PandasTypes.DECIMAL, PandasTypes.MIXED_INTEGER_FLOAT):
206
216
  return 'double precision'
207
- elif dtype == PandasTypes.INTEGER:
217
+ elif dtype == PandasTypes.INTEGER or dtype == PandasTypes.INT64:
208
218
  max_int, min_int = column.max(), column.min()
209
219
  if np.int16(max_int) == max_int and np.int16(min_int) == min_int:
210
220
  return 'smallint'
@@ -229,9 +239,21 @@ class Postgres(BaseSQL):
229
239
  self,
230
240
  cursor: _psycopg.cursor,
231
241
  df: DataFrame,
242
+ db_dtypes: List[str],
243
+ dtypes: List[str],
232
244
  full_table_name: str,
233
- buffer: Union[IO, None] = None
245
+ buffer: Union[IO, None] = None,
246
+ allow_reserved_words: bool = False,
247
+ unique_conflict_method: str = None,
248
+ unique_constraints: List[str] = None,
234
249
  ) -> None:
250
+ def clean_array_value(val):
251
+ if val is None or type(val) is not str or len(val) < 2:
252
+ return val
253
+ if val[0] == '[' and val[-1] == ']':
254
+ return '{' + val[1:-1] + '}'
255
+ return val
256
+
235
257
  df_ = df.copy()
236
258
  columns = df_.columns
237
259
 
@@ -239,28 +261,49 @@ class Postgres(BaseSQL):
239
261
  df_col_dropna = df_[col].dropna()
240
262
  if df_col_dropna.count() == 0:
241
263
  continue
242
- if PandasTypes.OBJECT == df_[col].dtype and type(df_col_dropna.iloc[0]) != str:
264
+ if dtypes[col] in JSON_SERIALIZABLE_TYPES \
265
+ or (df_[col].dtype == PandasTypes.OBJECT and
266
+ type(df_col_dropna.iloc[0]) != str):
243
267
  df_[col] = df_[col].apply(lambda x: simplejson.dumps(
244
268
  x,
245
269
  default=encode_complex,
246
270
  ignore_nan=True,
247
271
  ))
272
+ if '[]' in db_dtypes[col]:
273
+ df_[col] = df_[col].apply(lambda x: clean_array_value(x))
248
274
 
249
- df_.to_csv(
250
- buffer,
251
- header=False,
252
- index=False,
253
- na_rep='',
254
- )
275
+ values = []
255
276
 
256
- buffer.seek(0)
277
+ for _, row in df_.iterrows():
278
+ t = tuple(row)
279
+ if len(t) == 1:
280
+ value = f'({str(t[0])})'
281
+ else:
282
+ value = str(t)
283
+ values.append(value.replace('None', 'NULL'))
284
+ values_string = ', '.join(values)
285
+ insert_columns = ', '.join([f'"{col}"'for col in columns])
286
+
287
+ commands = [
288
+ f'INSERT INTO {full_table_name} ({insert_columns})',
289
+ f'VALUES {values_string}',
290
+ ]
291
+ if unique_constraints and unique_conflict_method:
292
+ unique_constraints = \
293
+ [f'"{self._clean_column_name(col, allow_reserved_words=allow_reserved_words)}"'
294
+ for col in unique_constraints]
295
+ columns_cleaned = \
296
+ [f'"{self._clean_column_name(col, allow_reserved_words=allow_reserved_words)}"'
297
+ for col in columns]
257
298
 
258
- columns_names = ', '.join(columns)
259
- cursor.copy_expert(f"""
260
- COPY {full_table_name} FROM STDIN (
261
- FORMAT csv
262
- , DELIMITER \',\'
263
- , NULL \'\'
264
- , FORCE_NULL({columns_names})
265
- );
266
- """, buffer)
299
+ commands.append(f"ON CONFLICT ({', '.join(unique_constraints)})")
300
+ if UNIQUE_CONFLICT_METHOD_UPDATE == unique_conflict_method:
301
+ update_command = [f'{col} = EXCLUDED.{col}' for col in columns_cleaned]
302
+ commands.append(
303
+ f"DO UPDATE SET {', '.join(update_command)}",
304
+ )
305
+ else:
306
+ commands.append('DO NOTHING')
307
+ cursor.execute(
308
+ '\n'.join(commands)
309
+ )
mage_ai/io/sql.py CHANGED
@@ -43,9 +43,15 @@ class BaseSQL(BaseSQLConnection):
43
43
  self,
44
44
  dtypes: Mapping[str, str],
45
45
  schema_name: str,
46
- table_name: str
46
+ table_name: str,
47
+ unique_constraints: List[str] = [],
47
48
  ) -> str:
48
- return gen_table_creation_query(dtypes, schema_name, table_name)
49
+ return gen_table_creation_query(
50
+ dtypes,
51
+ schema_name,
52
+ table_name,
53
+ unique_constraints=unique_constraints,
54
+ )
49
55
 
50
56
  def build_create_table_as_command(
51
57
  self,
@@ -80,6 +86,8 @@ class BaseSQL(BaseSQLConnection):
80
86
  self,
81
87
  cursor,
82
88
  df: DataFrame,
89
+ db_dtypes: List[str],
90
+ dtypes: List[str],
83
91
  full_table_name: str,
84
92
  buffer: Union[IO, None] = None
85
93
  ) -> None:
@@ -183,6 +191,9 @@ class BaseSQL(BaseSQLConnection):
183
191
  query_string: Union[str, None] = None,
184
192
  drop_table_on_replace: bool = False,
185
193
  cascade_on_drop: bool = False,
194
+ allow_reserved_words: bool = False,
195
+ unique_conflict_method: str = None,
196
+ unique_constraints: List[str] = None,
186
197
  ) -> None:
187
198
  """
188
199
  Exports dataframe to the connected database from a Pandas data frame. If table doesn't
@@ -221,7 +232,10 @@ class BaseSQL(BaseSQLConnection):
221
232
  df = clean_df_for_export(df, self.clean, dtypes)
222
233
 
223
234
  # Clean column names
224
- col_mapping = {col: self._clean_column_name(col) for col in df.columns}
235
+ col_mapping = {col: self._clean_column_name(
236
+ col,
237
+ allow_reserved_words=allow_reserved_words)
238
+ for col in df.columns}
225
239
  df = df.rename(columns=col_mapping)
226
240
  dtypes = infer_dtypes(df)
227
241
 
@@ -263,12 +277,27 @@ class BaseSQL(BaseSQLConnection):
263
277
  )
264
278
  cur.execute(query)
265
279
  else:
280
+ db_dtypes = {col: self.get_type(df[col], dtypes[col]) for col in dtypes}
266
281
  if should_create_table:
267
- db_dtypes = {col: self.get_type(df[col], dtypes[col]) for col in dtypes}
268
- query = self.build_create_table_command(db_dtypes, schema_name, table_name)
282
+ query = self.build_create_table_command(
283
+ db_dtypes,
284
+ schema_name,
285
+ table_name,
286
+ unique_constraints=unique_constraints,
287
+ )
269
288
  cur.execute(query)
270
289
 
271
- self.upload_dataframe(cur, df, full_table_name, buffer)
290
+ self.upload_dataframe(
291
+ cur,
292
+ df,
293
+ db_dtypes,
294
+ dtypes,
295
+ full_table_name,
296
+ buffer,
297
+ allow_reserved_words=allow_reserved_words,
298
+ unique_conflict_method=unique_conflict_method,
299
+ unique_constraints=unique_constraints,
300
+ )
272
301
  self.conn.commit()
273
302
 
274
303
  if verbose:
mage_ai/io/trino.py CHANGED
@@ -15,7 +15,7 @@ from pandas import DataFrame, Series
15
15
  from trino.auth import BasicAuthentication
16
16
  from trino.dbapi import Connection, Cursor as CursorParent
17
17
  from trino.transaction import IsolationLevel
18
- from typing import IO, Mapping, Union
18
+ from typing import IO, List, Mapping, Union
19
19
  import pandas as pd
20
20
 
21
21
 
@@ -86,7 +86,8 @@ class Trino(BaseSQL):
86
86
  self,
87
87
  dtypes: Mapping[str, str],
88
88
  schema_name: str,
89
- table_name: str
89
+ table_name: str,
90
+ unique_constraints: List[str] = [],
90
91
  ):
91
92
  query = []
92
93
  for cname in dtypes:
@@ -123,8 +124,11 @@ class Trino(BaseSQL):
123
124
  self,
124
125
  cursor: Cursor,
125
126
  df: DataFrame,
127
+ db_dtypes: List[str],
128
+ dtypes: List[str],
126
129
  full_table_name: str,
127
- buffer: Union[IO, None] = None
130
+ buffer: Union[IO, None] = None,
131
+ **kwargs,
128
132
  ) -> None:
129
133
  values = []
130
134
  for _, row in df.iterrows():
@@ -81,35 +81,11 @@ class PipelineScheduler:
81
81
  self.schedule()
82
82
 
83
83
  def stop(self) -> None:
84
- if self.pipeline_run.status not in [PipelineRun.PipelineRunStatus.INITIAL,
85
- PipelineRun.PipelineRunStatus.RUNNING]:
86
- return
87
-
88
- self.pipeline_run.update(status=PipelineRun.PipelineRunStatus.CANCELLED)
89
-
90
- # Cancel all the block runs
91
- block_runs_to_cancel = []
92
- running_blocks = []
93
- for b in self.pipeline_run.block_runs:
94
- if b.status in [
95
- BlockRun.BlockRunStatus.INITIAL,
96
- BlockRun.BlockRunStatus.QUEUED,
97
- BlockRun.BlockRunStatus.RUNNING,
98
- ]:
99
- block_runs_to_cancel.append(b)
100
- if b.status == BlockRun.BlockRunStatus.RUNNING:
101
- running_blocks.append(b)
102
- BlockRun.batch_update_status(
103
- [b.id for b in block_runs_to_cancel],
104
- BlockRun.BlockRunStatus.CANCELLED,
84
+ stop_pipeline_run(
85
+ self.pipeline_run,
86
+ self.pipeline,
105
87
  )
106
88
 
107
- if self.pipeline.type in [PipelineType.INTEGRATION, PipelineType.STREAMING]:
108
- job_manager.kill_pipeline_run_job(self.pipeline_run.id)
109
- else:
110
- for b in running_blocks:
111
- job_manager.kill_block_run_job(b.id)
112
-
113
89
  def schedule(self, block_runs: List[BlockRun] = None) -> None:
114
90
  self.__run_heartbeat()
115
91
 
@@ -806,6 +782,40 @@ def run_pipeline(pipeline_run_id, variables, tags):
806
782
  )
807
783
 
808
784
 
785
+ def stop_pipeline_run(
786
+ pipeline_run: PipelineRun,
787
+ pipeline: Pipeline = None,
788
+ ) -> None:
789
+ if pipeline_run.status not in [PipelineRun.PipelineRunStatus.INITIAL,
790
+ PipelineRun.PipelineRunStatus.RUNNING]:
791
+ return
792
+
793
+ pipeline_run.update(status=PipelineRun.PipelineRunStatus.CANCELLED)
794
+
795
+ # Cancel all the block runs
796
+ block_runs_to_cancel = []
797
+ running_blocks = []
798
+ for b in pipeline_run.block_runs:
799
+ if b.status in [
800
+ BlockRun.BlockRunStatus.INITIAL,
801
+ BlockRun.BlockRunStatus.QUEUED,
802
+ BlockRun.BlockRunStatus.RUNNING,
803
+ ]:
804
+ block_runs_to_cancel.append(b)
805
+ if b.status == BlockRun.BlockRunStatus.RUNNING:
806
+ running_blocks.append(b)
807
+ BlockRun.batch_update_status(
808
+ [b.id for b in block_runs_to_cancel],
809
+ BlockRun.BlockRunStatus.CANCELLED,
810
+ )
811
+
812
+ if pipeline and pipeline.type in [PipelineType.INTEGRATION, PipelineType.STREAMING]:
813
+ job_manager.kill_pipeline_run_job(pipeline_run.id)
814
+ else:
815
+ for b in running_blocks:
816
+ job_manager.kill_block_run_job(b.id)
817
+
818
+
809
819
  def check_sla():
810
820
  repo_pipelines = set(Pipeline.get_all_pipelines(get_repo_path()))
811
821
  pipeline_schedules = \
@@ -1,6 +1,9 @@
1
1
  from jupyter_client import KernelClient, KernelManager
2
2
  from jupyter_client.kernelspec import NoSuchKernel
3
3
  from mage_ai.server.kernels import DEFAULT_KERNEL_NAME, KernelName, kernel_managers
4
+ from mage_ai.server.logger import Logger
5
+
6
+ logger = Logger().new_server_logger(__name__)
4
7
 
5
8
 
6
9
  class ActiveKernel():
@@ -13,14 +16,14 @@ active_kernel = ActiveKernel()
13
16
 
14
17
 
15
18
  def switch_active_kernel(kernel_name: KernelName) -> None:
16
- print(f'Switch active kernel: {kernel_name}')
19
+ logger.info(f'Switch active kernel: {kernel_name}')
17
20
  if kernel_managers[kernel_name].is_alive():
18
- print(f'Kernel {kernel_name} is already alive.')
21
+ logger.info(f'Kernel {kernel_name} is already alive.')
19
22
  return
20
23
 
21
24
  for kernel in kernel_managers.values():
22
25
  if kernel.is_alive():
23
- print(f'Shut down current kernel {kernel}.')
26
+ logger.info(f'Shut down current kernel {kernel}.')
24
27
  kernel.request_shutdown()
25
28
 
26
29
  try:
@@ -10,9 +10,12 @@ from mage_ai.cluster_manager.constants import (
10
10
  KUBE_NAMESPACE,
11
11
  KUBE_STORAGE_CLASS_NAME
12
12
  )
13
+ from mage_ai.server.logger import Logger
13
14
  import os
14
15
  import yaml
15
16
 
17
+ logger = Logger().new_server_logger(__name__)
18
+
16
19
 
17
20
  class ClusterType(str, Enum):
18
21
  EMR = 'emr'
@@ -32,7 +35,7 @@ class ApiInstancesHandler(BaseHandler):
32
35
  ecs_instance_manager = EcsTaskManager(cluster_name)
33
36
  instances = ecs_instance_manager.list_tasks()
34
37
  except Exception as e:
35
- print(str(e))
38
+ logger.error(str(e))
36
39
  instances = list()
37
40
  elif cluster_type == ClusterType.CLOUD_RUN:
38
41
  from mage_ai.cluster_manager.gcp.cloud_run_service_manager import CloudRunServiceManager
@@ -1,10 +1,13 @@
1
1
  from mage_ai.data_preparation.models.block import PYTHON_COMMAND
2
+ from mage_ai.server.logger import Logger
2
3
  from typing import List, Dict
3
4
  import importlib
4
5
  import json
5
6
  import subprocess
6
7
  import traceback
7
8
 
9
+ logger = Logger().new_server_logger(__name__)
10
+
8
11
 
9
12
  def get_collection(key: str, available_options: List[Dict]):
10
13
  collection = []
@@ -38,8 +41,8 @@ def get_collection(key: str, available_options: List[Dict]):
38
41
  except Exception:
39
42
  pass
40
43
  except Exception as err:
41
- print(f"Failed to load source {d['uuid']}: {err}")
42
- print(traceback.format_exc())
44
+ logger.error(f"Failed to load source {d['uuid']}: {err}")
45
+ logger.error(traceback.format_exc())
43
46
  continue
44
47
 
45
48
  collection.append(d)
@@ -1,10 +1,10 @@
1
1
  from mage_ai.shared.hash import merge_dict
2
+ from mage_ai.server.logger import Logger
2
3
 
3
4
  import json
4
- import logging
5
5
  import requests
6
6
 
7
- logger = logging.getLogger(__name__)
7
+ logger = Logger().new_server_logger(__name__)
8
8
 
9
9
 
10
10
  class Mage: