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

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

@@ -1,23 +1,25 @@
 import os
-import re
 import shlex
-from functools import reduce
 from logging import Logger
-from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 import simplejson
 from jinja2 import Template
 
 from mage_ai.data_preparation.models.block.dbt import DBTBlock
-from mage_ai.data_preparation.models.block.dbt.constants import DBT_DIRECTORY_NAME
 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
 from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.data_preparation.templates.utils import get_variable_for_template
 from mage_ai.orchestration.constants import PIPELINE_RUN_MAGE_VARIABLES_KEY
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.custom_logger import DX_PRINTER
+from mage_ai.shared.files import (
+    find_file_from_another_file_path,
+    get_full_file_paths_containing_item,
+)
 from mage_ai.shared.hash import merge_dict
+from mage_ai.shared.path_fixer import add_absolute_path
 
 
 class DBTBlockYAML(DBTBlock):
@@ -32,150 +34,60 @@ class DBTBlockYAML(DBTBlock):
         # Example:
         # demo
         # default_repo/dbt/demo
-        project_name_init = self.configuration.get('dbt_project_name')
-        if project_name_init:
-            try:
-                diff = Path(project_name_init).relative_to(DBT_DIRECTORY_NAME)
-                project_name_init = diff
-            except ValueError:
-                pass
 
-            project_name = self.get_project_path_from_project_name(project_name_init)
-            if project_name:
-                return project_name
-
-            # Adds demo to /home/src/default_repo/dbt
-            # /home/src/default_repo/dbt/demo
-            return str(Path(self.get_base_project_from_source()) / project_name_init)
-
-    def tags(self) -> List[str]:
-        """
-        Get the tags associated with the DBT block.
-
-        Returns:
-            List[str]: The list of tags.
-        """
-        arr = super().tags()
-        command = self._dbt_configuration.get('command', 'run')
-        if command:
-            arr.append(command)
-        return arr
-
-    async def metadata_async(self) -> Dict[str, Any]:
-        """
-        Retrieves metadata needed to configure the block.
-        - available local dbt projects
-        - target to use and other target of the local dbt projects
-
-        Returns:
-            Dict: The metadata of the DBT block.
-        """
-        projects = {}
-
-        project_dirs = Project(self.get_base_project_from_source()).local_packages
-        for project_dir in (project_dirs or []):
-            project_full_path = str(Path(self.get_base_project_from_source()) / project_dir)
-
-            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_dir] = {
-                'project_name': project_name,
-                'target': target,
-                'targets': targets,
-            }
-
-        return {
-            'dbt': {
-                'block': {},
-                'project': None,
-                'projects': projects
-            }
-        }
-
-    def _hydrate_block_outputs(
-        self,
-        content: str,
-        outputs_from_input_vars: Dict,
-        variables: Dict = None,
-    ) -> str:
-        def _block_output(
-            block_uuid: str = None,
-            parse: str = None,
-            outputs_from_input_vars=outputs_from_input_vars,
-            upstream_block_uuids=self.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:
-                        acc.append(outputs_from_input_vars.get(upstream_block_uuid))
-                        return acc
-
-                    data = reduce(
-                        _build_positional_arguments,
-                        upstream_block_uuids,
-                        [],
-                    )
-
-                try:
-                    return parse(data, variables)
-                except Exception as err:
-                    print(f'[WARNING] block_output: {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 self.configuration.get(''dbt_project_path''):
+        #     return add_absolute_path(self.configuration.get(''dbt_project_path''))
+
+        file_path = None
+        for key in [
+            'dbt_profiles_file_path',
+            'dbt_project_name',
+        ]:
+            if not file_path and self.configuration.get(key):
+                path = add_absolute_path(self.configuration.get(key))
+                file_path = find_file_from_another_file_path(
+                    path,
+                    lambda x: os.path.basename(x) in [
+                        'dbt_project.yml',
+                        'dbt_project.yaml'
+                    ],
+                )
+                if file_path:
+                    return os.path.dirname(file_path)
+
+        if file_path is None and self.configuration and self.configuration.get('dbt_project_name'):
+            file_paths = get_full_file_paths_containing_item(
+                os.path.join(
+                    get_repo_path(root_project=False),
+                    'dbt',
+                    self.configuration.get('dbt_project_name'),
+                ),
+                lambda x: x and (
+                    str(x).endswith('dbt_project.yml') or
+                    str(x).endswith('dbt_project.yaml'),
+                ),
             )
-
-            if not match:
-                continue
-
-            si, ei = match.span()
-            substring = content[si:ei]
-
-            match2 = re.match(
-                r'{}([^{}{}]+){}'.format(r'\{\{', r'\{\{', r'\}\}', r'\}\}'),
-                substring,
+            if file_paths:
+                file_path = file_paths[0]
+                return os.path.dirname(file_path)
+
+    def set_default_configurations(self):
+        self.configuration = self.configuration or {}
+        if not self.configuration.get('file_source'):
+            self.configuration['file_source'] = {}
+
+        pp = self.project_path
+        if pp:
+            self.configuration['file_source']['project_path'] = add_absolute_path(
+                pp,
+                add_base_repo_path=False,
             )
-            if match2:
-                groups = match2.groups()
-                if groups:
-                    function_string = groups[0].strip()
-                    results = 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 content
+        if not self.configuration['file_source'].get('path'):
+            self.configuration['file_source']['path'] = add_absolute_path(
+                self.file_path,
+                add_base_repo_path=False,
+            )
 
     def _execute_block(
         self,
@@ -194,21 +106,19 @@ class DBTBlockYAML(DBTBlock):
             runtime_arguments (Optional[Dict[str, Any]], optional): The runtime arguments.
         """
         # Which dbt task should be executed
+        DX_PRINTER.label = 'DBTYamlBlock'
+
         task = self._dbt_configuration.get('command', 'run')
 
         # Get variables
         variables = merge_dict(global_vars, runtime_arguments or {})
 
         content = self.content or ''
-
-        content = self._hydrate_block_outputs(content, outputs_from_input_vars, variables)
-        content = Template(content).render(
-            variables=lambda x, p=None, v=variables: get_variable_for_template(
-                x,
-                parse=p,
-                variables=v,
-            ),
-            **get_template_vars(),
+        content = self.interpolate_content(
+            content,
+            outputs_from_input_vars=outputs_from_input_vars,
+            variables=variables,
+            **kwargs,
         )
 
         # Get args from block content and split them by word, just like a shell does
@@ -236,6 +146,14 @@ class DBTBlockYAML(DBTBlock):
         except Exception:
             pass
 
+        DX_PRINTER.debug(
+            f'execute block content: {content}',
+            block=self,
+            args=args,
+            file_path=self.file_path,
+            project_path=self.project_path,
+        )
+
         if vars_index:
             variables2 = Template(args[vars_index]).render(
                 variables=lambda x, p=None, v=variables: get_variable_for_template(
@@ -267,6 +185,11 @@ class DBTBlockYAML(DBTBlock):
             args += ([
                 "--profiles-dir", str(profiles.profiles_dir)
             ])
-            _res, success = DBTCli([task] + args, logger).invoke()
-            if not success:
-                raise Exception('DBT exited with a non 0 exit status.')
+
+            cli = DBTCli(logger=logger)
+
+            cli.invoke(['deps'] + args)
+
+            res = cli.invoke([task] + args)
+            if not res.success:
+                raise Exception(str(res.exception))

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

@@ -1,3 +1,5 @@
+from enum import Enum
+
 DBT_DIRECTORY_NAME = 'dbt'
 
 SQL_SERVER_ADAPTER_NAME = 'SQLServerAdapter'
@@ -9,3 +11,13 @@ SKIP_LIMIT_ADAPTER_NAMES = [
     SQL_SERVER_ADAPTER_NAME,
     SYNAPSE_ADAPTER_NAME,
 ]
+
+
+class Flag(str, Enum):
+    PROFILES_DIR = 'profiles-dir'
+    PROJECT_DIR = 'project-dir'
+
+
+class LogLevel(str, Enum):
+    DEBUG = 'debug'
+    INFO = 'info'

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

@@ -19,6 +19,7 @@ from dbt.contracts.connection import AdapterResponse
 from dbt.contracts.relation import RelationType
 
 from mage_ai.data_preparation.models.block.dbt.profiles import Profiles
+from mage_ai.shared.environments import is_debug
 
 
 @dataclass
@@ -30,6 +31,7 @@ class DBTAdapterConfig:
     profiles_dir: Union[str, os.PathLike]
     profile: Optional[str] = None
     target: Optional[str] = None
+    threads = None
 
 
 class DBTAdapter:
@@ -220,18 +222,29 @@ class DBTAdapter:
             target=self.target
         )
         flags.set_from_args(adapter_config, user_config)
-        config = RuntimeConfig.from_args(adapter_config)
-        reset_adapters()
-        # register the correct adapter from config
-        register_adapter(config)
-        # load the adapter
-        self.__adapter = get_adapter(config)
-        # connect
-        self.__adapter.acquire_connection('mage_dbt_adapter_' + uuid.uuid4().hex)
-        return self
+
+        try:
+            config = RuntimeConfig.from_args(adapter_config)
+            reset_adapters()
+            # register the correct adapter from config
+            register_adapter(config)
+            # load the adapter
+            self.__adapter = get_adapter(config)
+            # connect
+            self.__adapter.acquire_connection('mage_dbt_adapter_' + uuid.uuid4().hex)
+            return self
+        except Exception as err:
+            if is_debug():
+                raise err
+            print(f'[DBTAdapter] open: {err}.')
 
     def __enter__(self):
         return self.open()
 
     def __exit__(self, *args):
-        self.close()
+        try:
+            self.close()
+        except Exception as err:
+            if is_debug():
+                raise err
+            print(f'[DBTAdapter] __exit__: {err}.')