mage-ai 0.9.58__py3-none-any.whl → 0.9.59__py3-none-any.whl

mage_ai/data_preparation/models/block/content.py

@@ -0,0 +1,111 @@
+import datetime
+import json
+import math
+import re
+from functools import reduce
+from typing import Any, Dict, List
+
+import inflection
+import pandas as pd
+from jinja2 import Template
+
+from mage_ai.data_preparation.shared.utils import (
+    get_template_vars,
+    get_template_vars_no_db,
+)
+from mage_ai.data_preparation.templates.utils import get_variable_for_template
+from mage_ai.shared.hash import merge_dict
+
+
+def hydrate_block_outputs(
+    content: str,
+    outputs_from_input_vars: Dict = None,
+    upstream_block_uuids: List[str] = None,
+    variables: Dict = None,
+) -> str:
+    include_python_libraries = dict(
+        datetime=datetime,
+        inflection=inflection,
+        json=json,
+        math=math,
+    )
+
+    def _block_output(
+        block_uuid: str = None,
+        parse: str = None,
+        outputs_from_input_vars=outputs_from_input_vars,
+        include_python_libraries=include_python_libraries,
+        upstream_block_uuids=upstream_block_uuids,
+        variables=variables,
+    ) -> Any:
+        data = outputs_from_input_vars
+
+        if parse:
+            if block_uuid:
+                data = outputs_from_input_vars.get(block_uuid)
+            elif upstream_block_uuids:
+                def _build_positional_arguments(acc: List, upstream_block_uuid: str) -> List:
+                    results = outputs_from_input_vars.get(upstream_block_uuid)
+                    if isinstance(results, pd.DataFrame):
+                        results = results.to_dict(orient='records')
+                    acc.append(results)
+
+                    return acc
+
+                data = reduce(
+                    _build_positional_arguments,
+                    upstream_block_uuids,
+                    [],
+                )
+
+            try:
+                return parse(data, variables)
+            except Exception as err:
+                print(f'[ERROR] Parsing block_output function: {err}')
+
+        return data
+
+    variable_pattern = r'{}[ ]*block_output[^{}]*[ ]*{}'.format(r'\{\{', r'\}\}', r'\}\}')
+
+    match = 1
+    while match is not None:
+        match = None
+
+        match = re.search(
+            variable_pattern,
+            content,
+            re.IGNORECASE,
+        )
+
+        if not match:
+            continue
+
+        si, ei = match.span()
+        substring = content[si:ei]
+
+        match2 = re.match(
+            r'{}([^{}{}]+){}'.format(r'\{\{', r'\{\{', r'\}\}', r'\}\}'),
+            substring,
+        )
+        if match2:
+            groups = match2.groups()
+            if groups:
+                function_string = groups[0].strip()
+                results = merge_dict(get_template_vars_no_db(
+                    include_python_libraries=include_python_libraries,
+                ), dict(
+                    block_output=_block_output,
+                ))
+                exec(f'value_hydrated = {function_string}', results)
+
+                value_hydrated = results['value_hydrated']
+                content = f'{content[0:si]}{value_hydrated}{content[ei:]}'
+
+    return Template(content).render(
+        variables=lambda x, p=None, v=variables: get_variable_for_template(
+            x,
+            parse=p,
+            variables=v,
+        ),
+        **get_template_vars(include_python_libraries=include_python_libraries),
+    )

mage_ai/data_preparation/models/block/data_integration/mixins.py

@@ -419,6 +419,7 @@ class DataIntegrationMixin:
         upstream_block_uuids: List[str] = None,
         all_catalogs: bool = False,
         all_streams: bool = False,
+        **kwargs,
     ) -> Tuple[List, List, List]:
         block_uuids_to_fetch = upstream_block_uuids or self.upstream_block_uuids_for_inputs
 

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

@@ -20,6 +20,7 @@ from mage_ai.data_preparation.models.block.dbt.sources import Sources
 from mage_ai.data_preparation.models.constants import BlockLanguage
 from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.orchestration.constants import PIPELINE_RUN_MAGE_VARIABLES_KEY
+from mage_ai.shared.environments import is_debug
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.parsers import encode_complex
 
@@ -244,25 +245,31 @@ class DBTBlock(Block):
             if block_uuids:
                 for (project_path, target) in targets:
                     try:
-                        with DBTAdapter(
-                            project_path,
-                            variables=variables,
-                            target=target
-                        ) as dbt_adapter:
-                            credentials = dbt_adapter.credentials
-                            # some databases use other default schema names
-                            # e.g. duckdb uses main schema as default
-                            schema = getattr(credentials, 'schema', None)
-
-                        Sources(project_path).reset_pipeline(
-                            project_name=Path(project_path).stem,
-                            pipeline_uuid=pipeline_uuid,
-                            block_uuids=block_uuids,
-                            schema=schema
-                        )
-                    # project not yet configured correctly, so just skip that step for now
-                    except FileNotFoundError:
-                        pass
+                        try:
+                            with DBTAdapter(
+                                project_path,
+                                variables=variables,
+                                target=target
+                            ) as dbt_adapter:
+                                credentials = dbt_adapter.credentials
+                                # some databases use other default schema names
+                                # e.g. duckdb uses main schema as default
+                                schema = getattr(credentials, 'schema', None)
+
+                            Sources(project_path).reset_pipeline(
+                                project_name=Path(project_path).stem,
+                                pipeline_uuid=pipeline_uuid,
+                                block_uuids=block_uuids,
+                                schema=schema
+                            )
+                        # project not yet configured correctly, so just skip that step for now
+                        except FileNotFoundError:
+                            pass
+                    except Exception as err:
+                        if is_debug():
+                            raise err
+                        else:
+                            print(f'[DBTBlock] update_sources: {err}')
 
     @contextmanager
     def _redirect_streams(
@@ -270,3 +277,6 @@ class DBTBlock(Block):
         **kwargs
     ) -> Generator[None, None, None]:
         yield (None, None)
+
+    def set_default_configurations(self):
+        pass

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

@@ -5,15 +5,11 @@ from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import pandas as pd
-import simplejson
 from jinja2 import Template
 
 from mage_ai.data_preparation.models.block import Block
 from mage_ai.data_preparation.models.block.dbt import DBTBlock
-from mage_ai.data_preparation.models.block.dbt.constants import (
-    DBT_DIRECTORY_NAME,
-    LogLevel,
-)
+from mage_ai.data_preparation.models.block.dbt.constants import LogLevel
 from mage_ai.data_preparation.models.block.dbt.dbt_cli import DBTCli
 from mage_ai.data_preparation.models.block.dbt.profiles import Profiles
 from mage_ai.data_preparation.models.block.dbt.project import Project
@@ -27,9 +23,10 @@ from mage_ai.data_preparation.models.block.platform.mixins import (
 from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.orchestration.constants import PIPELINE_RUN_MAGE_VARIABLES_KEY
+from mage_ai.settings.utils import base_repo_path
+from mage_ai.shared.files import find_file_from_another_file_path
 from mage_ai.shared.hash import merge_dict
-from mage_ai.shared.strings import remove_extension_from_filename
-from mage_ai.shared.utils import clean_name
+from mage_ai.shared.path_fixer import add_absolute_path
 
 
 class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
@@ -43,7 +40,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
         """
         project = Project(self.project_path).project
         target_path = project.get('target-path', 'target')
-        file_path = self.__get_original_file_path()
+        file_path = self.file_path
         compiled_path = Path(self.project_path) / target_path / 'compiled' / file_path
 
         compiled_sql = None
@@ -63,16 +60,9 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
         Returns:
             Union[str, os.PathLike]: The file path of the DBT block.
         """
-        if self.project_platform_activated:
-            file_path = self.get_file_path_from_source()
-            if file_path:
-                return file_path
-
-        file_path = self.configuration.get('file_path') or ''
-        if file_path and Path(file_path).parts[0] == 'dbt':
-            file_path = os.path.join(*Path(file_path).parts[1:]) or ''
-
-        return str((Path(self.repo_path)) / DBT_DIRECTORY_NAME / file_path)
+        return add_absolute_path(
+            self.get_file_path_from_source() or self.configuration.get('file_path'),
+        )
 
     @property
     def project_path(self) -> Union[str, os.PathLike]:
@@ -82,21 +72,46 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
         Returns:
             Union[str, os.PathLike]: Path of the dbt project, being used
         """
-        if self.project_platform_activated:
-            project_path = self.get_project_path_from_source()
-            if project_path:
-                return project_path
+        if not self.file_path:
+            return
+
+        project_file_path = find_file_from_another_file_path(
+            self.file_path,
+            lambda x: os.path.basename(x) in [
+                'dbt_project.yml',
+                'dbt_project.yaml',
+            ],
+        )
+
+        if project_file_path:
+            return os.path.dirname(project_file_path)
+
+    def set_default_configurations(self):
+        self.configuration = self.configuration or {}
+        file_path = self.configuration.get('file_path') or \
+            (self.configuration.get('file_source') or {}).get('path')
+
+        file_path = add_absolute_path(file_path, add_base_repo_path=False) if file_path else None
+        if not file_path:
+            return
+
+        self.configuration['file_path'] = file_path
+        self.configuration['file_source'] = dict(
+            path=file_path,
+        )
 
         try:
-            diff = Path(self.file_path).relative_to(self.base_project_path)
-            if diff:
-                diff = Path(diff).parts[0]
+            self.uuid = str(Path(self.file_path).relative_to(base_repo_path()).with_suffix(''))
         except ValueError:
-            diff = self.file_path
+            pass
 
-        return str(
-            Path(self.base_project_path) / diff,
-        )
+        if self.project_path:
+            try:
+                pp = str(Path(self.project_path).relative_to(base_repo_path()))
+                self.configuration['dbt_project_name'] = pp
+                self.configuration['file_source']['project_path'] = pp
+            except ValueError:
+                pass
 
     @property
     def __node_type(self) -> str:
@@ -115,237 +130,12 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
         }
         search_paths = {_path: k for k, v in search_paths.items() for _path in v}
 
-        node_path = self.__get_original_file_path()
+        node_path = self.file_path
         node_path_parts = node_path.split(os.sep)
         node_path_type_part = node_path_parts[1]
 
         return search_paths.get(node_path_type_part)
 
-    async def metadata_async(self) -> Dict[str, Any]:
-        """
-        Retrieves metadata needed to configure the block.
-        - dbt project to use
-        - dbt target to use and other dbt targets for the project
-
-        Returns:
-            Dict: The metadata of the DBT block.
-        """
-        project = None
-        projects = {}
-
-        project_full_path = self.project_path
-        project = Project(project_full_path).project
-        project_name = project.get('name')
-        project_profile = project.get('profile')
-
-        # ignore exception if no profiles.yml is found.
-        # Just means that the targets have no option
-        try:
-            profiles = Profiles(
-                project_full_path,
-                self.pipeline.variables if self.pipeline else None
-            ).profiles
-        except FileNotFoundError:
-            profiles = None
-        profile = profiles.get(project_profile) if profiles else None
-        target = profile.get('target') if profile else None
-        targets = sorted(list(profile.get('outputs', {}).keys())) if profile else None
-
-        projects = {
-            project_name: {
-                'target': target,
-                'targets': targets
-            }
-        }
-
-        node_type = self.__node_type
-
-        return {
-            DBT_DIRECTORY_NAME: {
-                'block': {
-                    'snapshot': node_type == 'snapshot'
-                },
-                'project': project_name,
-                'projects': projects
-            }
-        }
-
-    def tags(self) -> List[str]:
-        """
-        Get the tags associated with the DBT block.
-
-        Returns:
-            List[str]: The list of tags.
-        """
-        arr = super().tags()
-        node_type = self.__node_type
-        if node_type == 'snapshot':
-            arr.append(self.__node_type)
-        return arr
-
-    def upstream_dbt_blocks(self, read_only=False) -> List['DBTBlockSQL']:
-        """
-        Get an up to date list, which represents the upstream dbt graph.
-        It is using `dbt list` to generate the list.
-
-        Args:
-            read_only (bool):
-                If True it does not read the Blocks from the model. Defaults to False
-
-        Returns:
-            List[DBTBlockSQL]: THe upstream dbt graph as DBTBlocksSQL objects
-        """
-        # Get upstream nodes via dbt list
-        with Profiles(
-            self.project_path,
-            self.pipeline.variables if self.pipeline else {},
-        ) as profiles:
-            try:
-                args = [
-                    'list',
-                    # project-dir
-                    # /home/src/default_repo/default_platform/default_repo/dbt/demo
-                    '--project-dir', self.project_path,
-                    '--profiles-dir', str(profiles.profiles_dir),
-                    '--select', '+' + Path(self.__get_original_file_path()).stem,
-                    '--output', 'json',
-                    '--output-keys', 'unique_id original_file_path depends_on',
-                    '--resource-type', 'model',
-                    '--resource-type', 'snapshot'
-                ]
-                res = DBTCli().invoke(args)
-
-                if res:
-                    nodes_init = [simplejson.loads(node) for node in res.result]
-                else:
-                    return []
-            except Exception as err:
-                print(f'[ERROR] DBTBlockSQL.upstream_dbt_blocks: {err}.')
-                return [
-                    self.build_dbt_block(
-                        block_class=DBTBlock,
-                        block_dict=dict(
-                            block_type=self.type,
-                            configuration=self.configuration,
-                            language=self.language,
-                            name=self.uuid,
-                            pipeline=self.pipeline,
-                            uuid=self.uuid,
-                        ),
-                        hydrate_configuration=False,
-                    )
-                ]
-
-        # transform List into dict and remove unnecessary fields
-        file_path = self.__get_original_file_path()
-        path_parts = file_path.split(os.sep)
-        project_dir = path_parts[0]
-
-        nodes_default = {
-            node['unique_id']: {
-                # file_path:
-                # default_repo/dbt/demo/models/example/model.sql
-                # default_platform/default_repo/dbt/demo/models/example/model.sql
-                'file_path': os.path.join(project_dir, node['original_file_path']),
-                'original_file_path': node['original_file_path'],
-                'upstream_nodes': set(node['depends_on']['nodes'])
-            }
-            for node in nodes_init
-        }
-
-        nodes = self.hydrate_dbt_nodes(nodes_default, nodes_init)
-
-        # calculate downstream_nodes
-        for unique_id, node in nodes.items():
-            for upstream_node in node['upstream_nodes']:
-                if nodes.get(upstream_node):
-                    downstream_nodes = nodes[upstream_node].get('downstream_nodes', set())
-                    downstream_nodes.add(unique_id)
-                    nodes[upstream_node]['downstream_nodes'] = downstream_nodes
-
-        # map dbt unique_id to mage uuid
-        uuids_default = {
-            unique_id: clean_name(
-                remove_extension_from_filename(node['file_path']),
-                allow_characters=[os.sep]
-            )
-            for unique_id, node in nodes.items()
-        }
-        uuids = self.node_uuids_mapping(uuids_default, nodes)
-
-        # replace dbt unique_ids with mage uuids
-        nodes = {
-            uuids[unique_id]: {
-                'file_path': node['file_path'],
-                'original_file_path': node['original_file_path'],
-                'upstream_nodes': {
-                    uuids[upstream_node]
-                    for upstream_node in node.get('upstream_nodes', set())
-                    if uuids.get(upstream_node)
-                },
-                'downstream_nodes': {
-                    uuids[downstream_node]
-                    for downstream_node in node.get('downstream_nodes', set())
-                    if uuids.get(downstream_node)
-                }
-            }
-            for unique_id, node in nodes.items()
-        }
-
-        # create dict of blocks
-        blocks = {}
-        for uuid, node in nodes.items():
-            block = None
-            if not read_only:
-                if uuid == self.uuid:
-                    block = self
-                elif self.pipeline:
-                    block = self.pipeline.get_block(
-                        uuid,
-                        self.type,
-                    )
-            # if not found create the block
-            block = block or self.build_dbt_block(
-                block_class=DBTBlock,
-                block_dict=dict(
-                    block_type=self.type,
-                    configuration=dict(file_path=node['file_path']),
-                    language=self.language,
-                    name=uuid,
-                    pipeline=self.pipeline,
-                    uuid=uuid,
-                ),
-                node=node,
-            )
-            # reset upstream dbt blocks
-            block.upstream_blocks = [
-                block
-                for block in block.upstream_blocks
-                # there must not be any other upstream dbt blocks that are not part of
-                # the upstream nodes, therefore delete all upstream dbt blocks
-                if not isinstance(block, DBTBlockSQL)
-            ]
-            # reset downstream dbt blocks
-            block.downstream_blocks = [
-                block
-                for block in block.downstream_blocks
-                # remove all downstream dbt block, which are part of the upstream nodes,
-                # in order to fix errors
-                if block.uuid not in nodes.keys()
-            ]
-            blocks[uuid] = block
-
-        # Update upstream and downstream dbt blocks
-        for uuid, node in nodes.items():
-            for upstream_node in node.get('upstream_nodes', set()):
-                blocks[uuid].upstream_blocks.append(blocks[upstream_node])
-            for downstream_node in node.get('downstream_nodes', set()):
-                blocks[uuid].downstream_blocks.append(blocks[downstream_node])
-
-        # transform into list
-        blocks = [block for _, block in blocks.items()]
-        return blocks
-
     def _execute_block(
         self,
         outputs_from_input_vars,
@@ -408,7 +198,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
         suffix = ''
         if runtime_configuration.get('suffix'):
             suffix = runtime_configuration['suffix']
-        args += ['--select', f"{prefix}{Path(self.__get_original_file_path()).stem}{suffix}"]
+        args += ['--select', f"{prefix}{Path(self.file_path).stem}{suffix}"]
 
         # Create --vars
         args += ['--vars', self._variables_json(variables)]
@@ -459,7 +249,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
                 res = cli.invoke([task] + args)
                 success = res.success
                 if not success:
-                    raise res.exception
+                    raise Exception(str(res.exception))
             # run show task, to get data for preview or downstream usage
             # test task does not have any data
             #
@@ -471,7 +261,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
                 if res.success:
                     df = cli.to_pandas(res)
                 else:
-                    raise res.exception
+                    raise Exception(str(res.exception))
 
         # provide df for downstream usage or data preview
         self.store_variables(
@@ -491,7 +281,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
             '--project-dir',
             self.project_path,
             '--select',
-            Path(self.__get_original_file_path()).stem,
+            Path(self.file_path).stem,
             '--limit',
             str(limit),
         ]
@@ -515,7 +305,7 @@ class DBTBlockSQL(DBTBlock, ProjectPlatformAccessible):
             if res.success:
                 return cli.to_pandas(res)
             else:
-                raise res.exception
+                raise Exception(str(res.exception))
 
     def __create_upstream_tables(
         self,