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.
- mage_ai/api/logging.py +2 -2
- mage_ai/api/resources/PipelineRunResource.py +6 -3
- mage_ai/data_integrations/sources/constants.py +1 -0
- mage_ai/data_preparation/models/block/dbt/__init__.py +20 -0
- mage_ai/data_preparation/models/block/dbt/utils/__init__.py +4 -1
- mage_ai/data_preparation/models/pipeline.py +10 -1
- mage_ai/io/base.py +2 -2
- mage_ai/io/constants.py +3 -0
- mage_ai/io/export_utils.py +14 -1
- mage_ai/io/mssql.py +5 -2
- mage_ai/io/mysql.py +7 -3
- mage_ai/io/postgres.py +64 -21
- mage_ai/io/sql.py +35 -6
- mage_ai/io/trino.py +7 -3
- mage_ai/orchestration/pipeline_scheduler.py +37 -27
- mage_ai/server/active_kernel.py +6 -3
- mage_ai/server/api/clusters.py +4 -1
- mage_ai/server/api/integration_sources.py +5 -2
- mage_ai/server/client/mage.py +2 -2
- mage_ai/server/constants.py +1 -1
- mage_ai/server/data/base.py +2 -2
- mage_ai/server/data/models.py +2 -2
- mage_ai/server/frontend_dist/404.html +2 -2
- mage_ai/server/frontend_dist/404.html.html +2 -2
- mage_ai/server/frontend_dist/_next/static/chunks/{2626-905774aafeb2c600.js → 2626-e7fa4f83f8214c97.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/4178-9103014b7dae3c49.js +1 -0
- mage_ai/server/frontend_dist/_next/static/chunks/{4538-8a3c3e47be976ede.js → 4538-347283088b83c6bf.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/5141-ddf4ba0a362d6f34.js +1 -0
- mage_ai/server/frontend_dist/_next/static/chunks/{5477-793cd2120261d023.js → 5477-b439f211b6146a11.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/{5872-103815a4a043489b.js → 5872-1767c45ee6690ae5.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/{5896-f84e336fb8877027.js → 5896-10a676bcc86978cc.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/{7400-365cb7888b6db7d9.js → 7400-f4db9b5d41f67f75.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/{9386-fb899ca8ecc2a350.js → 9386-d4cc11bab74eec8d.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/{9832-c8b8970bb522f302.js → 9832-f97919376d52e3bf.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/{manage-9e5f315db570ac77.js → manage-3046bc53d24917c7.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/{pipeline-runs-79e10a783afec3df.js → pipeline-runs-e64ba4e8b2bfe73c.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-017f6a2de38f8658.js +1 -0
- mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/{syncs-1767a2f57f887ef7.js → syncs-e1271453ed0c8d6e.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/{pipelines-52c3ee3817e5554b.js → pipelines-7446a70bdd8381a5.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/{preferences-a23b61bab04a16f3.js → preferences-997acba85f777259.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/settings/workspace/{sync-data-de28e502102defde.js → sync-data-8c903140c99e487c.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/terminal-6c0ea500b3bc6b61.js +1 -0
- mage_ai/server/frontend_dist/_next/static/chunks/pages/{triggers-8a2169d30b643ae7.js → triggers-783b9526167f1249.js} +1 -1
- mage_ai/server/frontend_dist/_next/static/{23nXkA2GRjJCZDwY3kuKd → ngCg1gM6urg83y2nV5KHK}/_buildManifest.js +1 -1
- mage_ai/server/frontend_dist/index.html +2 -2
- mage_ai/server/frontend_dist/manage.html +2 -2
- mage_ai/server/frontend_dist/pipeline-runs.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills/[...slug].html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/backfills.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/edit.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/logs.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runs.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors/block-runtime.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/monitors.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/runs/[run].html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/runs.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/syncs.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers/[...slug].html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline]/triggers.html +2 -2
- mage_ai/server/frontend_dist/pipelines/[pipeline].html +2 -2
- mage_ai/server/frontend_dist/pipelines.html +2 -2
- mage_ai/server/frontend_dist/settings/account/profile.html +2 -2
- mage_ai/server/frontend_dist/settings/workspace/preferences.html +2 -2
- mage_ai/server/frontend_dist/settings/workspace/sync-data.html +2 -2
- mage_ai/server/frontend_dist/settings/workspace/users.html +2 -2
- mage_ai/server/frontend_dist/settings.html +2 -2
- mage_ai/server/frontend_dist/sign-in.html +2 -2
- mage_ai/server/frontend_dist/terminal.html +2 -2
- mage_ai/server/frontend_dist/test.html +2 -2
- mage_ai/server/frontend_dist/triggers.html +2 -2
- mage_ai/server/logger.py +16 -0
- mage_ai/server/scheduler_manager.py +8 -4
- mage_ai/server/server.py +7 -14
- mage_ai/server/subscriber.py +6 -3
- mage_ai/server/terminal_server.py +65 -0
- mage_ai/server/utils/frontend_renderer.py +3 -2
- mage_ai/server/websocket_server.py +4 -1
- mage_ai/settings/__init__.py +2 -0
- mage_ai/shared/logger.py +4 -0
- {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/METADATA +1 -1
- {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/RECORD +87 -85
- mage_ai/server/frontend_dist/_next/static/chunks/4178-e17f37d21253b832.js +0 -1
- mage_ai/server/frontend_dist/_next/static/chunks/5141-2ae9eae00ec2cdfa.js +0 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/pipelines/[pipeline]/edit-6a52671f1853e1a5.js +0 -1
- mage_ai/server/frontend_dist/_next/static/chunks/pages/terminal-9c21edae8f1b6737.js +0 -1
- /mage_ai/server/frontend_dist/_next/static/{23nXkA2GRjJCZDwY3kuKd → ngCg1gM6urg83y2nV5KHK}/_middlewareManifest.js +0 -0
- /mage_ai/server/frontend_dist/_next/static/{23nXkA2GRjJCZDwY3kuKd → ngCg1gM6urg83y2nV5KHK}/_ssgManifest.js +0 -0
- {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/LICENSE +0 -0
- {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/WHEEL +0 -0
- {mage_ai-0.8.37.dist-info → mage_ai-0.8.39.dist-info}/entry_points.txt +0 -0
- {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
|
@@ -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
|
-
|
|
222
|
+
pipeline = Pipeline.get(
|
|
223
|
+
self.model.pipeline_uuid,
|
|
224
|
+
check_if_exists=True,
|
|
225
|
+
)
|
|
223
226
|
|
|
224
|
-
|
|
227
|
+
stop_pipeline_run(self.model, pipeline)
|
|
225
228
|
|
|
226
229
|
return self
|
|
@@ -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
|
-
|
|
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
mage_ai/io/export_utils.py
CHANGED
|
@@ -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
|
|
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
|
|
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
|
-
|
|
250
|
-
buffer,
|
|
251
|
-
header=False,
|
|
252
|
-
index=False,
|
|
253
|
-
na_rep='',
|
|
254
|
-
)
|
|
275
|
+
values = []
|
|
255
276
|
|
|
256
|
-
|
|
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
|
-
|
|
259
|
-
|
|
260
|
-
|
|
261
|
-
|
|
262
|
-
|
|
263
|
-
|
|
264
|
-
|
|
265
|
-
)
|
|
266
|
-
|
|
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(
|
|
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(
|
|
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
|
-
|
|
268
|
-
|
|
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(
|
|
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
|
-
|
|
85
|
-
|
|
86
|
-
|
|
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 = \
|
mage_ai/server/active_kernel.py
CHANGED
|
@@ -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
|
-
|
|
19
|
+
logger.info(f'Switch active kernel: {kernel_name}')
|
|
17
20
|
if kernel_managers[kernel_name].is_alive():
|
|
18
|
-
|
|
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
|
-
|
|
26
|
+
logger.info(f'Shut down current kernel {kernel}.')
|
|
24
27
|
kernel.request_shutdown()
|
|
25
28
|
|
|
26
29
|
try:
|
mage_ai/server/api/clusters.py
CHANGED
|
@@ -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
|
-
|
|
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
|
-
|
|
42
|
-
|
|
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)
|
mage_ai/server/client/mage.py
CHANGED