mage-ai 0.9.68__py3-none-any.whl → 0.9.69__py3-none-any.whl

mage_ai/data_preparation/models/pipeline.py

@@ -62,6 +62,8 @@ from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.data_preparation.templates.utils import copy_template_directory
 from mage_ai.data_preparation.variable_manager import VariableManager
 from mage_ai.orchestration.constants import Entity
+from mage_ai.orchestration.notification.config import NotificationConfig
+from mage_ai.orchestration.notification.sender import NotificationSender
 from mage_ai.settings.platform import build_repo_path_for_all_projects
 from mage_ai.settings.platform.constants import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
@@ -339,7 +341,7 @@ class Pipeline:
         duplicate_pipeline_dict['name'] = duplicate_pipeline_uuid
         safe_write(
             duplicate_pipeline.config_path,
-            yaml.dump(duplicate_pipeline_dict)
+            yaml.dump(duplicate_pipeline_dict),
         )
 
         tags = duplicate_pipeline_dict.get('tags')
@@ -514,7 +516,7 @@ class Pipeline:
 
         if not os.path.exists(config_path):
             raise Exception(f'Pipeline {uuid} does not exist.')
-        async with aiofiles.open(config_path, mode='r') as f:
+        async with aiofiles.open(config_path, mode='r', encoding='utf-8') as f:
             config = yaml.safe_load(await f.read()) or {}
 
         if PipelineType.INTEGRATION == config.get('type'):
@@ -673,6 +675,7 @@ class Pipeline:
         analyze_outputs: bool = False,
         build_block_output_stdout: Callable[..., object] = None,
         global_vars=None,
+        retry_config=None,
         run_sensors: bool = True,
         run_tests: bool = True,
         update_status: bool = True,
@@ -689,6 +692,7 @@ class Pipeline:
             StreamingPipelineExecutor(self).execute(
                 build_block_output_stdout=build_block_output_stdout,
                 global_vars=global_vars,
+                retry_config=retry_config,
             )
         else:
             root_blocks = []
@@ -715,7 +719,7 @@ class Pipeline:
     def get_config_from_yaml(self):
         if not os.path.exists(self.config_path):
             raise Exception(f'Pipeline {self.uuid} does not exist in repo_path {self.repo_path}.')
-        with open(self.config_path) as fp:
+        with open(self.config_path, encoding='utf-8') as fp:
             config = yaml.full_load(fp) or {}
         return config
 
@@ -726,6 +730,16 @@ class Pipeline:
             config = json.load(f)
         return config
 
+    def get_notification_sender(self):
+        return NotificationSender(
+            NotificationConfig.load(
+                config=merge_dict(
+                    self.repo_config.notification_config,
+                    self.notification_config,
+                ),
+            ),
+        )
+
     def load_config_from_yaml(self):
         catalog = None
         if os.path.exists(self.catalog_config_path):
@@ -1294,7 +1308,9 @@ class Pipeline:
                         file_path = (configuration.get('file_source') or {}).get('path')
                         if file_path:
                             # Check for block name with period to avoid replacing a directory name
-                            new_file_path = file_path.replace(f'{block.name}.', f'{name}.')
+                            new_file_path = file_path.replace(
+                                f'{clean_name(block.name)}.', f'{clean_name(name)}.'
+                            )
                             configuration['file_source']['path'] = new_file_path
                             block_update_payload['configuration'] = configuration
                         blocks_to_remove_from_cache.append(block.to_dict())
@@ -2132,7 +2148,7 @@ class Pipeline:
                 'Blocks cannot be added or removed when saving content, please try again.',
             )
 
-        content = yaml.dump(pipeline_dict)
+        content = yaml.dump(pipeline_dict, allow_unicode=True)
 
         safe_write(self.config_path, content)
 
@@ -2205,7 +2221,7 @@ class Pipeline:
                 'Blocks cannot be added or removed when saving content, please try again.',
             )
 
-        content = yaml.dump(pipeline_dict)
+        content = yaml.dump(pipeline_dict, allow_unicode=True)
 
         test_path = f'{self.config_path}.test'
         async with aiofiles.open(test_path, mode='w', encoding='utf-8') as fp:

mage_ai/data_preparation/models/utils.py

@@ -19,6 +19,12 @@ STRING_SERIALIZABLE_COLUMN_TYPES = set([
     'ObjectId',
 ])
 
+AMBIGUOUS_COLUMN_TYPES = set([
+    'mixed-integer',
+    'complex',
+    'unknown-array',
+])
+
 CAST_TYPE_COLUMN_TYPES = set([
     'Int64',
     'int64',

mage_ai/data_preparation/models/variable.py

@@ -7,7 +7,7 @@ from typing import Any, Dict, List
 import numpy as np
 import pandas as pd
 import polars as pl
-from pandas.api.types import is_object_dtype
+from pandas.api.types import infer_dtype, is_object_dtype
 from pandas.core.indexes.range import RangeIndex
 
 from mage_ai.data_cleaner.shared.utils import is_geo_dataframe, is_spark_dataframe
@@ -18,6 +18,7 @@ from mage_ai.data_preparation.models.constants import (
     VARIABLE_DIR,
 )
 from mage_ai.data_preparation.models.utils import (  # dask_from_pandas,
+    AMBIGUOUS_COLUMN_TYPES,
     STRING_SERIALIZABLE_COLUMN_TYPES,
     apply_transform_pandas,
     cast_column_types,
@@ -282,7 +283,9 @@ class Variable:
         else:
             self.__write_json(data)
 
-        self.write_metadata()
+        if self.variable_type != VariableType.SPARK_DATAFRAME:
+            # Not write json file in spark data directory to avoid read error
+            self.write_metadata()
 
     async def write_data_async(self, data: Any) -> None:
         """
@@ -315,7 +318,9 @@ class Variable:
         else:
             await self.__write_json_async(data)
 
-        self.write_metadata()
+        if self.variable_type != VariableType.SPARK_DATAFRAME:
+            # Not write json file in spark data directory to avoid read error
+            self.write_metadata()
 
     def write_metadata(self) -> None:
         """
@@ -578,10 +583,19 @@ class Variable:
             else:
                 series_non_null = df_col.dropna()
                 if len(series_non_null) > 0:
-                    coltype = type(series_non_null.iloc[0])
+                    sample_element = series_non_null.iloc[0]
+                    coltype = type(sample_element)
+                    coltype_inferred = infer_dtype(series_non_null)
                     if is_object_dtype(series_non_null.dtype):
                         if coltype.__name__ in STRING_SERIALIZABLE_COLUMN_TYPES:
                             cast_coltype = str
+                        # If the column is a "primitive" type, i.e. int/bool/etc and there is
+                        # a mix of types in the column, cast to string
+                        elif (
+                            not hasattr(sample_element, '__dict__')
+                            and coltype_inferred in AMBIGUOUS_COLUMN_TYPES
+                        ):
+                            cast_coltype = str
                         else:
                             cast_coltype = coltype
                         try:

mage_ai/data_preparation/repo_manager.py

@@ -310,8 +310,9 @@ def get_cluster_type(repo_path=None) -> Optional[ClusterType]:
 
 def set_project_uuid_from_metadata() -> None:
     global project_uuid
-    if os.path.exists(get_metadata_path()):
-        with open(get_metadata_path(), 'r', encoding='utf-8') as f:
+    metadata_path = get_metadata_path(root_project=True)
+    if os.path.exists(metadata_path):
+        with open(metadata_path, 'r', encoding='utf-8') as f:
             config = yml.load(f) or {}
             project_uuid = config.get('project_uuid')
 

mage_ai/data_preparation/shared/utils.py

@@ -25,7 +25,7 @@ def get_template_vars_no_db(include_python_libraries: Dict = None) -> Dict[str,
     try:
         from mage_ai.services.aws.secrets_manager.secrets_manager import get_secret
         kwargs['aws_secret_var'] = get_secret
-    except ImportError:
+    except Exception:
         pass
 
     try:

mage_ai/data_preparation/templates/data_exporters/mysql.py

@@ -23,8 +23,8 @@ def export_data_to_mysql(df: DataFrame, **kwargs) -> None:
     with MySQL.with_config(ConfigFileLoader(config_path, config_profile)) as loader:
         loader.export(
             df,
-            None,
-            table_name,
+            schema_name=None,
+            table_name=table_name,
             index=False,  # Specifies whether to include index in exported table
             if_exists='replace',  # Specify resolution policy if table name already exists
         )

mage_ai/data_preparation/templates/data_exporters/oracledb.py

@@ -0,0 +1,27 @@
+from os import path
+
+from pandas import DataFrame
+
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.io.config import ConfigFileLoader
+from mage_ai.io.oracledb import OracleDB
+
+if 'data_exporter' not in globals():
+    from mage_ai.data_preparation.decorators import data_exporter
+
+
+@data_exporter
+def export_data_to_oracledb(df: DataFrame, **kwargs) -> None:
+    """
+    Template to export data to Oracledb.
+    """
+    config_path = path.join(get_repo_path(), 'io_config.yaml')
+    config_profile = 'default'
+    table_name = 'your_table_name'
+
+    with OracleDB.with_config(ConfigFileLoader(config_path, config_profile)) as exporter:
+        exporter.export(
+            df,
+            table_name=table_name,
+            if_exists='replace',  # Specify resolution policy if table name already exists
+        )

mage_ai/data_preparation/templates/repo/metadata.yaml

@@ -43,6 +43,7 @@ spark_config:
   # e.g. kwargs['context']['spark'] = spark_session
   custom_session_var_name: 'spark'
 
+help_improve_mage: true
 notification_config:
   alert_on:
     - trigger_failure

mage_ai/io/bigquery.py

@@ -1,3 +1,4 @@
+import uuid
 from typing import Dict, List, Mapping, Union
 
 import numpy as np
@@ -14,6 +15,7 @@ from sqlglot import exp, parse_one
 
 from mage_ai.io.base import QUERY_ROW_LIMIT, BaseSQLDatabase, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
 from mage_ai.io.export_utils import infer_dtypes
 from mage_ai.shared.custom_logger import DX_PRINTER
 from mage_ai.shared.environments import is_debug
@@ -200,6 +202,9 @@ WHERE TABLE_NAME = '{table_name}'
         overwrite_types: Dict = None,
         query_string: Union[str, None] = None,
         verbose: bool = True,
+        unique_conflict_method: str = None,
+        unique_constraints: List[str] = None,
+        write_disposition: str = None,
         **configuration_params,
     ) -> None:
         """
@@ -230,24 +235,23 @@ WHERE TABLE_NAME = '{table_name}'
         elif type(df) is list:
             df = DataFrame(df)
 
-        def __process(database: Union[str, None]):
-            if query_string:
-                parts = table_id.split('.')
-                if len(parts) == 2:
-                    schema, table_name = parts
-                elif len(parts) == 3:
-                    database, schema, table_name = parts
+        def __process(database: Union[str, None], write_disposition: str = None):
+            parts = table_id.split('.')
+            if len(parts) == 2:
+                schema, table_name = parts
+            elif len(parts) == 3:
+                database, schema, table_name = parts
 
-                df_existing = self.client.query(f"""
+            df_existing = self.client.query(f"""
 SELECT 1
 FROM `{database}.{schema}.__TABLES_SUMMARY__`
 WHERE table_id = '{table_name}'
 """).to_dataframe()
 
-                full_table_name = f'`{database}.{schema}.{table_name}`'
-
-                table_doesnt_exist = df_existing.empty
+            full_table_name = f'`{database}.{schema}.{table_name}`'
 
+            table_doesnt_exist = df_existing.empty
+            if query_string:
                 if ExportWritePolicy.FAIL == if_exists and not table_doesnt_exist:
                     raise ValueError(
                         f'Table \'{full_table_name}\' already exists in database.',
@@ -268,60 +272,129 @@ WHERE table_id = '{table_name}'
                 self.client.query(sql)
 
             else:
-                config = LoadJobConfig(**configuration_params)
-                if overwrite_types is not None:
-                    config.schema = [SchemaField(k, v) for k, v in overwrite_types.items()]
-                if 'write_disposition' not in configuration_params:
-                    if if_exists == 'replace':
-                        config.write_disposition = WriteDisposition.WRITE_TRUNCATE
-                    elif if_exists == 'append':
-                        config.write_disposition = WriteDisposition.WRITE_APPEND
-                    elif if_exists == 'fail':
-                        config.write_disposition = WriteDisposition.WRITE_EMPTY
-                    else:
-                        raise ValueError(
-                            f'Invalid policy specified for handling existence of '
-                            f'table: \'{if_exists}\''
+                if (
+                    if_exists == ExportWritePolicy.APPEND
+                    and not table_doesnt_exist
+                    and unique_constraints
+                    and unique_conflict_method
+                ):
+                    temp_table_id = f'{table_id}_{uuid.uuid4().hex}'
+
+                    try:
+                        self.__write_table(
+                            df,
+                            temp_table_id,
+                            overwrite_types=overwrite_types,
+                            **configuration_params,
+                        )
+
+                        parts = temp_table_id.split('.')
+                        if len(parts) == 2:
+                            temp_table_name = parts[1]
+                        elif len(parts) == 3:
+                            temp_table_name = parts[2]
+                        column_types = self.get_column_types(schema, temp_table_name)
+                        columns = list(column_types.keys())
+                        if not columns:
+                            columns = df.columns.str.replace(' ', '_')
+
+                        on_conditions = []
+                        for col in unique_constraints:
+                            on_conditions.append(
+                                f'((a.{col} IS NULL AND b.{col} IS NULL) OR a.{col} = b.{col})',
+                            )
+
+                        insert_columns = ', '.join([f'`{col}`' for col in columns])
+
+                        merge_commands = [
+                            f'MERGE INTO `{table_id}` AS a',
+                            f'USING (SELECT * FROM `{temp_table_id}`) AS b',
+                            f"ON {' AND '.join(on_conditions)}",
+                        ]
+
+                        if UNIQUE_CONFLICT_METHOD_UPDATE == unique_conflict_method:
+                            set_command = ', '.join(
+                                [f'a.`{col}` = b.`{col}`' for col in columns],
+                            )
+                            merge_commands.append(f'WHEN MATCHED THEN UPDATE SET {set_command}')
+
+                        merge_values = f"({', '.join([f'b.`{col}`' for col in columns])})"
+                        merge_commands.append(
+                            f'WHEN NOT MATCHED THEN INSERT ({insert_columns}) VALUES {merge_values}',  # noqa: E501
                         )
-                parts = table_id.split('.')
-                if len(parts) == 2:
-                    schema = parts[0]
-                    table_name = parts[1]
-                elif len(parts) == 3:
-                    schema = parts[1]
-                    table_name = parts[2]
-
-                self.client.create_dataset(dataset=schema, exists_ok=True)
-
-                column_types = self.get_column_types(schema, table_name)
-
-                if df is not None:
-                    df.fillna(value=np.NaN, inplace=True)
-                    for col in df.columns:
-                        col_type = column_types.get(col)
-                        if not col_type:
-                            continue
-
-                        null_rows = df[col].isnull()
-                        if col_type.startswith('ARRAY<STRUCT'):
-                            df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [{}])
-                        elif col_type.startswith('ARRAY'):
-                            df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [])
-                        elif col_type.startswith('STRUCT'):
-                            df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: {})
-
-                    # Clean column names
-                    if type(df) is DataFrame:
-                        df.columns = df.columns.str.replace(' ', '_')
-
-                    self.client.load_table_from_dataframe(df, table_id, job_config=config).result()
+
+                        merge_command = '\n'.join(merge_commands)
+
+                        self.client.query(merge_command).result()
+                    finally:
+                        self.client.query(f'DROP TABLE IF EXISTS {temp_table_id}').result()
+                else:
+                    if not write_disposition:
+                        if if_exists == ExportWritePolicy.APPEND:
+                            write_disposition = WriteDisposition.WRITE_APPEND
+                        elif if_exists == ExportWritePolicy.REPLACE:
+                            write_disposition = WriteDisposition.WRITE_TRUNCATE
+                        elif if_exists == ExportWritePolicy.FAIL:
+                            write_disposition = WriteDisposition.WRITE_EMPTY
+                    self.__write_table(
+                        df,
+                        table_id,
+                        overwrite_types=overwrite_types,
+                        write_disposition=write_disposition,
+                        **configuration_params,
+                    )
 
         if verbose:
             with self.printer.print_msg(f'Exporting data to table \'{table_id}\''):
-                __process(database=database)
+                __process(database=database, write_disposition=write_disposition)
         else:
             __process(database=database)
 
+    def __write_table(
+        self,
+        df: DataFrame,
+        table_id: str,
+        overwrite_types: Dict = None,
+        **configuration_params,
+    ):
+        config = LoadJobConfig(**configuration_params)
+        if overwrite_types is not None:
+            config.schema = [SchemaField(k, v) for k, v in overwrite_types.items()]
+        if not config.write_disposition:
+            config.write_disposition = WriteDisposition.WRITE_APPEND
+        parts = table_id.split('.')
+        if len(parts) == 2:
+            schema = parts[0]
+            table_name = parts[1]
+        elif len(parts) == 3:
+            schema = parts[1]
+            table_name = parts[2]
+
+        self.client.create_dataset(dataset=schema, exists_ok=True)
+
+        column_types = self.get_column_types(schema, table_name)
+
+        if df is not None:
+            df.fillna(value=np.NaN, inplace=True)
+            for col in df.columns:
+                col_type = column_types.get(col)
+                if not col_type:
+                    continue
+
+                null_rows = df[col].isnull()
+                if col_type.startswith('ARRAY<STRUCT'):
+                    df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [{}])
+                elif col_type.startswith('ARRAY'):
+                    df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: [])
+                elif col_type.startswith('STRUCT'):
+                    df.loc[null_rows, col] = df.loc[null_rows, col].apply(lambda x: {})
+
+            # Clean column names
+            if type(df) is DataFrame:
+                df.columns = df.columns.str.replace(' ', '_')
+
+            return self.client.load_table_from_dataframe(df, table_id, job_config=config).result()
+
     def execute(self, query_string: str, **kwargs) -> None:
         """
         Sends query to the connected BigQuery warehouse.

mage_ai/io/mysql.py

@@ -8,6 +8,7 @@ from mysql.connector.cursor import MySQLCursor
 from pandas import DataFrame, Series
 
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
 from mage_ai.io.export_utils import PandasTypes
 from mage_ai.io.sql import BaseSQL
 from mage_ai.shared.parsers import encode_complex
@@ -60,7 +61,7 @@ class MySQL(BaseSQL):
     ) -> str:
         if unique_constraints is None:
             unique_constraints = []
-        query = []
+        columns_and_types = []
         for cname in dtypes:
             if overwrite_types is not None and cname in overwrite_types.keys():
                 dtypes[cname] = overwrite_types[cname]
@@ -68,9 +69,22 @@ class MySQL(BaseSQL):
                 cleaned_col_name = clean_name(cname, case_sensitive=case_sensitive)
             else:
                 cleaned_col_name = cname
-            query.append(f'`{cleaned_col_name}` {dtypes[cname]} NULL')
-
-        return f'CREATE TABLE {table_name} (' + ','.join(query) + ');'
+            columns_and_types.append(f'`{cleaned_col_name}` {dtypes[cname]} NULL')
+
+        if unique_constraints:
+            unique_constraints = [
+                clean_name(col, case_sensitive=case_sensitive)
+                for col in unique_constraints
+            ]
+            index_name = '_'.join([
+                clean_name(table_name, case_sensitive=case_sensitive),
+            ] + unique_constraints)
+            index_name = f'unique{index_name}'[:64]
+            columns_and_types.append(
+                f"CONSTRAINT {index_name} Unique({', '.join(unique_constraints)})"
+            )
+
+        return f'CREATE TABLE {table_name} (' + ','.join(columns_and_types) + ');'
 
     def open(self) -> None:
         with self.printer.print_msg('Opening connection to MySQL database'):
@@ -94,6 +108,9 @@ class MySQL(BaseSQL):
         dtypes: List[str],
         full_table_name: str,
         buffer: Union[IO, None] = None,
+        case_sensitive: bool = False,
+        unique_constraints: List[str] = None,
+        unique_conflict_method: str = None,
         **kwargs,
     ) -> None:
         def serialize_obj(val):
@@ -133,9 +150,24 @@ class MySQL(BaseSQL):
         for _, row in df_.iterrows():
             values.append(tuple([str(val) if type(val) is pd.Timestamp else val for val in row]))
 
-        insert_columns = ', '.join([f'`{col}`'for col in columns])
+        cleaned_columns = [clean_name(col, case_sensitive=case_sensitive) for col in columns]
+        insert_columns = ', '.join([f'`{col}`'for col in cleaned_columns])
+
+        query = [
+            f'INSERT INTO {full_table_name} ({insert_columns})',
+            f'VALUES ({values_placeholder})',
+        ]
+
+        if unique_constraints and unique_conflict_method:
+            if UNIQUE_CONFLICT_METHOD_UPDATE == unique_conflict_method:
+                update_command = [f'{col} = new.{col}' for col in cleaned_columns]
+                query += [
+                    'AS new',
+                    f"ON DUPLICATE KEY UPDATE {', '.join(update_command)}",
+                ]
+
+        sql = '\n'.join(query)
 
-        sql = f'INSERT INTO {full_table_name} ({insert_columns}) VALUES ({values_placeholder})'
         cursor.executemany(sql, values)
 
     def get_type(self, column: Series, dtype: str) -> str: