mage-ai 0.9.13__py3-none-any.whl → 0.9.15__py3-none-any.whl

mage_ai/ai/generator.py

@@ -33,7 +33,7 @@ class Generator:
         elif use_case == LLMUseCase.GENERATE_PIPELINE_WITH_DESCRIPTION:
             from mage_ai.ai.llm_pipeline_wizard import LLMPipelineWizard
 
-            return await LLMPipelineWizard().async_generate_pipeline_with_description(
+            return await LLMPipelineWizard().async_generate_pipeline_from_description(
                 request.get('pipeline_description'),
             )
         elif use_case == LLMUseCase.GENERATE_COMMENT_FOR_CODE:

mage_ai/ai/generator_cmds.py

@@ -49,7 +49,7 @@ def generate_block_with_description(block_description: str):
 
 @app.command()
 def generate_pipeline_with_description(pipeline_description: str):
-    print(asyncio.run(LLMPipelineWizard().async_generate_pipeline_with_description(
+    print(asyncio.run(LLMPipelineWizard().async_generate_pipeline_from_description(
                                             pipeline_description)))
 
 

mage_ai/ai/llm_pipeline_wizard.py

@@ -2,6 +2,7 @@ import ast
 import asyncio
 import json
 import os
+import re
 from typing import Dict
 
 import openai
@@ -53,23 +54,24 @@ PROMPT_TO_SPLIT_BLOCKS = """
 A BLOCK does one action either reading data from one data source, transforming the data from
 one format to another or exporting data into a data source.
 Based on the code description delimited by triple backticks, your task is to identify
-how many BLOCKS required and function for each BLOCK.
+how many BLOCKS required, function for each BLOCK and upstream blocks between BLOCKs.
 
 Use the following format:
-BLOCK 1: <block function>
-BLOCK 2: <block function>
-BLOCK 3: <block function>
+BLOCK 1: function: <block function>. upstream: <upstream blocks>
+BLOCK 2: function: <block function>. upstream: <upstream blocks>
+BLOCK 3: function: <block function>. upstream: <upstream blocks>
 ...
 
 Example:
 <code description>: ```
-Read data from MySQL, filter out rows with book_price > 100, and save data to BigQuery.
+Read data from MySQL and Postgres, filter out rows with book_price > 100, and save data to BigQuery.
 ```
 
 Answer:
-BLOCK 1: load data from MySQL
-BLOCK 2: filter out rows with book_price > 100
-BLOCK 3: export data to BigQuery
+BLOCK 1: function: load data from MySQL. upstream:
+BLOCK 2: function: load data from Postgres. upstream:
+BLOCK 3: function: filter out rows with book_price > 100. upstream: 1, 2
+BLOCK 4: function: export data to BigQuery. upstream: 3
 
 <code description>: ```{code_description}```"""
 PROMPT_FOR_FUNCTION_COMMENT = """
@@ -81,6 +83,7 @@ Your task is to write comments for each function inside.
 The comment should follow Google Docstring format.
 Return your response in JSON format with function name as key and the comment as value.
 """
+BLOCK_SPLIT_PATTERN = r"BLOCK\s+(\w+):\s+function:\s+(.*?)\.\s+upstream:\s*(.*?)$"
 TRANSFORMERS_FOLDER = 'transformers'
 CLASSIFICATION_FUNCTION_NAME = "classify_description"
 TEMPLATE_CLASSIFICATION_FUNCTION = [
@@ -197,7 +200,10 @@ class LLMPipelineWizard:
                                     function_args.get(DataSource.__name__))
         return block_type, block_language, pipeline_type, config
 
-    async def async_generate_block_with_description(self, block_description: str) -> dict:
+    async def async_generate_block_with_description(
+            self,
+            block_description: str,
+            upstream_blocks: [str]) -> dict:
         messages = [{"role": "user", "content": block_description}]
         response = await openai.ChatCompletion.acreate(
             model="gpt-3.5-turbo-0613",
@@ -220,6 +226,7 @@ class LLMPipelineWizard:
                     pipeline_type=pipeline_type,
                 ),
                 language=block_language,
+                upstream_blocks=upstream_blocks,
             )
         else:
             logger.error("Failed to interpret the description as a block template.")
@@ -238,11 +245,12 @@ class LLMPipelineWizard:
     async def __async_generate_blocks(self,
                                       block_dict: dict,
                                       block_id: int,
-                                      block_description: str) -> dict:
-        block = await self.async_generate_block_with_description(block_description)
+                                      block_description: str,
+                                      upstream_blocks: [str]) -> dict:
+        block = await self.async_generate_block_with_description(block_description, upstream_blocks)
         block_dict[block_id] = block
 
-    async def async_generate_pipeline_with_description(self, pipeline_description: str) -> dict:
+    async def async_generate_pipeline_from_description(self, pipeline_description: str) -> dict:
         splited_block_descriptions = await self.__async_split_description_by_blocks(
             pipeline_description)
         blocks = {}
@@ -250,11 +258,17 @@ class LLMPipelineWizard:
         for line in splited_block_descriptions.strip().split('\n'):
             if line.startswith("BLOCK") and ":" in line:
                 # Extract the block_id and block_description from the line
-                raw_block_id, block_description = line.split(":", 1)
-                block_id = raw_block_id.strip().split(" ")[-1]
-                block_description = block_description.strip()
-                block_tasks.append(
-                    self.__async_generate_blocks(blocks, block_id, block_description))
+                match = re.search(BLOCK_SPLIT_PATTERN, line)
+                if match:
+                    block_id = match.group(1)
+                    block_description = match.group(2).strip()
+                    upstream_blocks = match.group(3).split(", ")
+                    block_tasks.append(
+                        self.__async_generate_blocks(
+                            blocks,
+                            block_id,
+                            block_description,
+                            upstream_blocks))
         await asyncio.gather(*block_tasks)
         return blocks
 

mage_ai/api/policies/PipelineRunPolicy.py

@@ -47,6 +47,7 @@ PipelineRunPolicy.allow_read(PipelineRunPresenter.default_attributes + [
 PipelineRunPolicy.allow_read(PipelineRunPresenter.default_attributes + [
     'block_runs',
     'block_runs_count',
+    'completed_block_runs_count',
     'pipeline_schedule_name',
     'pipeline_schedule_token',
     'pipeline_schedule_type',

mage_ai/api/presenters/BlockRunPresenter.py

@@ -17,4 +17,7 @@ class BlockRunPresenter(BasePresenter):
     ]
 
     def present(self, **kwargs):
-        return self.model.to_dict()
+        if isinstance(self.model, dict):
+            return self.model
+        else:
+            return self.model.to_dict()

mage_ai/api/presenters/PipelineRunPresenter.py

@@ -27,6 +27,7 @@ class PipelineRunPresenter(BasePresenter):
             additional_attributes = [
                 'block_runs',
                 'block_runs_count',
+                'completed_block_runs_count',
                 'pipeline_schedule_name',
                 'pipeline_schedule_token',
                 'pipeline_schedule_type',
@@ -66,6 +67,7 @@ PipelineRunPresenter.register_format(
     PipelineRunPresenter.default_attributes + [
         'block_runs',
         'block_runs_count',
+        'completed_block_runs_count',
         'pipeline_schedule_name',
         'pipeline_schedule_token',
         'pipeline_schedule_type',

mage_ai/api/resources/BlockRunResource.py

@@ -1,7 +1,12 @@
+from sqlalchemy.orm import aliased
+
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
-from mage_ai.orchestration.db.models.schedules import BlockRun, PipelineRun, PipelineSchedule
 from mage_ai.orchestration.db import safe_db_query
-from sqlalchemy.orm import aliased
+from mage_ai.orchestration.db.models.schedules import (
+    BlockRun,
+    PipelineRun,
+    PipelineSchedule,
+)
 
 
 class BlockRunResource(DatabaseResource):
@@ -24,16 +29,17 @@ class BlockRunResource(DatabaseResource):
                 pipeline_schedule_name,
             ) = tup
 
-            block_run = BlockRun()
-            block_run.block_uuid = block_uuid
-            block_run.completed_at = completed_at
-            block_run.created_at = created_at
-            block_run.id = id
-            block_run.pipeline_run_id = pipeline_run_id
-            block_run.status = status
-            block_run.updated_at = updated_at
-            block_run.pipeline_schedule_id = pipeline_schedule_id
-            block_run.pipeline_schedule_name = pipeline_schedule_name
+            block_run = dict(
+                block_uuid=block_uuid,
+                completed_at=completed_at,
+                created_at=created_at,
+                id=id,
+                pipeline_run_id=pipeline_run_id,
+                status=status,
+                updated_at=updated_at,
+                pipeline_schedule_id=pipeline_schedule_id,
+                pipeline_schedule_name=pipeline_schedule_name,
+            )
 
             block_runs.append(block_run)
 

mage_ai/api/resources/GitBranchResource.py

@@ -265,7 +265,14 @@ class GitBranchResource(GenericResource):
         untracked_files: List[str],
         limit: int = None
     ) -> Dict:
-        arr = modified_files + staged_files + untracked_files
+        arr = []
+
+        if modified_files and isinstance(modified_files, list):
+            arr += modified_files
+        if staged_files and isinstance(staged_files, list):
+            arr += staged_files
+        if untracked_files and isinstance(untracked_files, list):
+            arr += untracked_files
 
         if limit:
             arr = arr[:limit]

mage_ai/api/resources/PipelineResource.py

@@ -1,4 +1,5 @@
 import asyncio
+from typing import Dict, List
 
 from sqlalchemy import or_
 from sqlalchemy.orm import aliased
@@ -312,17 +313,28 @@ class PipelineResource(BaseResource):
                 schedule.update(status=status)
 
         @safe_db_query
-        def cancel_pipeline_runs(status, pipeline_uuid):
-            pipeline_runs = (
-                PipelineRun.
-                query.
-                filter(PipelineRun.pipeline_uuid == pipeline_uuid).
-                filter(PipelineRun.status.in_([
-                    PipelineRun.PipelineRunStatus.INITIAL,
-                    PipelineRun.PipelineRunStatus.RUNNING,
-                ]))
-            )
-            for pipeline_run in pipeline_runs:
+        def cancel_pipeline_runs(
+            pipeline_uuid: str = None,
+            pipeline_runs: List[Dict] = None,
+        ):
+            if pipeline_runs is not None:
+                pipeline_run_ids = [run.get('id') for run in pipeline_runs]
+                pipeline_runs_to_cancel = (
+                    PipelineRun.
+                    query.
+                    filter(PipelineRun.id.in_(pipeline_run_ids))
+                )
+            else:
+                pipeline_runs_to_cancel = (
+                    PipelineRun.
+                    query.
+                    filter(PipelineRun.pipeline_uuid == pipeline_uuid).
+                    filter(PipelineRun.status.in_([
+                        PipelineRun.PipelineRunStatus.INITIAL,
+                        PipelineRun.PipelineRunStatus.RUNNING,
+                    ]))
+                )
+            for pipeline_run in pipeline_runs_to_cancel:
                 PipelineScheduler(pipeline_run).stop()
 
         def retry_pipeline_runs(pipeline_runs):
@@ -334,15 +346,19 @@ class PipelineResource(BaseResource):
 
         def _update_callback(resource):
             if status:
+                pipeline_runs = payload.get('pipeline_runs')
                 if status in [
                     ScheduleStatus.ACTIVE.value,
                     ScheduleStatus.INACTIVE.value,
                 ]:
                     update_schedule_status(status, pipeline_uuid)
                 elif status == PipelineRun.PipelineRunStatus.CANCELLED.value:
-                    cancel_pipeline_runs(status, pipeline_uuid)
-                elif status == 'retry' and payload.get('pipeline_runs'):
-                    retry_pipeline_runs(payload.get('pipeline_runs'))
+                    if pipeline_runs is None:
+                        cancel_pipeline_runs(pipeline_uuid=pipeline_uuid)
+                    else:
+                        cancel_pipeline_runs(pipeline_runs=pipeline_runs)
+                elif status == 'retry' and pipeline_runs:
+                    retry_pipeline_runs(pipeline_runs)
 
         self.on_update_callback = _update_callback
 

mage_ai/api/resources/PipelineRunResource.py

@@ -60,12 +60,15 @@ class PipelineRunResource(DatabaseResource):
                 else:
                     order_by.append((parts[0], 'asc'))
 
+        repo_pipeline_schedule_ids = [s.id for s in PipelineSchedule.repo_query]
+
         results = (
-            PipelineRun.
-            query.
-            options(selectinload(PipelineRun.block_runs)).
-            options(selectinload(PipelineRun.pipeline_schedule)).
-            join(PipelineSchedule, PipelineRun.pipeline_schedule_id == PipelineSchedule.id)
+            PipelineRun
+            .query
+            .filter(PipelineRun.pipeline_schedule_id.in_(repo_pipeline_schedule_ids))
+            .options(selectinload(PipelineRun.block_runs))
+            .options(selectinload(PipelineRun.pipeline_schedule))
+            .join(PipelineSchedule, PipelineRun.pipeline_schedule_id == PipelineSchedule.id)
         )
 
         if global_data_product_uuid is not None:

mage_ai/data_integrations/sources/constants.py

@@ -38,6 +38,7 @@ SOURCES = sorted([
     dict(name='Pipedrive'),
     dict(name='Postmark'),
     dict(name='Salesforce'),
+    dict(name='Sftp'),
     dict(name='Stripe'),
     dict(name='Twitter Ads'),
     dict(name='Zendesk'),

mage_ai/data_preparation/executors/block_executor.py

@@ -252,20 +252,18 @@ class BlockExecutor:
                     )
 
                 result = __execute_with_retry()
-            except Exception as e:
+            except Exception as error:
                 self.logger.exception(
                     f'Failed to execute block {self.block.uuid}',
                     **merge_dict(tags, dict(
-                        block_type=self.block.type,
-                        block_uuid=self.block.uuid,
-                        error=e
+                        error=error,
                     )),
                 )
                 if on_failure is not None:
                     on_failure(
                         self.block_uuid,
                         error=dict(
-                            error=str(e),
+                            error=error,
                             errors=traceback.format_stack(),
                             message=traceback.format_exc(),
                         ),
@@ -285,7 +283,7 @@ class BlockExecutor:
                     logging_tags=tags,
                     pipeline_run=pipeline_run,
                 )
-                raise e
+                raise error
             self.logger.info(f'Finish executing block with {self.__class__.__name__}.', **tags)
             if on_complete is not None:
                 on_complete(self.block_uuid)
@@ -537,7 +535,7 @@ class BlockExecutor:
         tags: Dict = None,
     ):
         """
-        Update the status of block run by edither updating the BlockRun db object or making
+        Update the status of block run by either updating the BlockRun db object or making
         API call
 
         Args:

mage_ai/data_preparation/git/__init__.py

@@ -434,6 +434,24 @@ class Git:
                 for url in remote.urls:
                     if url.lower().startswith('https'):
                         repository_names.append('/'.join(url.split('/')[-2:]).replace('.git', ''))
+
+                        # Remove the token from the URL
+                        # e.g. https://[user]:[token]@[netloc]
+                        parts = url.split('@')
+
+                        parts_arr = []
+                        if len(parts) >= 2:
+                            # https://[user]:[token]
+                            parts2 = parts[0].split(':')
+                            # ['https', '', 'user', 'token']
+                            parts2[len(parts2) - 1] = '[token]'
+                            parts_arr.append(':'.join(parts2))
+                            parts_arr += parts[1:]
+                        else:
+                            parts_arr += parts
+
+                        url = '@'.join(parts_arr)
+
                     urls.append(url)
             except GitCommandError as err:
                 print('WARNING (mage_ai.data_preparation.git.remotes):')

mage_ai/data_preparation/models/block/__init__.py

@@ -17,6 +17,7 @@ import pandas as pd
 import simplejson
 from jinja2 import Template
 
+import mage_ai.data_preparation.decorators
 from mage_ai.cache.block import BlockCache
 from mage_ai.data_cleaner.shared.utils import is_geo_dataframe, is_spark_dataframe
 from mage_ai.data_preparation.logging.logger import DictLogger
@@ -1110,6 +1111,8 @@ class Block:
     ) -> List:
         if logging_tags is None:
             logging_tags = dict()
+        if input_vars is None:
+            input_vars = list()
 
         decorated_functions = []
         test_functions = []
@@ -1177,9 +1180,18 @@ class Block:
                         block_uuid = self.replicated_block
                         block_file_path = self.replicated_block_object.file_path
                     spec = importlib.util.spec_from_file_location(
-                        block_uuid, block_file_path,
+                        block_uuid,
+                        block_file_path,
                     )
                     module = importlib.util.module_from_spec(spec)
+                    # Set the decorators in the module in case they are not defined in the block
+                    # code
+                    setattr(
+                        module,
+                        self.type,
+                        getattr(mage_ai.data_preparation.decorators, self.type),
+                    )
+                    module.test = mage_ai.data_preparation.decorators.test
                     spec.loader.exec_module(module)
                     block_function_updated = getattr(module, block_function.__name__)
                     self.module = module

mage_ai/data_preparation/models/block/dbt/utils/__init__.py

@@ -501,9 +501,12 @@ def config_file_loader_and_configuration(
         database = kwargs.get('database') or profile.get('project')
         schema = profile.get('dataset')
 
-        config_file_loader = ConfigFileLoader(config=dict(
+        config_file_loader_kwargs = dict(
             GOOGLE_SERVICE_ACC_KEY_FILEPATH=keyfile,
-        ))
+        )
+        if profile.get('location'):
+            config_file_loader_kwargs['GOOGLE_LOCATION'] = profile.get('location')
+        config_file_loader = ConfigFileLoader(config=config_file_loader_kwargs)
         configuration = dict(
             data_provider=profile_type,
             data_provider_database=database,

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

@@ -122,7 +122,7 @@ class IntegrationBlock(Block):
                 proc = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
 
                 for line in proc.stdout:
-                    f.write(filter_out_config_values(line.decode(), config)),
+                    f.write(line.decode()),
                     print_log_from_line(
                         line,
                         config=config,

mage_ai/data_preparation/models/custom_templates/utils.py

@@ -56,13 +56,23 @@ def group_and_hydrate_files(
     file_dicts: List[Dict],
     custom_template_class,
 ) -> List:
-    groups = group_by(lambda x: os.path.join(*x.get('parent_names', [])), file_dicts)
+    def _func(x):
+        arr = ['']
 
-    arr = []
+        if x:
+            parent_names = x.get('parent_names', []) or []
+            if parent_names and len(parent_names) >= 1:
+                arr = [str(parent_name) for parent_name in parent_names]
+
+        return os.path.join(*arr)
+
+    groups = group_by(_func, file_dicts)
+
+    custom_templates = []
 
     for template_uuid, _ in groups.items():
         custom_template = custom_template_class.load(template_uuid=template_uuid)
         if custom_template:
-            arr.append(custom_template)
+            custom_templates.append(custom_template)
 
-    return arr
+    return custom_templates

mage_ai/data_preparation/models/variable.py

@@ -430,11 +430,14 @@ class Variable:
         df_output = data.copy()
         # Clean up data types since parquet doesn't support mixed data types
         for c in df_output.columns:
-            c_dtype = df_output[c].dtype
+            df_col = df_output[c]
+            if type(df_col) is pd.DataFrame:
+                raise Exception(f'Please do not use duplicate column name: "{c}"')
+            c_dtype = df_col.dtype
             if not is_object_dtype(c_dtype):
                 column_types[c] = str(c_dtype)
             else:
-                series_non_null = df_output[c].dropna()
+                series_non_null = df_col.dropna()
                 if len(series_non_null) > 0:
                     coltype = type(series_non_null.iloc[0])
                     if is_object_dtype(series_non_null.dtype):

mage_ai/data_preparation/templates/constants.py

@@ -51,6 +51,14 @@ TEMPLATES = [
         name='MongoDB',
         path='data_loaders/mongodb.py',
     ),
+    dict(
+        block_type=BlockType.DATA_LOADER,
+        description='Load data from MSSQL.',
+        groups=[GROUP_DATABASES],
+        language=BlockLanguage.PYTHON,
+        name='MSSQL',
+        path='data_loaders/mssql.py',
+    ),
     dict(
         block_type=BlockType.DATA_EXPORTER,
         description='Export data to a Delta Table in Amazon S3.',
@@ -82,6 +90,14 @@ TEMPLATES = [
         name='MongoDB',
         path='data_exporters/mongodb.py',
     ),
+    dict(
+        block_type=BlockType.DATA_EXPORTER,
+        description='Export data to MSSQL.',
+        groups=[GROUP_DATABASES],
+        language=BlockLanguage.PYTHON,
+        name='MSSQL',
+        path='data_exporters/mssql.py',
+    ),
     dict(
         block_type=BlockType.DATA_LOADER,
         description='Trigger another pipeline to run.',

mage_ai/data_preparation/templates/data_exporters/mssql.py

@@ -0,0 +1,32 @@
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.io.config import ConfigFileLoader
+from mage_ai.io.mssql import MSSQL
+from pandas import DataFrame
+from os import path
+
+if 'data_exporter' not in globals():
+    from mage_ai.data_preparation.decorators import data_exporter
+
+
+@data_exporter
+def export_data_to_mssql(df: DataFrame, **kwargs) -> None:
+    """
+    Template for exporting data to a MSSQL database.
+    Specify your configuration settings in 'io_config.yaml'.
+    Set the following in your io_config:
+
+    Docs: https://docs.mage.ai/integrations/databases/MicrosoftSQLServer
+    """
+    schema_name = 'dbo'  # Specify the name of the schema to export data to
+    table_name = 'your_table_name'  # Specify the name of the table to export data to
+    config_path = path.join(get_repo_path(), 'io_config.yaml')
+    config_profile = 'default'
+
+    with MSSQL.with_config(ConfigFileLoader(config_path, config_profile)) as loader:
+        loader.export(
+            df,
+            schema_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_loaders/mssql.py

@@ -0,0 +1,27 @@
+{% 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.mssql import MSSQL
+from os import path
+{{ super() -}}
+{% endblock %}
+
+
+{% block content %}
+@data_loader
+def load_data_from_mssql(*args, **kwargs):
+    """
+    Template for loading data from a MSSQL database.
+    Specify your configuration settings in 'io_config.yaml'.
+    Set the following in your io_config:
+
+    Docs: https://docs.mage.ai/integrations/databases/MicrosoftSQLServer
+    """
+    query = 'Your MSSQL query'  # Specify your SQL query here
+    config_path = path.join(get_repo_path(), 'io_config.yaml')
+    config_profile = 'default'
+
+    with MSSQL.with_config(ConfigFileLoader(config_path, config_profile)) as loader:
+        return loader.load(query)
+{% endblock %}

mage_ai/io/base.py

@@ -35,6 +35,7 @@ class FileFormat(str, Enum):
     JSON = 'json'
     PARQUET = 'parquet'
     HDF5 = 'hdf5'
+    XML = 'xml'
 
 
 class ExportWritePolicy(str, Enum):
@@ -137,6 +138,8 @@ class BaseFile(BaseIO):
             return pd.read_parquet
         elif format == FileFormat.HDF5:
             return pd.read_hdf
+        elif format == FileFormat.XML:
+            return pd.read_xml
         else:
             raise ValueError(f'Invalid format \'{format}\' specified.')
 
@@ -226,6 +229,8 @@ class BaseFile(BaseIO):
             return df.to_json
         elif format == FileFormat.HDF5:
             return df.to_hdf
+        elif format == FileFormat.XML:
+            return df.to_xml
         return df.to_parquet
 
     def __del__(self):

mage_ai/io/mssql.py

@@ -142,6 +142,7 @@ class MSSQL(BaseSQL):
             # Remove extraneous surrounding double quotes
             # that get added while performing conversion to string.
             df_[col] = df_[col].apply(lambda x: x.strip('"') if x and isinstance(x, str) else x)
+        df_.replace({np.NaN: None}, inplace=True)
         for _, row in df_.iterrows():
             values.append(tuple(row))