mage-ai 0.8.4__py3-none-any.whl → 0.8.6__py3-none-any.whl

mage_ai/api/policies/BlockPolicy.py

@@ -53,6 +53,7 @@ BlockPolicy.allow_write([
     'content',
     'converted_from',
     'language',
+    'metadata',
     'name',
     'priority',
     'type',
@@ -73,6 +74,7 @@ BlockPolicy.allow_write([
     'executor_config',
     'executor_type',
     'language',
+    'metadata',
     'name',
     'outputs',
     'status',

mage_ai/api/policies/PipelinePolicy.py

@@ -65,6 +65,7 @@ PipelinePolicy.allow_write([
 ], condition=lambda policy: policy.has_at_least_editor_role())
 
 PipelinePolicy.allow_query([
+    'includes_block_metadata',
     'includes_content',
     'includes_outputs',
 ], scopes=[

mage_ai/api/presenters/BlockPresenter.py

@@ -11,6 +11,7 @@ class BlockPresenter(BasePresenter):
         'executor_config',
         'executor_type',
         'language',
+        'metadata',
         'name',
         'status',
         'type',

mage_ai/api/presenters/OutputPresenter.py

@@ -5,6 +5,7 @@ class OutputPresenter(BasePresenter):
     default_attributes = [
         'sample_data',
         'shape',
+        'text_data',
         'type',
         'variable_uuid',
     ]

mage_ai/api/presenters/PipelinePresenter.py

@@ -21,13 +21,18 @@ class PipelinePresenter(BasePresenter):
             include_content = query.get('includes_content', [True])
             if include_content:
                 include_content = include_content[0]
+
             include_outputs = query.get('includes_outputs', [True])
             if include_outputs:
                 include_outputs = include_outputs[0]
 
+            include_block_metadata = query.get('includes_block_metadata', [True])
+            if include_block_metadata:
+                include_block_metadata = include_block_metadata[0]
+
             return await self.model.to_dict_async(
+                include_block_metadata=include_block_metadata,
                 include_content=include_content,
-                include_outputs=include_outputs,
                 sample_count=DATAFRAME_SAMPLE_COUNT_PREVIEW,
             )
 

mage_ai/data_preparation/models/block/__init__.py

@@ -267,6 +267,9 @@ class Block:
             self._content = await self.file.content_async()
         return self._content
 
+    async def metadata_async(self):
+        return {}
+
     @property
     def executable(self):
         return (
@@ -566,6 +569,7 @@ class Block:
         dynamic_block_index: int = None,
         dynamic_block_uuid: str = None,
         dynamic_upstream_block_uuids: List[str] = None,
+        run_settings: Dict = None,
     ) -> Dict:
         try:
             if not run_all_blocks:
@@ -595,6 +599,7 @@ class Block:
                 runtime_arguments=runtime_arguments,
                 dynamic_block_index=dynamic_block_index,
                 dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                run_settings=run_settings,
             )
             block_output = output['output'] or []
             variable_mapping = dict()
@@ -773,6 +778,7 @@ class Block:
         runtime_arguments: Dict = None,
         dynamic_block_index: int = None,
         dynamic_upstream_block_uuids: List[str] = None,
+        run_settings: Dict = None,
     ) -> Dict:
         # Add pipeline uuid and block uuid to global_vars
         global_vars = merge_dict(
@@ -821,7 +827,8 @@ class Block:
                 test_execution=test_execution,
                 input_from_output=input_from_output,
                 runtime_arguments=runtime_arguments,
-                upstream_block_uuids=upstream_block_uuids
+                upstream_block_uuids=upstream_block_uuids,
+                run_settings=run_settings,
             )
 
         output_message = dict(output=outputs)
@@ -841,6 +848,7 @@ class Block:
         input_from_output: Dict = None,
         runtime_arguments: Dict = None,
         upstream_block_uuids: List[str] = None,
+        run_settings: Dict = None,
     ) -> List:
         decorated_functions = []
         test_functions = []
@@ -1219,15 +1227,17 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
 
     async def to_dict_async(
         self,
-        include_content=False,
-        include_outputs=False,
-        sample_count=None,
+        include_block_metadata: bool = False,
+        include_content: bool = False,
+        include_outputs: bool = False,
+        sample_count: int = None,
         check_if_file_exists: bool = False,
     ):
         data = self.to_dict_base()
 
         if include_content:
             data['content'] = await self.content_async()
+
         if include_outputs:
             data['outputs'] = await self.outputs_async()
             if check_if_file_exists:
@@ -1239,6 +1249,10 @@ df = get_variable('{self.pipeline.uuid}', '{block_uuid}', 'df')
                         'Delete the current block to remove it from the pipeline or write code ' +
                         f'and save the pipeline to create a new file at {file_path}.',
                     )
+
+        if include_block_metadata:
+            data['metadata'] = await self.metadata_async()
+
         return data
 
     def update(self, data):

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

@@ -2,6 +2,8 @@ from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.block.dbt.utils import (
     build_command_line_arguments,
     create_upstream_tables,
+    load_profiles_async,
+    parse_attributes,
     query_from_compiled_sql,
     run_dbt_tests,
     update_model_settings,
@@ -10,7 +12,6 @@ from mage_ai.data_preparation.models.constants import BlockLanguage
 from mage_ai.data_preparation.repo_manager import get_repo_path
 from mage_ai.shared.hash import merge_dict
 from typing import Any, Dict, List
-
 import json
 import os
 import subprocess
@@ -31,6 +32,49 @@ class DBTBlock(Block):
 
         return super().file_path
 
+    async def metadata_async(self) -> Dict:
+        project = None
+        projects = {}
+
+        if self.configuration.get('file_path'):
+            attributes_dict = parse_attributes(self)
+            profiles_full_path = attributes_dict['profiles_full_path']
+            project_name = attributes_dict['project_name']
+            project = project_name
+
+            targets = []
+            profiles = await load_profiles_async(project_name, profiles_full_path)
+            outputs = profiles.get('outputs')
+            if outputs:
+                targets += sorted(list(outputs.keys()))
+
+            projects[project_name] = dict(
+                target=profiles.get('target'),
+                targets=targets,
+            )
+        else:
+            dbt_dir = f'{get_repo_path()}/dbt'
+            project_names = [
+                name for name in os.listdir(dbt_dir) if os.path.isdir(f'{dbt_dir}/{name}')
+            ]
+            for project_name in project_names:
+                profiles_full_path = f'{dbt_dir}/{project_name}/profiles.yml'
+                targets = []
+                profiles = await load_profiles_async(project_name, profiles_full_path)
+                outputs = profiles.get('outputs')
+                if outputs:
+                    targets += sorted(list(outputs.keys()))
+
+                projects[project_name] = dict(
+                    target=profiles.get('target'),
+                    targets=targets,
+                )
+
+        return dict(dbt=dict(
+            project=project,
+            projects=projects,
+        ))
+
     def run_tests(
         self,
         build_block_output_stdout=None,
@@ -56,10 +100,14 @@ class DBTBlock(Block):
         global_vars: Dict = None,
         test_execution: bool = False,
         runtime_arguments: Dict = None,
+        run_settings: Dict = None,
         **kwargs,
     ) -> List:
         variables = merge_dict(global_vars, runtime_arguments or {})
 
+        if run_settings and run_settings.get('run_model', False):
+            test_execution = False
+
         dbt_command, args, command_line_dict = build_command_line_arguments(
             self,
             variables,
@@ -95,6 +143,7 @@ class DBTBlock(Block):
             df = query_from_compiled_sql(
                 self,
                 dbt_profile_target,
+                limit=self.configuration.get('limit'),
             )
             self.store_variables(
                 dict(df=df),

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

@@ -24,6 +24,7 @@ from mage_ai.shared.parsers import encode_complex
 from mage_ai.shared.utils import clean_name, files_in_path
 from pandas import DataFrame
 from typing import Callable, Dict, List, Tuple
+import aiofiles
 import os
 import re
 import simplejson
@@ -252,11 +253,65 @@ def add_table_to_source(block: 'Block', settings: Dict, source_name: str, table_
     return settings
 
 
-def get_profile(block, profile_target: str = None) -> Dict:
-    attr = parse_attributes(block)
-    project_name = attr['project_name']
-    profiles_full_path = attr['profiles_full_path']
-    return load_profile(project_name, profiles_full_path, profile_target)
+def load_profiles_file(profiles_full_path: str) -> Dict:
+    try:
+        with open(profiles_full_path, 'r') as f:
+            try:
+                text = Template(f.read()).render(
+                    **get_template_vars(),
+                )
+                return yaml.safe_load(text)
+            except Exception as err:
+                print(
+                    f'Error loading file {profiles_full_path}, check file content syntax: {err}.',
+                )
+                return {}
+    except OSError as err:
+        print(
+            f'Error loading file {profiles_full_path}, check file content syntax: {err}.',
+        )
+        return {}
+
+
+async def load_profiles_file_async(profiles_full_path: str) -> Dict:
+    try:
+        async with aiofiles.open(profiles_full_path, mode='r') as fp:
+            try:
+                file_content = await fp.read()
+                text = Template(file_content).render(
+                    **get_template_vars(),
+                )
+                return yaml.safe_load(text)
+            except Exception as err:
+                print(
+                    f'Error loading file {profiles_full_path}, check file content syntax: {err}.',
+                )
+                return {}
+    except OSError as err:
+        print(
+            f'Error loading file {profiles_full_path}, check file content syntax: {err}.',
+        )
+        return {}
+
+
+def load_profiles(project_name: str, profiles_full_path: str) -> Dict:
+    profiles = load_profiles_file(profiles_full_path)
+
+    if not profiles or project_name not in profiles:
+        print(f'Project name {project_name} does not exist in profile file {profiles_full_path}.')
+        return {}
+
+    return profiles[project_name]
+
+
+async def load_profiles_async(project_name: str, profiles_full_path: str) -> Dict:
+    profiles = await load_profiles_file_async(profiles_full_path)
+
+    if not profiles or project_name not in profiles:
+        print(f'Project name {project_name} does not exist in profile file {profiles_full_path}.')
+        return {}
+
+    return profiles[project_name]
 
 
 def load_profile(
@@ -264,19 +319,19 @@ def load_profile(
     profiles_full_path: str,
     profile_target: str = None,
 ) -> Dict:
-    with open(profiles_full_path, 'r') as f:
-        try:
-            text = Template(f.read()).render(
-                **get_template_vars(),
-            )
-            profile = yaml.safe_load(text)[project_name]
-            outputs = profile['outputs']
-            target = profile.get('target')
 
-            return outputs.get(profile_target or target)
-        except Exception as err:
-            print(f'Error loading file {profiles_full_path}, please check file content '
-                  f'syntax: {err}.')
+    profile = load_profiles(project_name, profiles_full_path)
+    outputs = profile.get('outputs', {})
+    target = profile.get('target', None)
+
+    return outputs.get(profile_target or target)
+
+
+def get_profile(block, profile_target: str = None) -> Dict:
+    attr = parse_attributes(block)
+    project_name = attr['project_name']
+    profiles_full_path = attr['profiles_full_path']
+    return load_profile(project_name, profiles_full_path, profile_target)
 
 
 def config_file_loader_and_configuration(block, profile_target: str) -> Dict:
@@ -565,7 +620,50 @@ def interpolate_input(
     return query
 
 
-def query_from_compiled_sql(block, profile_target: str) -> DataFrame:
+def interpolate_refs_with_table_names(
+    query_string: str,
+    block: Block,
+    profile_target: str,
+    configuration: Dict,
+):
+    profile = get_profile(block, profile_target)
+
+    profile_type = profile.get('type')
+    quote_str = ''
+    if DataSource.POSTGRES == profile_type:
+        database = profile['dbname']
+        schema = profile['schema']
+        quote_str = '"'
+    elif DataSource.MYSQL == profile_type:
+        database = configuration['data_provider_database']
+        schema = None
+        quote_str = '`'
+    elif DataSource.BIGQUERY == profile_type:
+        database = profile['project']
+        schema = profile['dataset']
+        quote_str = '`'
+    elif DataSource.REDSHIFT == profile_type:
+        database = profile['dbname']
+        schema = profile['schema']
+        quote_str = '"'
+    elif DataSource.SNOWFLAKE == profile_type:
+        database = profile['database']
+        schema = profile['schema']
+    elif DataSource.TRINO == profile_type:
+        database = profile['catalog']
+        schema = profile['schema']
+
+    return interpolate_input(
+        block,
+        query_string,
+        configuration=configuration,
+        profile_database=database,
+        profile_schema=schema,
+        quote_str=quote_str,
+    )
+
+
+def query_from_compiled_sql(block, profile_target: str, limit: int = None) -> DataFrame:
     attr = parse_attributes(block)
 
     config_file_loader, configuration = config_file_loader_and_configuration(
@@ -577,77 +675,55 @@ def query_from_compiled_sql(block, profile_target: str) -> DataFrame:
     project_full_path = attr['project_full_path']
     file_path = attr['file_path']
 
-    profile = get_profile(block, profile_target)
-
     file = f'{project_full_path}/target/compiled/{file_path}'
 
     with open(file, 'r') as f:
         query_string = f.read()
 
-        profile_type = profile.get('type')
-        quote_str = ''
-        if DataSource.POSTGRES == profile_type:
-            database = profile['dbname']
-            schema = profile['schema']
-            quote_str = '"'
-        elif DataSource.MYSQL == profile_type:
-            database = configuration['data_provider_database']
-            schema = None
-            quote_str = '`'
-        elif DataSource.BIGQUERY == profile_type:
-            database = profile['project']
-            schema = profile['dataset']
-            quote_str = '`'
-        elif DataSource.REDSHIFT == profile_type:
-            database = profile['dbname']
-            schema = profile['schema']
-            quote_str = '"'
-        elif DataSource.SNOWFLAKE == profile_type:
-            database = profile['database']
-            schema = profile['schema']
-        elif DataSource.TRINO == profile_type:
-            database = profile['catalog']
-            schema = profile['schema']
-
-        query_string = interpolate_input(
-            block,
-            query_string,
-            configuration=configuration,
-            profile_database=database,
-            profile_schema=schema,
-            quote_str=quote_str,
-        )
+        # TODO (tommy dang): this was needed because we didn’t want to create model tables and
+        # so we’d create a table to store the model results without creating the model.
+        # However, we’re requiring people to run the model and create the model table to use ref.
+        # query_string = interpolate_refs_with_table_names(
+        #     query_string,
+        #     block,
+        #     profile_target=profile_target,
+        #     configuration=configuration,
+        # )
+
+        shared_kwargs = {}
+        if limit is not None:
+            shared_kwargs['limit'] = limit
 
         if DataSource.POSTGRES == data_provider:
             from mage_ai.io.postgres import Postgres
 
             with Postgres.with_config(config_file_loader) as loader:
-                return loader.load(query_string)
+                return loader.load(query_string, **shared_kwargs)
         elif DataSource.MYSQL == data_provider:
             from mage_ai.io.mysql import MySQL
 
             with MySQL.with_config(config_file_loader) as loader:
-                return loader.load(query_string)
+                return loader.load(query_string, **shared_kwargs)
         elif DataSource.BIGQUERY == data_provider:
             from mage_ai.io.bigquery import BigQuery
 
             loader = BigQuery.with_config(config_file_loader)
-            return loader.load(query_string)
+            return loader.load(query_string, **shared_kwargs)
         elif DataSource.REDSHIFT == data_provider:
             from mage_ai.io.redshift import Redshift
 
             with Redshift.with_config(config_file_loader) as loader:
-                return loader.load(query_string)
+                return loader.load(query_string, **shared_kwargs)
         elif DataSource.SNOWFLAKE == data_provider:
             from mage_ai.io.snowflake import Snowflake
 
             with Snowflake.with_config(config_file_loader) as loader:
-                return loader.load(query_string)
+                return loader.load(query_string, **shared_kwargs)
         elif DataSource.TRINO == data_provider:
             from mage_ai.io.trino import Trino
 
             with Trino.with_config(config_file_loader) as loader:
-                return loader.load(query_string)
+                return loader.load(query_string, **shared_kwargs)
 
 
 def build_command_line_arguments(
@@ -657,7 +733,7 @@ def build_command_line_arguments(
     test_execution: bool = False,
 ) -> Tuple[str, List[str], Dict]:
     variables = merge_dict(
-        variables,
+        variables or {},
         get_global_variables(block.pipeline.uuid) if block.pipeline else {},
     )
     dbt_command = 'test' if run_tests else 'run'

mage_ai/data_preparation/models/block/utils.py

@@ -45,7 +45,8 @@ def create_block_runs_from_dynamic_block(
 
     values = []
     block_metadata = []
-    for idx, output_name in enumerate(block.output_variables()):
+    output_vars = block.output_variables(execution_partition=execution_partition)
+    for idx, output_name in enumerate(output_vars):
         if idx == 0:
             values = block.pipeline.variable_manager.get_variable(
                 block.pipeline.uuid,

mage_ai/data_preparation/models/pipeline.py

@@ -470,12 +470,14 @@ class Pipeline:
 
     async def to_dict_async(
         self,
-        include_content=False,
-        include_outputs=False,
-        sample_count=None,
+        include_block_metadata: bool = False,
+        include_content: bool = False,
+        include_outputs: bool = False,
+        sample_count: int = None,
     ):
         blocks_data = await asyncio.gather(
             *[b.to_dict_async(
+                include_block_metadata=include_block_metadata,
                 include_content=include_content,
                 include_outputs=include_outputs,
                 sample_count=sample_count,

mage_ai/data_preparation/repo_manager.py

@@ -37,7 +37,7 @@ class RepoConfig:
             # Priority:
             # 1. 'variables_dir' from config_dict
             # 1. os.getenv(MAGE_DATA_DIR_ENV_VAR)
-            # 2. 'variables_dir' from repo_config file
+            # 2. 'variables_dir' from project's metadata.yaml file
             # 3. DEFAULT_MAGE_DATA_DIR
             if config_dict is not None and config_dict.get('variables_dir'):
                 self.variables_dir = config_dict.get('variables_dir')

mage_ai/data_preparation/storage/local_storage.py

@@ -31,7 +31,7 @@ class LocalStorage(BaseStorage):
         os.remove(path)
 
     def remove_dir(self, path: str) -> None:
-        shutil.rmtree(path)
+        shutil.rmtree(path, ignore_errors=True)
 
     def read_json_file(self, file_path: str, default_value={}) -> Dict:
         if not self.path_exists(file_path):

mage_ai/data_preparation/templates/custom/python/default.jinja

@@ -9,7 +9,7 @@ if 'custom' not in globals():
 def transform_custom(*args, **kwargs):
     """
     Returns:
-        Anything
+        Anything (e.g. data frame, dictionary, array, int, str, etc.)
     """
     # Specify your custom logic here
 {{ code }}

mage_ai/data_preparation/templates/data_exporters/default.jinja

@@ -1,16 +1,14 @@
-from pandas import DataFrame
-
 if 'data_exporter' not in globals():
     from mage_ai.data_preparation.decorators import data_exporter
 
 
 @data_exporter
-def export_data(df: DataFrame, **kwargs):
+def export_data(*args, **kwargs):
     """
     Exports data to some source
 
     Args:
-        df (DataFrame): Data frame to export to
+        args: The input variables from upstream blocks
 
     Output (optional):
         Optionally return any object and it'll be logged and

mage_ai/data_preparation/templates/data_exporters/pyspark/default.jinja

@@ -3,12 +3,12 @@ if 'data_exporter' not in globals():
 
 
 @data_exporter
-def export_data(df, **kwargs):
+def export_data(*args, **kwargs):
     """
     Exports data to some source
 
     Args:
-        df: Data frame to export to
+        args: The input variables from upstream blocks
 
     Output (optional):
         Optionally return any object and it'll be logged and

mage_ai/data_preparation/templates/data_loaders/api.py

@@ -13,7 +13,7 @@ def load_data_from_api(*args, **kwargs):
     Template for loading data from API
     """
     url = ''
-
     response = requests.get(url)
+
     return pd.read_csv(io.StringIO(response.text), sep=',')
 {% endblock %}

mage_ai/data_preparation/templates/data_loaders/default.jinja

@@ -11,7 +11,7 @@ def load_data(*args, **kwargs):
     Template code for loading data from any source.
 
     Returns:
-        Anything
+        Anything (e.g. data frame, dictionary, array, int, str, etc.)
     """
     # Specify your data loading logic here
 {{ code }}

mage_ai/data_preparation/templates/data_loaders/file.py

@@ -18,5 +18,6 @@ def load_data_from_file(*args, **kwargs):
     Docs: https://docs.mage.ai/design/data-loading#fileio
     """
     filepath = 'path/to/your/file.csv'
+
     return FileIO().load(filepath)
 {% endblock %}

mage_ai/data_preparation/templates/data_loaders/pyspark/default.jinja

@@ -11,7 +11,7 @@ def load_data(**kwargs):
     Template code for loading data from any source.
 
     Returns:
-        DataFrame: Returned pandas data frame.
+        Anything (e.g. data frame, dictionary, array, int, str, etc.)
     """
     # Specify your data loading logic here
 {{ code }}

mage_ai/data_preparation/templates/testable.jinja

@@ -9,8 +9,8 @@ if 'test' not in globals():
 
 
 @test
-def test_output(df, *args) -> None:
+def test_output(output, *args) -> None:
     """
     Template code for testing the output of the block.
     """
-    assert df is not None, 'The output is undefined'
+    assert output is not None, 'The output is undefined'

mage_ai/data_preparation/templates/transformers/data_warehouse_transformer.jinja

@@ -4,6 +4,8 @@ from mage_ai.data_preparation.repo_manager import get_repo_path
 from mage_ai.io.config import ConfigFileLoader
 from mage_ai.io.{{ data_source }} import {{ data_source_handler }}
 from os import path
+from pandas import DataFrame
+
 {{ super() -}}
 {% endblock %}