mage-ai 0.9.44__py3-none-any.whl → 0.9.46__py3-none-any.whl

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

@@ -75,7 +75,7 @@ def convert_dataframe_to_output(
             if log_message:
                 log_message(
                     f'Writing {len(records_in_batch)} records from stream {stream} to {file_path} '
-                    f'for batch index {index} ({index + 1 } out of {batches} batch(s)).',
+                    f'for batch index {index} ({index + 1} out of {batches} batch(s)).',
                 )
 
             output_file_paths.append(file_path)

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

@@ -1,15 +1,22 @@
 import os
+import re
 from logging import Logger
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Union
+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.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.models.block.dbt.utils import (
+    get_source_name,
+    get_source_table_name_for_block,
+)
 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
@@ -305,6 +312,15 @@ class DBTBlockSQL(DBTBlock):
         Returns:
             List: The list of outputs.
         """
+        # Create upstream tables
+        self.__create_upstream_tables(
+            execution_partition=execution_partition,
+            global_vars=global_vars,
+            logger=logger,
+            outputs_from_input_vars=outputs_from_input_vars,
+            runtime_arguments=runtime_arguments,
+        )
+
         # Which dbt task should be executed
         task = self.__task(from_notebook, run_settings)
 
@@ -400,6 +416,65 @@ class DBTBlockSQL(DBTBlock):
 
         return [df]
 
+    def __create_upstream_tables(
+        self,
+        execution_partition: Optional[str] = None,
+        global_vars: Dict = None,
+        logger: Logger = None,
+        outputs_from_input_vars: List = None,
+        runtime_arguments: Optional[Dict[str, Any]] = None,
+    ):
+        if outputs_from_input_vars is None:
+            outputs_from_input_vars = []
+
+        upstream_blocks = self.__upstream_blocks_from_sources(global_vars=global_vars)
+
+        for ublock in upstream_blocks:
+
+            output = None
+            if ublock.uuid in outputs_from_input_vars:
+                output = outputs_from_input_vars[ublock.uuid]
+            else:
+                output = self.pipeline.variable_manager.get_variable(
+                    self.pipeline.uuid,
+                    ublock.uuid,
+                    'output_0',
+                    partition=execution_partition,
+                )
+
+            # normalize output
+            if isinstance(output, pd.DataFrame):
+                df = output
+            elif isinstance(output, dict):
+                df = pd.DataFrame([output])
+            elif isinstance(output, list):
+                df = pd.DataFrame(output)
+            else:
+                df = pd.DataFrame()
+
+            if df.empty:
+                if logger:
+                    logger.info('No data for dbt to materialize.')
+            else:
+                DBTBlock.materialize_df(
+                    df=df,
+                    pipeline_uuid=self.pipeline.uuid,
+                    block_uuid=ublock.uuid,
+                    targets=[(self.project_path, self.target(variables=global_vars))],
+                    logger=logger,
+                    global_vars=global_vars,
+                    runtime_arguments=runtime_arguments,
+                )
+
+    def __extract_sources(self) -> List[Tuple[str, str]]:
+        return re.findall(
+            r"{}[ ]*source\(['\"]+([\w]+)['\"]+[,]+[ ]*['\"]+([\w]+)['\"]+\)[ ]*{}".format(
+                r'\{\{',
+                r'\}\}',
+            ),
+            self.content,
+        )
+
     def __task(
         self,
         from_notebook: bool,
@@ -452,3 +527,22 @@ class DBTBlockSQL(DBTBlock):
             else:
                 return 'run'
         return 'build'
+
+    def __upstream_blocks_from_sources(self, global_vars: Dict = None) -> List[Block]:
+        mapping = {}
+        sources = self.__extract_sources()
+        for tup in sources:
+            source_name, table_name = tup
+            if source_name not in mapping:
+                mapping[source_name] = {}
+            mapping[source_name][table_name] = True
+
+        source_name = get_source_name(Path(self.project_path).stem)
+
+        arr = []
+        for b in self.upstream_blocks:
+            table_name = get_source_table_name_for_block(b)
+            if mapping.get(source_name, {}).get(table_name):
+                arr.append(b)
+
+        return arr

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

@@ -0,0 +1,6 @@
+def get_source_name(dbt_project_name: str) -> str:
+    return f'mage_{dbt_project_name}'
+
+
+def get_source_table_name_for_block(block) -> str:
+    return f'{block.pipeline.uuid}_{block.uuid}'

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

@@ -12,6 +12,7 @@ from mage_ai.data_preparation.models.block.spark.constants import (
 from mage_ai.data_preparation.models.project import Project
 from mage_ai.data_preparation.models.project.constants import FeatureUUID
 from mage_ai.services.spark.api.service import API
+from mage_ai.services.spark.constants import ComputeServiceUUID
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.services.spark.models.jobs import Job
 from mage_ai.services.spark.models.sqls import Sql
@@ -30,6 +31,16 @@ class SparkBlock:
 
         return self._spark_session_current
 
+    @property
+    def compute_service_uuid(self) -> ComputeServiceUUID:
+        if self._compute_service_uuid:
+            return self._compute_service_uuid
+        self._compute_service_uuid = get_compute_service(
+            ignore_active_kernel=True,
+            repo_config=self.repo_config,
+        )
+        return self._compute_service_uuid
+
     def spark_session_application(self) -> Application:
         if not self.spark_session:
             return
@@ -40,7 +51,7 @@ class SparkBlock:
         if value_tup:
             application_id = value_tup[1]
 
-            return Application(
+            return Application.load(
                 id=application_id,
                 spark_ui_url=self.spark_session.sparkContext.uiWebUrl,
             )
@@ -51,7 +62,13 @@ class SparkBlock:
         )
 
     def is_using_spark(self) -> bool:
-        return get_compute_service()
+        return self.compute_service_uuid in [
+            ComputeServiceUUID.AWS_EMR,
+            ComputeServiceUUID.STANDALONE_CLUSTER,
+        ]
+
+    def should_track_spark(self) -> bool:
+        return self.is_using_spark() and self.compute_management_enabled()
 
     def execution_states(self, cache: bool = False) -> Dict:
         jobs_cache = self.__load_cache()
@@ -82,33 +99,40 @@ class SparkBlock:
     def jobs_during_execution(self) -> List[Job]:
         self.__load_spark_job_submission_timestamps()
 
-        if self.execution_timestamp_start:
-            jobs = self.__get_jobs(application=self.execution_start_application)
-
-            def _filter(
-                job: Job,
-                execution_timestamp_start: float = self.execution_timestamp_start,
-                execution_timestamp_end: float = self.execution_timestamp_end,
-            ) -> bool:
-                if not job.submission_time:
-                    return False
-
-                if isinstance(job.submission_time, str):
-                    submission_timestamp = dateutil.parser.parse(job.submission_time).timestamp()
-                elif isinstance(job.submission_time, float) or isinstance(job.submission_time, int):
-                    submission_timestamp = datetime.fromtimestamp(job.submission_time)
-
-                return execution_timestamp_start and \
-                    execution_timestamp_end and \
-                    submission_timestamp >= execution_timestamp_start and \
-                    (
-                        not execution_timestamp_end or
-                        submission_timestamp <= execution_timestamp_end
-                    )
+        if not self.execution_timestamp_start and self.execution_uuid_start:
+            return []
+
+        def _filter(
+            job: Job,
+            block_uuid=self.uuid,
+            compute_service_uuid=self.compute_service_uuid,
+            execution_timestamp_end: float = self.execution_timestamp_end,
+            execution_timestamp_start: float = self.execution_timestamp_start,
+            execution_uuid_start=self.execution_uuid_start,
+        ) -> bool:
+            if ComputeServiceUUID.AWS_EMR == compute_service_uuid:
+                key = self.execution_uuid_start or self.execution_timestamp_start
+                return job.name == f'{block_uuid}:{key}'
+
+            if not job.submission_time:
+                return False
+
+            if isinstance(job.submission_time, str):
+                submission_timestamp = dateutil.parser.parse(job.submission_time).timestamp()
+            elif isinstance(job.submission_time, float) or isinstance(job.submission_time, int):
+                submission_timestamp = datetime.fromtimestamp(job.submission_time)
+
+            return execution_timestamp_start and \
+                execution_timestamp_end and \
+                submission_timestamp >= execution_timestamp_start and \
+                (
+                    not execution_timestamp_end or
+                    submission_timestamp <= execution_timestamp_end
+                )
 
-            return list(filter(_filter, jobs))
+        jobs = self.__get_jobs(application=self.execution_start_application)
 
-        return []
+        return list(filter(_filter, jobs or []))
 
     def stages_during_execution(self, jobs: List[Job]):
         if not jobs:
@@ -122,7 +146,7 @@ class SparkBlock:
 
         stages = self.__get_stages(application=self.execution_start_application)
 
-        return list(filter(_filter, stages))
+        return list(filter(_filter, stages or []))
 
     def sqls_during_execution(self, jobs: List[Job]):
         if not jobs:
@@ -140,7 +164,7 @@ class SparkBlock:
 
         sqls = self.__get_sqls(application=self.execution_start_application)
 
-        return list(filter(_filter, sqls))
+        return list(filter(_filter, sqls or []))
 
     def clear_spark_jobs_cache(self) -> None:
         if os.path.exists(self.spark_jobs_full_path):
@@ -155,13 +179,24 @@ class SparkBlock:
         for application in applications:
             Application.cache_application(application)
 
-    def set_spark_job_execution_start(self) -> None:
+    def set_spark_job_execution_start(self, execution_uuid: str = None) -> None:
         self.execution_timestamp_start = datetime.utcnow().timestamp()
         application = self.spark_session_application()
 
+        if execution_uuid:
+            self.execution_uuid = execution_uuid
+
+        if self.spark_session and self.spark_session.sparkContext:
+            key = f'{self.uuid}:{self.execution_uuid or self.execution_timestamp_start}'
+            # For jobs
+            self.spark_session.sparkContext.setLocalProperty('callSite.short', key)
+            # For stages
+            self.spark_session.sparkContext.setLocalProperty('callSite.long', key)
+
         self.__update_spark_jobs_cache(
             dict(
                 application=application.to_dict() if application else None,
+                execution_uuid=self.execution_uuid,
                 submission_timestamp=self.execution_timestamp_start,
             ),
             'before',
@@ -187,7 +222,7 @@ class SparkBlock:
 
         if application:
             build_options.update(dict(
-                application_id=application.id,
+                application_id=application.calculated_id(),
                 application_spark_ui_url=application.spark_ui_url,
             ))
         else:
@@ -207,7 +242,7 @@ class SparkBlock:
         else:
             applications = api.applications_sync()
             if applications:
-                jobs = api.jobs_sync(applications[0].id)
+                jobs = api.jobs_sync(applications[0].calculated_id())
 
         return sorted(
             jobs,
@@ -231,7 +266,7 @@ class SparkBlock:
             applications = api.applications_sync()
 
             if applications:
-                stages = api.stages_sync(applications[0].id, dict(
+                stages = api.stages_sync(applications[0].calculated_id(), dict(
                     quantiles='0.01,0.25,0.5,0.75,0.99',
                     withSummaries=True,
                 ))
@@ -253,7 +288,7 @@ class SparkBlock:
 
             sqls = []
             if applications:
-                sqls = api.sqls_sync(applications[0].id, dict(
+                sqls = api.sqls_sync(applications[0].calculated_id(), dict(
                     length=9999,
                 ))
 
@@ -304,6 +339,13 @@ class SparkBlock:
             f.write(json.dumps(data))
 
     def __load_spark_job_submission_timestamps(self) -> None:
+        self.execution_end_application = None
+        self.execution_start_application = None
+        self.execution_timestamp_end = None
+        self.execution_timestamp_start = None
+        self.execution_uuid_end = None
+        self.execution_uuid_start = None
+
         jobs_cache = self.__load_cache()
         if jobs_cache:
             before = jobs_cache.get('before')
@@ -311,6 +353,9 @@ class SparkBlock:
                 self.execution_timestamp_start = (before or {}).get(
                     'submission_timestamp',
                 )
+                self.execution_uuid_start = (before or {}).get(
+                    'execution_uuid',
+                )
                 self.execution_start_application = (before or {}).get(
                     'application',
                 )
@@ -324,6 +369,9 @@ class SparkBlock:
                 self.execution_timestamp_end = (after or {}).get(
                     'submission_timestamp',
                 )
+                self.execution_uuid_end = (after or {}).get(
+                    'execution_uuid',
+                )
                 self.execution_end_application = (after or {}).get(
                     'application',
                 )

mage_ai/data_preparation/models/download/__init__.py

@@ -0,0 +1,8 @@
+class Download():
+    def __init__(self, token):
+        self.token = token
+
+    def to_dict(self, **kwargs):
+        data = {}
+        data['token'] = self.token
+        return data

mage_ai/data_preparation/models/global_hooks/constants.py

@@ -0,0 +1,44 @@
+from enum import Enum
+
+from mage_ai.authentication.permissions.constants import EntityName
+
+DISABLED_RESOURCE_TYPES = [
+    EntityName.ALL,
+    EntityName.ALL_EXCEPT_RESERVED,
+]
+
+RESTRICTED_RESOURCE_TYPES = [
+    EntityName.File,
+    EntityName.FileContent,
+    EntityName.Folder,
+    EntityName.GlobalHook,
+    EntityName.Oauth,
+    EntityName.OauthAccessToken,
+    EntityName.OauthApplication,
+    EntityName.Permission,
+    EntityName.Role,
+    EntityName.RolePermission,
+    EntityName.Secret,
+    EntityName.Session,
+    EntityName.User,
+    EntityName.UserRole,
+    EntityName.Workspace,
+]
+
+RESOURCE_TYPES = [en for en in EntityName if en not in DISABLED_RESOURCE_TYPES]
+
+
+class HookOutputKey(str, Enum):
+    ERROR = 'error'
+    META = 'meta'
+    METADATA = 'metadata'
+    PAYLOAD = 'payload'
+    QUERY = 'query'
+    RESOURCE = 'resource'
+    RESOURCES = 'resources'
+
+
+VALID_KEYS_FOR_INPUT_OUTPUT_DATA_RESTRICTED = [key.value for key in HookOutputKey]
+VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED = ['hook']
+VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL = \
+    VALID_KEYS_FOR_INPUT_OUTPUT_DATA_RESTRICTED + VALID_KEYS_FOR_INPUT_OUTPUT_DATA_UNRESTRICTED