mage-ai 0.9.74__py3-none-any.whl → 0.9.79__py3-none-any.whl

mage_ai/data_preparation/models/block/dbt/dbt_adapter.py

@@ -5,7 +5,9 @@ from typing import Any, Dict, Optional, Tuple, Union
 
 import dbt.flags as flags
 import pandas as pd
-from dbt.adapters.base import BaseRelation, Credentials
+from dbt.adapters.base.relation import BaseRelation
+from dbt.adapters.contracts.connection import AdapterResponse, Credentials
+from dbt.adapters.contracts.relation import RelationType
 from dbt.adapters.factory import (
     Adapter,
     cleanup_connections,
@@ -13,10 +15,10 @@ from dbt.adapters.factory import (
     register_adapter,
     reset_adapters,
 )
-from dbt.config.profile import read_user_config
+from dbt.config.project import read_project_flags
 from dbt.config.runtime import RuntimeConfig
-from dbt.contracts.connection import AdapterResponse
-from dbt.contracts.relation import RelationType
+from dbt.context.providers import generate_runtime_macro_context
+from dbt.mp_context import get_mp_context
 
 from mage_ai.data_preparation.models.block.dbt.profiles import Profiles
 from mage_ai.shared.environments import is_debug
@@ -79,7 +81,11 @@ class DBTAdapter:
         # remove interpolated profiles.yml
         self.__profiles.clean()
 
-    def execute(self, sql: str, fetch: bool = False) -> Tuple[AdapterResponse, pd.DataFrame]:
+    def execute(
+        self,
+        sql: str,
+        fetch: bool = False
+    ) -> Tuple[AdapterResponse, pd.DataFrame]:
         """
         Executes any sql statement using the dbt adapter.
 
@@ -135,6 +141,8 @@ class DBTAdapter:
             package
         )
 
+        self.__adapter.set_macro_resolver(manifest)
+
         # create a context for the macro (e.g. downstream macros)
         from dbt.context.providers import generate_runtime_macro_context
         macro_context = generate_runtime_macro_context(
@@ -203,7 +211,7 @@ class DBTAdapter:
         # set dbt flags
         # Need to add profiles.yml file
         try:
-            user_config = read_user_config(profiles_path)
+            user_config = read_project_flags(self.project_path, profiles_path)
         except Exception as err:
             print(f'[ERROR] DBTAdapter.open: {err}.')
 
@@ -211,7 +219,10 @@ class DBTAdapter:
                     not profiles_path.endswith('profiles.yml'):
 
                 try:
-                    user_config = read_user_config(os.path.join(profiles_path, 'profiles.yml'))
+                    user_config = read_project_flags(
+                        self.project_path,
+                        os.path.join(profiles_path, 'profiles.yml')
+                    )
                 except Exception as err2:
                     print(f'[ERROR] DBTAdapter.open: {err2}.')
                     raise err
@@ -227,9 +238,10 @@ class DBTAdapter:
             config = RuntimeConfig.from_args(adapter_config)
             reset_adapters()
             # register the correct adapter from config
-            register_adapter(config)
+            register_adapter(config, mp_context=get_mp_context())
             # load the adapter
             self.__adapter = get_adapter(config)
+            self.__adapter.set_macro_context_generator(generate_runtime_macro_context)
             # connect
             self.__adapter.acquire_connection('mage_dbt_adapter_' + uuid.uuid4().hex)
             return self

mage_ai/data_preparation/models/block/dynamic/constants.py

@@ -1 +0,0 @@
-CHILD_DATA_VARIABLE_UUID = 'output_0'

mage_ai/data_preparation/models/block/dynamic/counter.py

@@ -1,14 +1,12 @@
 from typing import Any, Iterable, List, Optional, Sequence
 
 from mage_ai.data.tabular.reader import read_metadata
-from mage_ai.data_preparation.models.block.dynamic.constants import (
-    CHILD_DATA_VARIABLE_UUID,
-)
 from mage_ai.data_preparation.models.block.dynamic.utils import (
     is_dynamic_block,
     is_dynamic_block_child,
     should_reduce_output,
 )
+from mage_ai.data_preparation.models.constants import CHILD_DATA_VARIABLE_UUID
 from mage_ai.data_preparation.models.variables.cache import VariableAggregateCache
 from mage_ai.data_preparation.models.variables.constants import (
     VariableAggregateDataType,

mage_ai/data_preparation/models/block/outputs.py

@@ -343,7 +343,13 @@ def format_output_data(
                 columns=columns_to_display,
                 rows=[
                     list(row.values())
-                    for row in json.loads(data[columns_to_display].write_json(row_oriented=True))
+                    for row in json.loads(
+                        simplejson.dumps(
+                            data[columns_to_display].to_dicts(),
+                            default=encode_complex,
+                            ignore_nan=True,
+                        )
+                    )
                 ],
             ),
             resource_usage=resource_usage,

mage_ai/data_preparation/models/block/r/__init__.py

@@ -8,7 +8,10 @@ import pandas as pd
 import simplejson
 
 from mage_ai.data_preparation.models.block import Block
-from mage_ai.data_preparation.models.constants import BlockType
+from mage_ai.data_preparation.models.constants import (
+    CHILD_DATA_VARIABLE_UUID,
+    BlockType,
+)
 from mage_ai.data_preparation.models.variables.constants import (
     DATAFRAME_CSV_FILE,
     VariableType,
@@ -81,7 +84,10 @@ def execute_r_code(
 
     if len(output_variable_objects) > 0:
         df = pd.read_csv(
-            os.path.join(output_variable_objects[0].variable_path, DATAFRAME_CSV_FILE)
+            os.path.join(
+                output_variable_objects[0].variable_path,
+                DATAFRAME_CSV_FILE
+            )
         )
     else:
         df = None
@@ -129,14 +135,19 @@ def __render_r_script(
         raise Exception(
             f"Block execution for {block.type} with R language is not supported.",
         )
-    template = template_env.get_template(BLOCK_TYPE_TO_EXECUTION_TEMPLATE[block.type])
+    template = template_env.get_template(
+        BLOCK_TYPE_TO_EXECUTION_TEMPLATE[block.type]
+    )
 
     output_variable_object = block.variable_object(
-        "output_0",
+        CHILD_DATA_VARIABLE_UUID,
         execution_partition=execution_partition,
     )
     os.makedirs(output_variable_object.variable_path, exist_ok=True)
-    output_path = os.path.join(output_variable_object.variable_path, DATAFRAME_CSV_FILE)
+    output_path = os.path.join(
+        output_variable_object.variable_path,
+        DATAFRAME_CSV_FILE
+    )
 
     global_vars_str = __render_global_vars(global_vars=global_vars)
 

mage_ai/data_preparation/models/block/sql/__init__.py

@@ -150,6 +150,7 @@ def execute_sql_code(
     interpolate_vars_options = dict(
         block=block,
         dynamic_block_index=dynamic_block_index,
+        execution_partition=execution_partition,
         global_vars=global_vars,
     )
 
@@ -368,6 +369,7 @@ def execute_sql_code(
                     loader,
                     block,
                     query_string,
+                    disable_query_preprocessing=disable_query_preprocessing,
                     configuration=configuration,
                     should_query=should_query,
                 )

mage_ai/data_preparation/models/block/sql/mssql.py

@@ -67,9 +67,17 @@ def execute_raw_sql(
     query_string: str,
     configuration: Dict = None,
     should_query: bool = False,
+    disable_query_preprocessing: bool = False,
 ) -> List[Any]:
     if configuration is None:
         configuration = {}
+
+    if disable_query_preprocessing:
+        return loader.execute_query_raw(
+            query_string,
+            configuration=configuration,
+        )
+
     queries = []
     fetch_query_at_indexes = []
 

mage_ai/data_preparation/models/block/sql/utils/shared.py

@@ -232,7 +232,8 @@ def interpolate_vars(
     content: str,
     global_vars: Dict = None,
     block=None,
-    dynamic_block_index: int = None
+    dynamic_block_index: int = None,
+    execution_partition: str = None,
 ) -> str :
     if not content:
         return content
@@ -241,7 +242,10 @@ def interpolate_vars(
 
     if block:
         content = block.interpolate_content(
-            content, variables=global_vars, dynamic_block_index=dynamic_block_index,
+            content,
+            variables=global_vars,
+            dynamic_block_index=dynamic_block_index,
+            execution_partition=execution_partition,
         )
 
     return Template(

mage_ai/data_preparation/models/constants.py

@@ -15,7 +15,7 @@ DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json'
 DATAFRAME_ANALYSIS_MAX_COLUMNS = 100
 DATAFRAME_ANALYSIS_MAX_ROWS = 100_000
 DATAFRAME_SAMPLE_COUNT = 1000
-DATAFRAME_SAMPLE_COUNT_PREVIEW = 10
+DATAFRAME_SAMPLE_COUNT_PREVIEW = int(os.getenv('DATAFRAME_SAMPLE_COUNT_PREVIEW', 10) or 10)
 DATAFRAME_SAMPLE_MAX_COLUMNS = 1000
 DYNAMIC_CHILD_BLOCK_SAMPLE_COUNT_PREVIEW = 10
 LOGS_SUBDIR = '.logs'
@@ -27,6 +27,7 @@ PIPELINES_FOLDER = 'pipelines'
 PREFERENCES_FILE = '.preferences.yaml'
 REPO_CONFIG_FILE = 'metadata.yaml'
 VARIABLE_DIR = '.variables'
+CHILD_DATA_VARIABLE_UUID = 'output_0'
 
 PIPELINE_RUN_STATUS_LAST_RUN_FAILED = 'last_run_failed'
 
@@ -81,6 +82,7 @@ class BlockColor(StrEnum):
 
 
 class CallbackStatus(StrEnum):
+    CANCELLED = 'cancelled'
     FAILURE = 'failure'
     SUCCESS = 'success'
 
@@ -125,6 +127,7 @@ BLOCK_LANGUAGE_TO_FILE_EXTENSION = {
 
 
 CALLBACK_STATUSES = [
+    CallbackStatus.CANCELLED,
     CallbackStatus.FAILURE,
     CallbackStatus.SUCCESS,
 ]

mage_ai/data_preparation/models/pipeline.py

@@ -4,6 +4,7 @@ import json
 import os
 import shutil
 import tempfile
+import time
 import zipfile
 from datetime import datetime, timezone
 from io import BytesIO
@@ -77,6 +78,7 @@ 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
+from mage_ai.settings.server import RUN_PIPELINE_IN_ONE_PROCESS
 from mage_ai.shared.array import find
 from mage_ai.shared.hash import extract, ignore_keys, index_by, merge_dict
 from mage_ai.shared.io import safe_write, safe_write_async
@@ -127,7 +129,7 @@ class Pipeline:
 
         self.repo_path = repo_path or get_repo_path()
         self.retry_config = {}
-        self.run_pipeline_in_one_process = False
+        self.run_pipeline_in_one_process = RUN_PIPELINE_IN_ONE_PROCESS
         self.schedules = []
         self.settings = {}
         self.tags = tags or []
@@ -884,7 +886,10 @@ class Pipeline:
         self.executor_type = config.get('executor_type')
         self.notification_config = config.get('notification_config') or {}
         self.retry_config = config.get('retry_config') or {}
-        self.run_pipeline_in_one_process = config.get('run_pipeline_in_one_process', False)
+        self.run_pipeline_in_one_process = config.get(
+            'run_pipeline_in_one_process',
+            RUN_PIPELINE_IN_ONE_PROCESS,
+        )
         self.settings = PipelineSettings.load(**config.get('settings') or {})
         self.spark_config = config.get('spark_config') or {}
         self.tags = config.get('tags') or []
@@ -2328,6 +2333,10 @@ class Pipeline:
         extension_uuid: str = None,
         widget: bool = False,
     ):
+        # Introduce a small delay to prevent multiple changes from generating
+        # identical timestamps for the pipeline YAML file
+        time.sleep(0.0005)
+
         blocks_current = sorted([b.uuid for b in self.blocks_by_uuid.values()])
 
         if block_uuid is not None:

mage_ai/data_preparation/models/project/__init__.py

@@ -214,10 +214,12 @@ class Project:
         return [self.repo_path]
 
     def projects(self) -> Dict:
-        return project_platform_settings(
+        projects = project_platform_settings(
             context_data=self.context_data,
             mage_projects_only=True
         )
+        sorted_projects = {k: projects[k] for k in sorted(projects)}
+        return sorted_projects
 
     def is_feature_enabled(self, feature_name: FeatureUUID) -> bool:
         feature_enabled = self.features.get(feature_name.value, False)

mage_ai/data_preparation/models/triggers/__init__.py

@@ -5,12 +5,12 @@ from datetime import datetime, timezone
 from typing import Dict, List, Optional
 
 import yaml
-from croniter import croniter
 
 from mage_ai.data_preparation.models.constants import PIPELINES_FOLDER
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.config import BaseConfig
 from mage_ai.shared.constants import VALID_ENVS
+from mage_ai.shared.croniter import croniter
 from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import index_by
 from mage_ai.shared.io import safe_write

mage_ai/data_preparation/storage/local_storage.py

@@ -121,7 +121,10 @@ class LocalStorage(BaseStorage):
         return pd.read_parquet(file_path, engine='pyarrow')
 
     def read_polars_parquet(self, file_path: str, **kwargs) -> pl.DataFrame:
-        return pl.read_parquet(file_path, use_pyarrow=True)
+        try:
+            return pl.read_parquet(file_path, use_pyarrow=True)
+        except Exception:
+            return pl.read_parquet(file_path)
 
     def write_csv(self, df: pd.DataFrame, file_path: str) -> None:
         File.create_parent_directories(file_path)

mage_ai/data_preparation/templates/constants.py

@@ -19,6 +19,13 @@ GROUP_ROW_ACTIONS = 'Row actions'
 GROUP_SHIFT = 'Shift'
 
 TEMPLATES = [
+    dict(
+        block_type=BlockType.DATA_LOADER,
+        description='Load a Table from Airtable App.',
+        language=BlockLanguage.PYTHON,
+        name='Airtable',
+        path='data_loaders/airtable.py',
+    ),
     dict(
         block_type=BlockType.DATA_LOADER,
         description='Load a Delta Table from Amazon S3.',

mage_ai/data_preparation/templates/data_exporters/streaming/elasticsearch.yaml

@@ -2,6 +2,9 @@ connector_type: elasticsearch
 host: http://localhost
 index_name: elastic_index
 api_key: api_key_for_elastic_search
+# # When using Data streams this must be set to create as they are Append-only indexes
+# _op_type: None / index, create or update
+# _source: None (Default '_doc') / Specify the source field
 # # Whether to verify SSL certificates to authenticate.
 # verify_certs: true
 # ca_cert: "certificate.cert"

mage_ai/data_preparation/templates/data_loaders/airtable.py

@@ -0,0 +1,28 @@
+{% extends "data_loaders/default.jinja" %}
+{% block imports %}
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.io.config import ConfigFileLoader
+from mage_ai.io.airtable import Airtable
+from os import path
+{{ super() -}}
+{% endblock %}
+
+
+{% block content %}
+@data_loader
+def load_data_from_airtable(*args, **kwargs):
+    """
+    Template for loading data from Airtable.
+    Specify your configuration settings in 'io_config.yaml'.
+    """
+    config_path = path.join(get_repo_path(), 'io_config.yaml')
+    config_profile = 'default'
+
+    base_id = 'your_base_id'
+    table_name = 'your_table_name'
+
+    return Airtable.with_config(ConfigFileLoader(config_path, config_profile)).load(
+        base_id=base_id,
+        table_name=table_name
+    )
+{% endblock %}

mage_ai/data_preparation/templates/data_loaders/streaming/nats.yaml

@@ -17,9 +17,12 @@ timeout: 5
 
 # Optional: TLS configuration
 use_tls: false
-# cafile: "path/to/ca.pem"
-# certfile: "path/to/client-cert.pem"
-# keyfile: "path/to/client-key.pem"
+
+# Optional: SSL configuration
+# ssl_config:
+#   cafile: "path/to/ca.pem"
+#   certfile: "path/to/client-cert.pem"
+#   keyfile: "path/to/client-key.pem"
 
 # Optional: NKEY seed string
 # nkeys_seed_str: "SUAPAEYJQOQHJ4"

mage_ai/data_preparation/templates/repo/io_config.yaml

@@ -11,6 +11,8 @@ default:
   ALGOLIA_APP_ID: app_id
   ALGOLIA_API_KEY: api_key
   ALGOLIA_INDEX_NAME: index_name
+  # Airtable
+  AIRTABLE_ACCESS_TOKEN: token
   # Azure
   AZURE_CLIENT_ID: "{{ env_var('AZURE_CLIENT_ID') }}"
   AZURE_CLIENT_SECRET: "{{ env_var('AZURE_CLIENT_SECRET') }}"

mage_ai/io/airtable.py

@@ -0,0 +1,104 @@
+from typing import Union
+
+import pandas as pd
+import polars as pl
+from pyairtable.api import Api
+
+from mage_ai.io.base import BaseIO
+from mage_ai.io.config import BaseConfigLoader, ConfigKey
+
+
+class Airtable(BaseIO):
+    """
+    Handles data transfer between Airtable tables and the Mage app.
+    """
+
+    def __init__(
+            self,
+            token: str,
+            verbose: bool = True,
+            **kwargs) -> None:
+        """
+        Initializes a connection to Airtable.
+
+        Args:
+            token (str): Airtable API token.
+            verbose (bool, optional): Whether to print verbose output. Defaults to True.
+            **kwargs: Additional keyword arguments to pass to the pyairtable Api constructor.
+        """
+        super().__init__(verbose=verbose)
+        self.client = Api(token, **kwargs)  # Create the Airtable API client
+
+    @classmethod
+    def with_config(
+            cls,
+            config: BaseConfigLoader
+    ) -> 'Airtable':
+        """
+        Initializes an Airtable client from a configuration loader.
+
+        Args:
+            config (BaseConfigLoader): Configuration loader object.
+
+        Raises:
+            ValueError: If no valid Airtable API token is found in the configuration.
+
+        Returns:
+            Airtable: An instance of the Airtable class.
+        """
+        if ConfigKey.AIRTABLE_ACCESS_TOKEN not in config:
+            raise ValueError(
+                'No valid API token found for Airtable.'
+                'You must specify your access token in your config.'
+            )
+
+        return cls(
+            token=config[ConfigKey.AIRTABLE_ACCESS_TOKEN]
+        )
+
+    def load(
+            self,
+            base_id: str,
+            table_name: str,
+            **kwargs,
+    ) -> pd.DataFrame:
+        """
+        Loads data from an Airtable table into a Pandas DataFrame.
+
+        Args:
+            base_id (str): The ID of the Airtable base (e.g., 'app*****').
+            table_name (str): The name or ID of the Airtable table (e.g., 'tbl*****').
+            **kwargs: Additional keyword arguments to pass to the pyairtable table.all() method.
+
+        Returns:
+            DataFrame: A Pandas DataFrame containing the data from the Airtable table.
+        """
+        with self.printer.print_msg(
+                f'Loading data frame from table \'{table_name}\' at airtable app \'{base_id}\''
+        ):
+            table = self.client.table(base_id, table_name)  # Get the Airtable table
+            data = table.all(**kwargs)  # Fetch all records from the table
+
+            # Flatten the Airtable data structure into a list of dictionaries
+            flattened_data = []
+            for record in data:
+                flattened_record = {
+                    'id': record['id'],
+                    'createdTime': record['createdTime']
+                }
+                fields = record['fields']
+                flattened_record.update(fields)
+                flattened_data.append(flattened_record)
+
+            return pd.DataFrame(flattened_data)  # Create and return a DataFrame
+
+    def export(
+            self,
+            df: Union[pd.DataFrame, pl.DataFrame],
+            *args,
+            **kwargs
+    ) -> None:
+        """
+        Not implemented yet. This method is intended to export data to Airtable.
+        """
+        pass

mage_ai/io/base.py

@@ -16,6 +16,7 @@ QUERY_ROW_LIMIT = 10_000_000
 
 
 class DataSource(StrEnum):
+    AIRTABLE = 'airtable'
     ALGOLIA = 'algolia'
     API = 'api'
     BIGQUERY = 'bigquery'
@@ -47,6 +48,26 @@ class FileFormat(StrEnum):
     PARQUET = 'parquet'
     HDF5 = 'hdf5'
     XML = 'xml'
+    EXCEL = 'excel'
+
+    @classmethod
+    def from_extension(cls, ext: str):
+        if ext == 'csv':
+            return cls.CSV
+        elif ext == 'json':
+            return cls.JSON
+        elif ext == 'parquet':
+            return cls.PARQUET
+        elif ext == 'hdf5':
+            return cls.HDF5
+        elif ext == 'xml':
+            return cls.XML
+        elif ext in ('xls', 'xlsx'):
+            return cls.EXCEL
+        else:
+            raise ValueError(
+                f'None file format found for this file extension: {ext}'
+            )
 
 
 class ExportWritePolicy(BaseEnum):
@@ -126,7 +147,9 @@ class BaseFile(BaseIO):
         pass
 
     def _get_file_format(self, filepath: Union[os.PathLike, str]) -> str:
-        return os.path.splitext(os.path.basename(filepath))[-1][1:]
+        return FileFormat.from_extension(
+            os.path.splitext(os.path.basename(filepath))[-1][1:]
+        )
 
     def __get_reader(self, format: Union[FileFormat, str, None]) -> Callable:
         """
@@ -151,6 +174,8 @@ class BaseFile(BaseIO):
             return pd.read_hdf
         elif format == FileFormat.XML:
             return pd.read_xml
+        elif format == FileFormat.EXCEL:
+            return pd.read_excel
         else:
             raise ValueError(f"Invalid format '{format}' specified.")
 
@@ -249,6 +274,8 @@ class BaseFile(BaseIO):
                 return df.write_hdf5
             elif format == FileFormat.XML:
                 return df.write_xml
+            elif format == FileFormat.EXCEL:
+                return df.write_excel
             return df.write_parquet
 
         elif isinstance(df, DataFrame):  # pandas DataFrame
@@ -260,6 +287,8 @@ class BaseFile(BaseIO):
                 return df.to_hdf
             elif format == FileFormat.XML:
                 return df.to_xml
+            elif format == FileFormat.EXCEL:
+                return df.to_excel
             return df.to_parquet
 
     def __del__(self):

mage_ai/io/bigquery.py

@@ -1,7 +1,9 @@
+import traceback
 import uuid
 from typing import Dict, List, Mapping, Union
 
 import numpy as np
+import pandas as pd
 from google.cloud.bigquery import (
     Client,
     LoadJobConfig,
@@ -399,7 +401,41 @@ WHERE table_id = '{table_name}'
             # Clean column names
             if type(df) is DataFrame:
                 df.columns = df.columns.str.replace(' ', '_')
+            table = None
+            try:
+                # Cast column types
+                table = self.client.get_table(table_id)
+            except Exception:
+                print(f'Table {table_id} does not exist.')
+                pass
 
+            if table is not None:
+                try:
+                    timestamp_columns = [field.name for field in table.schema
+                                         if field.field_type == 'TIMESTAMP']
+
+                    # Convert TIMESTAMP columns in DataFrame
+                    for col in timestamp_columns:
+                        if col in df.columns:
+                            df[col] = pd.to_datetime(df[col])
+                except Exception:
+                    print('Fail to cast column types in dataframe.')
+                    traceback.print_exc()
+                if (
+                    not config.schema and
+                    config.write_disposition == WriteDisposition.WRITE_TRUNCATE
+                ):
+                    df_columns = df.columns.tolist()
+                    config.schema = [
+                        SchemaField(
+                            field.name,
+                            field.field_type,
+                            mode=field.mode,
+                            fields=field.fields,
+                        )
+                        for field in table.schema
+                        if field.name in df_columns
+                    ]
             return self.client.load_table_from_dataframe(df, table_id, job_config=config).result()
 
     def execute(self, query_string: str, **kwargs) -> None: