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

mage_ai/ai/llm_pipeline_wizard.py

@@ -156,16 +156,18 @@ TEMPLATE_CLASSIFICATION_FUNCTION = [
                     "type": "string",
                     "description": "Programming language of the code block. "
                                    f"Default value is {BlockLanguage.__name__}__python.",
-                    "enum": [f"{BlockLanguage.__name__}__{type.name.lower()}"
-                             for type in BlockLanguage]
+                    "enum": [
+                        f"{BlockLanguage.__name__}__{type.name.lower()}"
+                        for type in BlockLanguage]
                 },
                 PipelineType.__name__: {
                     "type": "string",
                     "description": "Type of pipeline to build. Default value is "
                                    f"{PipelineType.__name__}__python if pipeline type "
                                    "is not mentioned in the description.",
-                    "enum": [f"{PipelineType.__name__}__{type.name.lower()}"
-                             for type in PipelineType]
+                    "enum": [
+                        f"{PipelineType.__name__}__{type.name.lower()}"
+                        for type in PipelineType]
                 },
                 ActionType.__name__: {
                     "type": "string",

mage_ai/ai/openai_client.py

@@ -42,16 +42,18 @@ tools = [
                         "type": "string",
                         "description": "Programming language of the code block. "
                                        f"Default value is {BlockLanguage.__name__}__python.",
-                        "enum": [f"{BlockLanguage.__name__}__{type.name.lower()}"
-                                 for type in BlockLanguage]
+                        "enum": [
+                            f"{BlockLanguage.__name__}__{type.name.lower()}"
+                            for type in BlockLanguage]
                     },
                     PipelineType.__name__: {
                         "type": "string",
                         "description": "Type of pipeline to build. Default value is "
                                        f"{PipelineType.__name__}__python if pipeline type "
                                        "is not mentioned in the description.",
-                        "enum": [f"{PipelineType.__name__}__{type.name.lower()}"
-                                 for type in PipelineType]
+                        "enum": [
+                            f"{PipelineType.__name__}__{type.name.lower()}"
+                            for type in PipelineType]
                     },
                     ActionType.__name__: {
                         "type": "string",
@@ -85,7 +87,7 @@ class OpenAIClient(AIClient):
             open_ai_config.openai_api_key or os.getenv('OPENAI_API_KEY')
         openai.api_key = openai_api_key
         self.llm = OpenAI(openai_api_key=openai_api_key, temperature=0)
-        self.openai_client = OpenAILib()
+        self.openai_client = OpenAILib(api_key=openai_api_key)
 
     def __chat_completion_request(self, messages):
         try:

mage_ai/api/policies/PipelineSchedulePolicy.py

@@ -100,6 +100,7 @@ PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes +
 PipelineSchedulePolicy.allow_read(PipelineSchedulePresenter.default_attributes + [
     'event_matchers',
     'next_pipeline_run_date',
+    'rotate_token',
     'tags',
 ], scopes=[
     OauthScope.CLIENT_PRIVATE,

mage_ai/api/presenters/PipelineSchedulePresenter.py

@@ -1,5 +1,6 @@
 from mage_ai.api.operations import constants
 from mage_ai.api.presenters.BasePresenter import BasePresenter
+from mage_ai.settings.server import HIDE_API_TRIGGER_TOKEN
 
 
 class PipelineSchedulePresenter(BasePresenter):
@@ -32,8 +33,6 @@ class PipelineSchedulePresenter(BasePresenter):
             ])
             data['tags'] = sorted([tag.name for tag in self.get_tag_associations])
             data['next_pipeline_run_date'] = self.model.next_execution_date()
-
-            return data
         elif 'with_runtime_average' == display_format:
             data = self.model.to_dict()
             data['runtime_average'] = self.model.runtime_average()
@@ -47,6 +46,16 @@ class PipelineSchedulePresenter(BasePresenter):
         else:
             data = self.model.to_dict()
 
+        if display_format == constants.UPDATE:
+            rotate_token = kwargs.get(
+                'payload', dict(),
+            ).get(
+                'pipeline_schedule', dict(),
+            ).get('rotate_token')
+        else:
+            rotate_token = False
+        if HIDE_API_TRIGGER_TOKEN and not rotate_token:
+            data['token'] = '[API_TOKEN_PLACEHOLDER]'
         return data
 
 

mage_ai/api/resources/PipelineScheduleResource.py

@@ -29,7 +29,8 @@ from mage_ai.orchestration.db.models.tags import (
 )
 from mage_ai.settings.platform import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
-from mage_ai.shared.hash import merge_dict
+from mage_ai.settings.server import HIDE_API_TRIGGER_TOKEN
+from mage_ai.shared.hash import ignore_keys, merge_dict
 
 
 class PipelineScheduleResource(DatabaseResource):
@@ -282,6 +283,7 @@ class PipelineScheduleResource(DatabaseResource):
 
     @safe_db_query
     def update(self, payload, **kwargs):
+        # Update associated event matchers
         arr = payload.pop('event_matchers', None)
         event_matchers = []
         if arr is not None:
@@ -320,6 +322,7 @@ class PipelineScheduleResource(DatabaseResource):
                 ]
                 em.update(pipeline_schedules=ps)
 
+        # Update associated tags
         tag_names = payload.pop('tags', None)
         if tag_names is not None:
             # 1. Fetch all tag associations
@@ -408,7 +411,11 @@ class PipelineScheduleResource(DatabaseResource):
 
         old_name = self.model.name
 
-        resource = super().update(payload)
+        # Rotate token
+        if payload.get('rotate_token'):
+            payload['token'] = uuid.uuid4().hex
+
+        resource = super().update(ignore_keys(payload, ['rotate_token']))
         updated_model = resource.model
 
         repo_path = get_repo_path(user=self.current_user)
@@ -424,9 +431,10 @@ class PipelineScheduleResource(DatabaseResource):
                 sla=updated_model.sla,
                 start_time=updated_model.start_time,
                 status=updated_model.status,
-                token=updated_model.token,
                 variables=updated_model.variables,
             )
+            if not HIDE_API_TRIGGER_TOKEN:
+                trigger.token = updated_model.token
 
             update_only_if_exists = (
                 not pipeline.should_save_trigger_in_code_automatically()

mage_ai/api/resources/PipelineTriggerResource.py

@@ -8,6 +8,7 @@ from mage_ai.data_preparation.models.triggers import (
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.schedules import PipelineSchedule
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.settings.server import HIDE_API_TRIGGER_TOKEN
 
 
 class PipelineTriggerResource(GenericResource):
@@ -50,9 +51,10 @@ class PipelineTriggerResource(GenericResource):
                 sla=pipeline_schedule.sla,
                 start_time=pipeline_schedule.start_time,
                 status=pipeline_schedule.status,
-                token=pipeline_schedule.token,
                 variables=pipeline_schedule.variables,
             )
+            if not HIDE_API_TRIGGER_TOKEN:
+                trigger.token = pipeline_schedule.token
         else:
             trigger = Trigger(**payload)
 

mage_ai/api/resources/SessionResource.py

@@ -6,7 +6,7 @@ from mage_ai.authentication.ldap import new_ldap_connection
 from mage_ai.authentication.oauth2 import encode_token, generate_access_token
 from mage_ai.authentication.passwords import verify_password
 from mage_ai.authentication.providers.constants import NAME_TO_PROVIDER
-from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db import safe_db_query, safe_db_query_async
 from mage_ai.orchestration.db.models.oauth import Role, User
 from mage_ai.settings import (
     AUTHENTICATION_MODE,
@@ -20,7 +20,7 @@ from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 class SessionResource(BaseResource):
     @classmethod
-    @safe_db_query
+    @safe_db_query_async
     async def create(self, payload, _, **kwargs):
         email = payload.get('email')
         password = payload.get('password')

mage_ai/api/resources/SyncResource.py

@@ -163,7 +163,7 @@ class SyncResource(GenericResource):
         access_token = user_payload.pop('access_token', None)
         if access_token:
             secret_name = get_access_token_secret_name(user=user)
-            secret = Secret.query.filter(Secret.name == secret_name).one_or_none()
+            secret = Secret.repo_query.filter(Secret.name == secret_name).one_or_none()
             if secret:
                 secret.delete()
             create_secret(secret_name, access_token, repo_name=repo_name)

mage_ai/api/resources/UserResource.py

@@ -28,7 +28,7 @@ class UserResource(DatabaseResource):
         results = (
             User.
             query.
-            order_by(User.username.asc())
+            order_by(User.id.asc())
         )
 
         if user and user.is_admin:

mage_ai/data_integrations/destinations/constants.py

@@ -1,9 +1,11 @@
 DESTINATIONS = [
+    dict(name='Airtable'),
     dict(name='Amazon S3'),
     dict(name='BigQuery'),
     dict(name='Clickhouse'),
     dict(name='Delta Lake Azure'),
     dict(name='Delta Lake S3'),
+    dict(name='Doris'),
     dict(name='Elasticsearch'),
     dict(name='Google Cloud Storage'),
     dict(name='Kafka'),
@@ -20,5 +22,6 @@ DESTINATIONS = [
     dict(name='Redshift'),
     dict(name='Salesforce'),
     dict(name='Snowflake'),
+    dict(name='Teradata'),
     dict(name='Trino'),
 ]

mage_ai/data_integrations/sources/constants.py

@@ -3,6 +3,7 @@ from mage_ai.shared.hash import index_by
 
 SQL_SOURCES = [
     dict(name='BigQuery'),
+    dict(name='Doris'),
     dict(
         name='Microsoft SQL Server',
         uuid='mssql',
@@ -53,6 +54,7 @@ SOURCES = sorted([
     dict(name='Sftp'),
     dict(name='Stripe'),
     dict(name='Tableau'),
+    dict(name='Teradata'),
     dict(name='Twitter Ads'),
     dict(name='Zendesk'),
 ] + SQL_SOURCES, key=lambda x: x['name'])

mage_ai/data_preparation/models/block/__init__.py

@@ -878,7 +878,7 @@ class Block(
         relative_path: bool = False,
     ) -> str:
         file_extension = BLOCK_LANGUAGE_TO_FILE_EXTENSION[language]
-        block_directory = block_type + 's' if block_type != BlockType.CUSTOM else block_type
+        block_directory = f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
 
         parts = []
         if not relative_path:
@@ -1010,7 +1010,7 @@ class Block(
         self,
         content: str,
     ) -> str:
-        return '@' + self.type + '\n' + content
+        return f'@{self.type}\n{content}'
 
     @classmethod
     def after_create(cls, block: 'Block', **kwargs) -> None:
@@ -1199,7 +1199,7 @@ class Block(
 
     @classmethod
     def file_directory_name(self, block_type: BlockType) -> str:
-        return block_type + 's' if block_type != BlockType.CUSTOM else block_type
+        return f'{block_type}s' if block_type != BlockType.CUSTOM else block_type
 
     @classmethod
     def block_type_from_path(
@@ -1218,7 +1218,7 @@ class Block(
         for block_type in BlockType:
             if BlockType.CUSTOM == block_type and dir_name == block_type:
                 return BlockType.CUSTOM
-            elif dir_name == block_type + 's':
+            elif dir_name == f'{block_type}s':
                 return block_type
 
     @classmethod
@@ -3941,7 +3941,6 @@ class Block(
             for v in b.output_variables(execution_partition=execution_partition):
                 objs.append(
                     self.get_variable_object(
-                        self.pipeline_uuid,
                         b.uuid,
                         v,
                         partition=execution_partition,

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/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'
 

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
@@ -2328,6 +2329,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/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') }}"