mage-ai 0.9.67__py3-none-any.whl → 0.9.68__py3-none-any.whl

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

@@ -416,9 +416,10 @@ class SourceBlock(IntegrationBlock):
 class DestinationBlock(IntegrationBlock):
     def to_dict(
         self,
-        include_content=False,
-        include_outputs=False,
-        sample_count=None,
+        include_content: bool = False,
+        include_outputs: bool = False,
+        include_block_pipelines: bool = False,
+        sample_count: int = None,
         check_if_file_exists: bool = False,
         destination_table: str = None,
         state_stream: str = None,
@@ -444,6 +445,7 @@ class DestinationBlock(IntegrationBlock):
             super().to_dict(
                 include_content=include_content,
                 include_outputs=include_outputs,
+                include_block_pipelines=include_block_pipelines,
                 sample_count=sample_count,
                 check_if_file_exists=check_if_file_exists,
             ),
@@ -452,9 +454,10 @@ class DestinationBlock(IntegrationBlock):
 
     async def to_dict_async(
         self,
-        include_content=False,
-        include_outputs=False,
-        sample_count=None,
+        include_content: bool = False,
+        include_outputs: bool = False,
+        include_block_pipelines: bool = False,
+        sample_count: int = None,
         check_if_file_exists: bool = False,
         destination_table: str = None,
         state_stream: str = None,
@@ -463,13 +466,14 @@ class DestinationBlock(IntegrationBlock):
         return self.to_dict(
             include_content=include_content,
             include_outputs=include_outputs,
+            include_block_pipelines=include_block_pipelines,
             sample_count=sample_count,
             check_if_file_exists=check_if_file_exists,
             destination_table=destination_table,
             state_stream=state_stream,
         )
 
-    def update(self, data, update_state=False):
+    def update(self, data, update_state=False, **kwargs):
         if update_state:
             from mage_ai.data_preparation.models.pipelines.integration_pipeline import (
                 IntegrationPipeline,
@@ -493,7 +497,7 @@ class DestinationBlock(IntegrationBlock):
                     bookmark_values=bookmark_values
                 )
 
-        return super().update(data)
+        return super().update(data, **kwargs)
 
     def output_variables(self, execution_partition: str = None) -> List[str]:
         return []

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

@@ -233,7 +233,7 @@ class ProjectPlatformAccessible:
                 ),
             )
 
-        return block_class(
+        return block_class.create(
             name,
             uuid,
             block_type,

mage_ai/data_preparation/models/pipeline.py

@@ -15,8 +15,10 @@ import yaml
 from jinja2 import Template
 
 from mage_ai.authentication.permissions.constants import EntityName
+from mage_ai.cache.block import BlockCache
 from mage_ai.cache.pipeline import PipelineCache
 from mage_ai.data_preparation.models.block import Block, run_blocks, run_blocks_sync
+from mage_ai.data_preparation.models.block.block_factory import BlockFactory
 from mage_ai.data_preparation.models.block.data_integration.utils import (
     convert_outputs_to_data,
 )
@@ -82,6 +84,8 @@ class Pipeline:
         repo_config=None,
         catalog=None,
         use_repo_path: bool = False,
+        description: str = None,
+        tags: List[str] = None,
     ):
         self.block_configs = []
         self.blocks_by_uuid = {}
@@ -90,7 +94,7 @@ class Pipeline:
         self.concurrency_config = dict()
         self.created_at = None
         self.data_integration = None
-        self.description = None
+        self.description = description
         self.executor_config = dict()
         self.executor_type = None
         self.extensions = {}
@@ -101,7 +105,7 @@ class Pipeline:
         self.run_pipeline_in_one_process = False
         self.schedules = []
         self.settings = {}
-        self.tags = []
+        self.tags = tags or []
         self.type = PipelineType.PYTHON
         self.use_repo_path = use_repo_path
         self.uuid = uuid
@@ -215,7 +219,14 @@ class Pipeline:
             self.widget_configs
 
     @classmethod
-    def create(self, name, pipeline_type=PipelineType.PYTHON, repo_path=None):
+    def create(
+        self,
+        name: str,
+        description: str = None,
+        pipeline_type: PipelineType = PipelineType.PYTHON,
+        repo_path: str = None,
+        tags: List[str] = None,
+    ):
         """
         1. Create a new folder for pipeline
         2. Create a new yaml file to store pipeline config
@@ -231,14 +242,20 @@ class Pipeline:
         with open(os.path.join(pipeline_path, PIPELINE_CONFIG_FILE), 'w') as fp:
             yaml.dump(dict(
                 created_at=str(datetime.now(tz=pytz.UTC)),
+                description=description,
                 name=name,
+                tags=tags or [],
                 uuid=uuid,
                 type=format_enum(pipeline_type or PipelineType.PYTHON),
             ), fp)
+
         pipeline = Pipeline(
             uuid,
+            description=description,
             repo_path=repo_path,
+            tags=tags or [],
         )
+
         return pipeline
 
     @classmethod
@@ -368,6 +385,54 @@ class Pipeline:
             IntegrationPipeline,
         )
 
+        config_path, repo_path = self._get_config_path(
+            uuid,
+            repo_path=repo_path,
+            all_projects=all_projects,
+            use_repo_path=use_repo_path,
+        )
+
+        if check_if_exists and not os.path.exists(config_path):
+            return None
+
+        pipeline = self(uuid, repo_path=repo_path, use_repo_path=use_repo_path)
+        if PipelineType.INTEGRATION == pipeline.type:
+            pipeline = IntegrationPipeline(uuid, repo_path=repo_path)
+
+        return pipeline
+
+    @classmethod
+    def get_config(
+        self,
+        uuid,
+        repo_path: str = None,
+        all_projects: bool = False,
+        use_repo_path: bool = False,
+    ):
+        config_path, _ = self._get_config_path(
+            uuid,
+            repo_path=repo_path,
+            all_projects=all_projects,
+            use_repo_path=use_repo_path,
+        )
+
+        metadata_path = os.path.join(config_path, PIPELINE_CONFIG_FILE)
+
+        if not os.path.exists(metadata_path):
+            return None
+
+        with open(metadata_path) as fp:
+            config = yaml.full_load(fp) or {}
+        return config
+
+    @classmethod
+    def _get_config_path(
+        self,
+        uuid,
+        repo_path: str = None,
+        all_projects: bool = False,
+        use_repo_path: bool = False,
+    ) -> Tuple[str, str]:
         if all_projects and not use_repo_path and project_platform_activated():
             from mage_ai.settings.platform.utils import get_pipeline_config_path
 
@@ -379,15 +444,7 @@ class Pipeline:
                 PIPELINES_FOLDER,
                 uuid,
             )
-
-        if check_if_exists and not os.path.exists(config_path):
-            return None
-
-        pipeline = self(uuid, repo_path=repo_path, use_repo_path=use_repo_path)
-        if PipelineType.INTEGRATION == pipeline.type:
-            pipeline = IntegrationPipeline(uuid, repo_path=repo_path)
-
-        return pipeline
+        return config_path, repo_path
 
     @classmethod
     async def load_metadata(
@@ -717,7 +774,8 @@ class Pipeline:
                 )
 
             language = c.get('language')
-            return Block.block_class_from_type(block_type, language=language, pipeline=self)(
+
+            return BlockFactory.block_class_from_type(block_type, language=language, pipeline=self)(
                 c.get('name'),
                 c.get('uuid'),
                 block_type,
@@ -931,6 +989,8 @@ class Pipeline:
             include_outputs=include_outputs,
             sample_count=sample_count,
         )
+        if include_block_pipelines:
+            shared_kwargs['block_cache'] = BlockCache()
         blocks_data = await asyncio.gather(
             *[b.to_dict_async(**merge_dict(shared_kwargs, dict(
                 include_block_catalog=include_block_catalog,
@@ -991,6 +1051,7 @@ class Pipeline:
         old_uuid = None
         blocks_to_remove_from_cache = []
         block_uuids_to_add_to_cache = []
+        tags_to_remove_from_cache = []
         should_update_block_cache = False
         should_update_tag_cache = False
 
@@ -1043,6 +1104,10 @@ class Pipeline:
             old_tags = self.tags or []
 
             if sorted(new_tags) != sorted(old_tags):
+                tags_diff = set(old_tags).symmetric_difference(set(new_tags))
+                for tag in tags_diff:
+                    if tag in old_tags:
+                        tags_to_remove_from_cache.append(tag)
                 self.tags = new_tags
                 should_save = True
                 should_update_tag_cache = True
@@ -1296,6 +1361,8 @@ class Pipeline:
 
             cache = await TagCache.initialize_cache()
 
+            for tag_uuid in tags_to_remove_from_cache:
+                cache.remove_pipeline(tag_uuid, self.uuid, self.repo_path)
             for tag_uuid in self.tags:
                 if old_uuid:
                     cache.remove_pipeline(tag_uuid, old_uuid, self.repo_path)

mage_ai/data_preparation/storage/gcs_storage.py

@@ -148,7 +148,7 @@ class GCSStorage(BaseStorage):
         df.write_parquet(buffer)
         buffer.seek(0)
         blob = self.bucket.blob(gcs_url_path(file_path))
-        blob.upload_from_string(buffer)
+        blob.upload_from_string(buffer.getvalue())
 
     @contextmanager
     def open_to_write(self, file_path: str) -> None:

mage_ai/data_preparation/templates/constants.py

@@ -608,6 +608,13 @@ TEMPLATES_ONLY_FOR_V2 = [
         name='MySQL',
         path='data_exporters/mysql.py',
     ),
+    dict(
+        block_type=BlockType.DATA_EXPORTER,
+        groups=[GROUP_DATABASES],
+        language=BlockLanguage.PYTHON,
+        name='OracleDB',
+        path='data_exporters/oracledb.py',
+    ),
     dict(
         block_type=BlockType.DATA_EXPORTER,
         groups=[GROUP_DATABASES],

mage_ai/io/export_utils.py

@@ -108,6 +108,7 @@ def gen_table_creation_query(
     case_sensitive: bool = False,
     unique_constraints: List[str] = None,
     overwrite_types: Dict = None,
+    skip_semicolon_at_end: bool = False,
 ) -> str:
     """
     Generates a database table creation query from a data frame.
@@ -157,4 +158,6 @@ def gen_table_creation_query(
         query.append(
             f"CONSTRAINT {index_name} UNIQUE ({', '.join(unique_constraints_escaped)})",
         )
+    if skip_semicolon_at_end:
+        return f'CREATE TABLE {full_table_name} (' + ','.join(query) + ')'
     return f'CREATE TABLE {full_table_name} (' + ','.join(query) + ');'

mage_ai/io/oracledb.py

@@ -1,13 +1,17 @@
 import warnings
-from typing import Union
+from typing import IO, Any, List, Union
 
+import numpy as np
 import oracledb
-from pandas import DataFrame, read_sql
+import simplejson
+from pandas import DataFrame, Series, read_sql
 
-from mage_ai.io.base import QUERY_ROW_LIMIT
+from mage_ai.io.base import QUERY_ROW_LIMIT, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.export_utils import PandasTypes
 from mage_ai.io.sql import BaseSQL
 from mage_ai.server.logger import Logger
+from mage_ai.shared.parsers import encode_complex
 
 logger = Logger().new_server_logger(__name__)
 
@@ -122,3 +126,134 @@ SELECT *
 FROM subquery
 FETCH FIRST {limit} ROWS ONLY
                 """
+
+    def table_exists(self, schema_name: str, table_name: str) -> bool:
+        with self.conn.cursor() as cur:
+            try:
+                cur.execute(f"select * from {table_name} where rownum=1")
+            except Exception as exc:
+                logger.info(f"Table not existing: {table_name}. Exception: {exc}")
+                return False
+        return True
+
+    def get_type(self, column: Series, dtype: str) -> str:
+        if dtype in (
+            PandasTypes.MIXED,
+            PandasTypes.UNKNOWN_ARRAY,
+            PandasTypes.COMPLEX,
+        ):
+            return 'CHAR(255)'
+        elif dtype in (PandasTypes.DATETIME, PandasTypes.DATETIME64):
+            try:
+                if column.dt.tz:
+                    return 'TIMESTAMP'
+            except AttributeError:
+                pass
+            return 'TIMESTAMP'
+        elif dtype == PandasTypes.TIME:
+            try:
+                if column.dt.tz:
+                    return 'TIMESTAMP'
+            except AttributeError:
+                pass
+            return 'TIMESTAMP'
+        elif dtype == PandasTypes.DATE:
+            return 'DATE'
+        elif dtype == PandasTypes.STRING:
+            return 'CHAR(255)'
+        elif dtype == PandasTypes.CATEGORICAL:
+            return 'CHAR(255)'
+        elif dtype == PandasTypes.BYTES:
+            return 'CHAR(255)'
+        elif dtype in (PandasTypes.FLOATING, PandasTypes.DECIMAL, PandasTypes.MIXED_INTEGER_FLOAT):
+            return 'NUMBER'
+        elif dtype == PandasTypes.INTEGER:
+            max_int, min_int = column.max(), column.min()
+            if np.int16(max_int) == max_int and np.int16(min_int) == min_int:
+                return 'NUMBER'
+            elif np.int32(max_int) == max_int and np.int32(min_int) == min_int:
+                return 'NUMBER'
+            else:
+                return 'NUMBER'
+        elif dtype == PandasTypes.BOOLEAN:
+            return 'CHAR(52)'
+        elif dtype in (PandasTypes.TIMEDELTA, PandasTypes.TIMEDELTA64, PandasTypes.PERIOD):
+            return 'NUMBER'
+        elif dtype == PandasTypes.EMPTY:
+            return 'CHAR(255)'
+        else:
+            print(f'Invalid datatype provided: {dtype}')
+
+        return 'CHAR(255)'
+
+    def export(
+        self,
+        df: DataFrame,
+        table_name: str = None,
+        if_exists: ExportWritePolicy = ExportWritePolicy.REPLACE,
+        **kwargs,
+    ) -> None:
+        super().export(
+            df,
+            **kwargs,
+            table_name=table_name,
+            if_exists=if_exists,
+            # Oracle cursor execute will automatically add a semicolon at the end of the query.
+            skip_semicolon_at_end=True
+        )
+
+    def upload_dataframe(
+        self,
+        cursor: Any,
+        df: DataFrame,
+        db_dtypes: List[str],
+        dtypes: List[str],
+        full_table_name: str,
+        buffer: Union[IO, None] = None,
+        **kwargs,
+    ) -> None:
+        def serialize_obj(val):
+            if type(val) is dict or type(val) is np.ndarray:
+                return simplejson.dumps(
+                    val,
+                    default=encode_complex,
+                    ignore_nan=True,
+                )
+            elif type(val) is list and len(val) >= 1 and type(val[0]) is dict:
+                return simplejson.dumps(
+                    val,
+                    default=encode_complex,
+                    ignore_nan=True,
+                )
+            return val
+
+        # Create values
+        df_ = df.copy()
+        columns = df_.columns
+        for col in columns:
+            dtype = df_[col].dtype
+            if dtype == PandasTypes.OBJECT:
+                df_[col] = df_[col].apply(lambda x: serialize_obj(x))
+            elif dtype in (
+                PandasTypes.MIXED,
+                PandasTypes.UNKNOWN_ARRAY,
+                PandasTypes.COMPLEX,
+            ):
+                df_[col] = df_[col].astype('string')
+
+            # 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)
+        values = []
+        for i in range(0, len(df_)):
+            values.append(tuple(df_.fillna('').values[i]))
+
+        # Create values placeholder
+        colmn_names = df.columns.tolist()
+        values_placeholder = ""
+        for i in range(0, len(colmn_names)):
+            values_placeholder += f':{str(i + 1)},'
+
+        insert_sql = f'INSERT INTO {full_table_name} VALUES({values_placeholder.rstrip(",")})'
+        cursor.executemany(insert_sql, values)

mage_ai/io/sql.py

@@ -56,6 +56,7 @@ class BaseSQL(BaseSQLConnection):
         case_sensitive: bool = False,
         unique_constraints: List[str] = None,
         overwrite_types: Dict = None,
+        skip_semicolon_at_end: bool = False,
         **kwargs,
     ) -> str:
         if unique_constraints is None:
@@ -68,6 +69,7 @@ class BaseSQL(BaseSQLConnection):
             case_sensitive=case_sensitive,
             unique_constraints=unique_constraints,
             overwrite_types=overwrite_types,
+            skip_semicolon_at_end=skip_semicolon_at_end,
         )
 
     def build_create_table_as_command(
@@ -234,6 +236,7 @@ class BaseSQL(BaseSQLConnection):
         query_string: Union[str, None] = None,
         unique_conflict_method: str = None,
         unique_constraints: List[str] = None,
+        skip_semicolon_at_end: bool = False,
         **kwargs,
     ) -> None:
         """
@@ -350,6 +353,7 @@ class BaseSQL(BaseSQLConnection):
                             case_sensitive=case_sensitive,
                             unique_constraints=unique_constraints,
                             overwrite_types=overwrite_types,
+                            skip_semicolon_at_end=skip_semicolon_at_end,
                         )
                         cur.execute(query)
                     self.upload_dataframe(

mage_ai/orchestration/db/__init__.py

@@ -13,7 +13,7 @@ from mage_ai.orchestration.db.setup import get_postgres_connection_url
 from mage_ai.orchestration.db.utils import get_user_info_from_db_connection_url
 from mage_ai.settings import OTEL_EXPORTER_OTLP_ENDPOINT
 from mage_ai.settings.repo import get_variables_dir
-from mage_ai.shared.environments import is_dev, is_test
+from mage_ai.shared.environments import is_debug, is_test
 
 DB_RETRY_COUNT = 2
 TEST_DB = 'test.db'
@@ -170,5 +170,5 @@ def safe_db_query(func):
 
 logging.basicConfig()
 
-if is_dev() and not os.getenv('DISABLE_DATABASE_TERMINAL_OUTPUT'):
+if is_debug() and not os.getenv('DISABLE_DATABASE_TERMINAL_OUTPUT'):
     logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)

mage_ai/orchestration/db/models/schedules.py

@@ -934,9 +934,9 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
 
     @property
     def pipeline_tags(self):
-        pipeline = Pipeline.get(self.pipeline_uuid, check_if_exists=True)
+        pipeline_config = Pipeline.get_config(self.pipeline_uuid)
 
-        return pipeline.tags if pipeline is not None else []
+        return pipeline_config.get('tags') if pipeline_config is not None else []
 
     def executable_block_runs(
         self,

mage_ai/orchestration/notification/sender.py

@@ -108,6 +108,7 @@ class NotificationSender:
         pipeline,
         pipeline_run,
         error: str = None,
+        stacktrace: str = None,
         summary: str = None,
     ) -> None:
         if AlertOn.PIPELINE_RUN_FAILURE in self.config.alert_on:
@@ -122,6 +123,7 @@ class NotificationSender:
                 pipeline_run,
                 error=error,
                 message_template=message_template,
+                stacktrace=stacktrace,
                 summary=summary,
             )
 
@@ -145,6 +147,7 @@ class NotificationSender:
         pipeline,
         pipeline_run,
         error: str = None,
+        stacktrace: str = None,
     ):
         if text is None or pipeline is None or pipeline_run is None:
             return text
@@ -155,6 +158,7 @@ class NotificationSender:
             pipeline_schedule_id=pipeline_run.pipeline_schedule.id,
             pipeline_schedule_name=pipeline_run.pipeline_schedule.name,
             pipeline_uuid=pipeline.uuid,
+            stacktrace=stacktrace,
         )
 
     def __send_pipeline_run_message(
@@ -164,6 +168,7 @@ class NotificationSender:
         pipeline_run,
         error: str = None,
         message_template: MessageTemplate = None,
+        stacktrace: str = None,
         summary: str = None,
     ):
         """Shared method to send pipeline run message of multiple types (success, failure, etc.).
@@ -207,18 +212,21 @@ class NotificationSender:
                 pipeline,
                 pipeline_run,
                 error=error,
+                stacktrace=stacktrace,
             ),
             summary=self.__interpolate_vars(
                 summary or default_summary,
                 pipeline,
                 pipeline_run,
                 error=error,
+                stacktrace=stacktrace,
             ),
             details=self.__interpolate_vars(
                 details or default_details,
                 pipeline,
                 pipeline_run,
                 error=error,
+                stacktrace=stacktrace,
             ),
         )
 

mage_ai/orchestration/pipeline_scheduler_original.py

@@ -320,11 +320,19 @@ class PipelineScheduler:
     @safe_db_query
     def on_pipeline_run_failure(self, error_msg: str) -> None:
         failed_block_runs = self.pipeline_run.failed_block_runs
+        stacktrace = None
         for br in failed_block_runs:
             if br.metrics:
                 message = br.metrics.get('error', {}).get('message')
                 if message:
-                    error_msg += f'\nError for block {br.block_uuid}:\n{message}'
+                    message_split = message.split('\n')
+                    # Truncate the error message if it has too many lines, set max
+                    # lines at 50
+                    if len(message_split) > 50:
+                        message_split = message_split[-50:]
+                        message_split.insert(0, '... (error truncated)')
+                    message = '\n'.join(message_split)
+                    stacktrace = f'Error for block {br.block_uuid}:\n{message}'
                     break
 
         asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
@@ -332,6 +340,7 @@ class PipelineScheduler:
             pipeline=self.pipeline,
             pipeline_run=self.pipeline_run,
             error=error_msg,
+            stacktrace=stacktrace,
         )
         # Cancel block runs that are still in progress for the pipeline run.
         cancel_block_runs_and_jobs(self.pipeline_run, self.pipeline)

mage_ai/server/constants.py

@@ -12,4 +12,4 @@ DATAFRAME_OUTPUT_SAMPLE_COUNT = 10
 # Dockerfile depends on it because it runs ./scripts/install_mage.sh and uses
 # the last line to determine the version to install.
 VERSION = \
-'0.9.67'
+'0.9.68'