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

mage_ai/io/snowflake.py

@@ -11,6 +11,7 @@ from mage_ai.data_preparation.models.block.sql.utils.shared import (
 )
 from mage_ai.io.base import QUERY_ROW_LIMIT, BaseSQLConnection, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
+from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
 from mage_ai.shared.hash import merge_dict
 
 DEFAULT_LOGIN_TIMEOUT = 20
@@ -201,6 +202,8 @@ class Snowflake(BaseSQLConnection):
         schema: str = None,
         if_exists: str = 'append',
         query_string: Union[str, None] = None,
+        unique_conflict_method: str = None,
+        unique_constraints: List[str] = None,
         verbose: bool = True,
         **kwargs,
     ) -> None:
@@ -265,37 +268,28 @@ class Snowflake(BaseSQLConnection):
                         cur.execute(f'DROP TABLE "{schema}"."{table_name}"', timeout=self.timeout)
                         should_create_table = True
 
-                if query_string:
-                    cur.execute(f'USE DATABASE {database}', timeout=self.timeout)
-                    cur.execute(f'USE SCHEMA {schema}', timeout=self.timeout)
-
-                    if should_create_table:
-                        cur.execute(f"""
-CREATE TABLE IF NOT EXISTS "{database}"."{schema}"."{table_name}" AS
-{query_string}
-""", timeout=self.timeout)
-                    else:
-                        cur.execute(f"""
-INSERT INTO "{database}"."{schema}"."{table_name}"
-{query_string}
-""", timeout=self.timeout)
-
-                else:
-                    write_kwargs = merge_dict(
-                        dict(
-                            auto_create_table=should_create_table,
-                            database=database,
-                            schema=schema,
-                            # This param makes sure datetime column is written correctly
-                            use_logical_type=True,
-                        ),
-                        kwargs or dict(),
-                    )
-                    write_pandas(
-                        self.conn,
+                if unique_constraints and unique_conflict_method and df is not None:
+                    self.__upsert_df_into_table(
+                        table_name,
                         df,
+                        cur,
+                        database,
+                        schema,
+                        should_create_table=should_create_table,
+                        unique_conflict_method=unique_conflict_method,
+                        unique_constraints=unique_constraints,
+                        **kwargs,
+                    )
+                else:
+                    self.__write_table(
                         table_name,
-                        **write_kwargs,
+                        df,
+                        cur,
+                        database,
+                        schema,
+                        should_create_table=should_create_table,
+                        query_string=query_string,
+                        **kwargs,
                     )
 
         if verbose:
@@ -306,6 +300,135 @@ INSERT INTO "{database}"."{schema}"."{table_name}"
         else:
             __process()
 
+    def __upsert_df_into_table(
+        self,
+        table_name: str,
+        df: DataFrame,
+        cursor,
+        database: str,
+        schema: str,
+        should_create_table: bool = False,
+        unique_conflict_method: str = None,
+        unique_constraints: List[str] = None,
+        allow_reserved_words: bool = True,
+        auto_clean_name: bool = True,
+        case_sensitive: bool = True,
+        **kwargs
+    ):
+        write_kwargs = merge_dict(
+            dict(
+                auto_create_table=True,
+                database=database,
+                schema=schema,
+                # This param makes sure datetime column is written correctly
+                use_logical_type=True,
+            ),
+            kwargs or dict(),
+        )
+        # should_create_table is True when the table does not exist, so just create the
+        # table as normal.
+        if should_create_table:
+            write_pandas(
+                self.conn,
+                df,
+                table_name,
+                **write_kwargs,
+            )
+        else:
+            temp_table_name = f'temp_{table_name}'
+            write_pandas(
+                self.conn,
+                df,
+                temp_table_name,
+                table_type='temp',
+                **write_kwargs,
+            )
+
+            cleaned_unique_constraints = []
+            for col in unique_constraints:
+                cleaned_col = self._clean_column_name(
+                    col,
+                    allow_reserved_words=allow_reserved_words,
+                    auto_clean_name=auto_clean_name,
+                    case_sensitive=case_sensitive,
+                )
+                cleaned_unique_constraints.append(f'"{cleaned_col}"')
+
+            cleaned_columns = []
+            for col in df.columns:
+                cleaned_col = self._clean_column_name(
+                    col,
+                    allow_reserved_words=allow_reserved_words,
+                    auto_clean_name=auto_clean_name,
+                    case_sensitive=case_sensitive,
+                )
+                cleaned_columns.append(f'"{cleaned_col}"')
+
+            merge_commands = [
+                f'MERGE INTO "{database}"."{schema}"."{table_name}" AS a',
+                f'USING (SELECT * FROM "{database}"."{schema}"."{temp_table_name}") AS b',
+                f"ON {' AND '.join([f'a.{col} = b.{col}' for col in cleaned_unique_constraints])}",
+            ]
+
+            if unique_conflict_method == UNIQUE_CONFLICT_METHOD_UPDATE:
+                set_command = ', '.join(
+                    [f'a.{col} = b.{col}' for col in cleaned_columns],
+                )
+                merge_commands.append(f'WHEN MATCHED THEN UPDATE SET {set_command}')
+
+            insert_columns = ', '.join(cleaned_columns)
+            merge_values = f"({', '.join([f'b.{col}' for col in cleaned_columns])})"
+            merge_commands.append(
+                f"WHEN NOT MATCHED THEN INSERT ({insert_columns}) VALUES {merge_values}",
+            )
+            merge_command = '\n'.join(merge_commands)
+
+            cursor.execute(merge_command, timeout=self.timeout)
+
+    def __write_table(
+        self,
+        table_name: str,
+        df: DataFrame,
+        cursor,
+        database: str,
+        schema: str,
+        should_create_table: bool = False,
+        query_string: str = None,
+        **kwargs
+    ):
+        if query_string:
+            cursor.execute(f'USE DATABASE {database}', timeout=self.timeout)
+            cursor.execute(f'USE SCHEMA {schema}', timeout=self.timeout)
+
+            if should_create_table:
+                cursor.execute(f"""
+CREATE TABLE IF NOT EXISTS "{database}"."{schema}"."{table_name}" AS
+{query_string}
+""", timeout=self.timeout)
+            else:
+                cursor.execute(f"""
+INSERT INTO "{database}"."{schema}"."{table_name}"
+{query_string}
+""", timeout=self.timeout)
+
+        else:
+            write_kwargs = merge_dict(
+                dict(
+                    auto_create_table=should_create_table,
+                    database=database,
+                    schema=schema,
+                    # This param makes sure datetime column is written correctly
+                    use_logical_type=True,
+                ),
+                kwargs or dict(),
+            )
+            write_pandas(
+                self.conn,
+                df,
+                table_name,
+                **write_kwargs,
+            )
+
     @classmethod
     def with_config(
         cls,

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/oauth.py

@@ -1,7 +1,7 @@
 import enum
 import re
 from datetime import datetime
-from typing import Dict, List, Union
+from typing import Callable, Dict, List, Union
 
 from sqlalchemy import (
     JSON,
@@ -327,7 +327,7 @@ class Role(BaseModel):
         self,
         entity: Entity = None,
         entity_id: str = None,
-        prefix: str = None,
+        name_func: Callable[[str], str] = None,
     ) -> None:
         """
         Create default roles with associated permissions for a given entity and entity_id.
@@ -349,8 +349,8 @@ class Role(BaseModel):
         }
         for name, access in mapping.items():
             role_name = name
-            if prefix is not None:
-                role_name = f'{prefix}_{name}'
+            if name_func is not None:
+                role_name = name_func(name)
             role = self.query.filter(self.name == role_name).first()
             if not role:
                 self.create(

mage_ai/orchestration/db/models/schedules.py

@@ -48,6 +48,7 @@ from mage_ai.data_preparation.models.constants import (
     PipelineType,
 )
 from mage_ai.data_preparation.models.pipeline import Pipeline
+from mage_ai.data_preparation.models.project import Project
 from mage_ai.data_preparation.models.triggers import (
     ScheduleInterval,
     ScheduleStatus,
@@ -934,9 +935,12 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
 
     @property
     def pipeline_tags(self):
-        pipeline = Pipeline.get(self.pipeline_uuid, check_if_exists=True)
+        try:
+            pipeline_config = Pipeline.get_config(self.pipeline_uuid)
+        except Exception:
+            pipeline_config = dict()
 
-        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,
@@ -989,6 +993,9 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
         block_runs_all = []
 
         data_integration_block_uuids_mapping = {}
+
+        pipeline_project = Project(pipeline.repo_config)
+
         for block_run in self.block_runs:
             block_runs_all.append(block_run)
 
@@ -1021,7 +1028,7 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
                 "original": 1,
             }
             """
-            if metrics and block and block.is_data_integration():
+            if metrics and block and block.is_data_integration(pipeline_project=pipeline_project):
                 original_block_uuid = metrics.get('original_block_uuid')
 
                 if original_block_uuid and metrics.get('child'):

mage_ai/orchestration/job_manager.py

@@ -86,6 +86,12 @@ class JobManager:
         job_id = self.__job_id(JobType.INTEGRATION_STREAM, id)
         return self.queue.kill_job(job_id)
 
+    def start(self):
+        self.queue.start()
+
+    def stop(self):
+        self.queue.stop()
+
     def __job_id(self, job_type: JobType, uid: Union[str, int]):
         return f'{job_type}_{uid}'
 

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)
@@ -821,6 +830,8 @@ class PipelineScheduler:
         Returns:
             List[BlockRun]: A list of crashed block runs.
         """
+        for b in self.pipeline_run.block_runs:
+            b.refresh()
         running_or_queued_block_runs = [b for b in self.pipeline_run.block_runs if b.status in [
             BlockRun.BlockRunStatus.RUNNING,
             BlockRun.BlockRunStatus.QUEUED,
@@ -1670,7 +1681,11 @@ def gen_pipeline_with_schedules_single_project(
     # Iterate through pipeline schedules by pipeline to handle pipeline run limits for
     # each pipeline.
     for pipeline_uuid, active_schedules in pipeline_schedules_by_pipeline.items():
-        pipeline = Pipeline.get(pipeline_uuid)
+        try:
+            pipeline = Pipeline.get(pipeline_uuid)
+        except Exception as e:
+            print(f'Error fetching pipeline {pipeline_uuid}: {e}')
+            continue
         yield pipeline_uuid, pipeline, active_schedules
 
 
@@ -1731,11 +1746,15 @@ def gen_pipeline_with_schedules_project_platform(
     for pair in pipeline_schedules_by_pipeline_by_repo_path.items():
         repo_path, pipeline_schedules_by_pipeline = pair
         for pipeline_uuid, active_schedules in pipeline_schedules_by_pipeline.items():
-            pipeline = get_pipeline_from_platform(
-                pipeline_uuid,
-                repo_path=repo_path,
-                mapping=pipeline_schedule_repo_paths_to_repo_path_mapping,
-            )
+            try:
+                pipeline = get_pipeline_from_platform(
+                    pipeline_uuid,
+                    repo_path=repo_path,
+                    mapping=pipeline_schedule_repo_paths_to_repo_path_mapping,
+                )
+            except Exception as e:
+                print(f'Error fetching pipeline {pipeline_uuid}: {e}')
+                continue
             yield pipeline_uuid, pipeline, active_schedules
 
 

mage_ai/orchestration/queue/celery_queue.py

@@ -1,6 +1,7 @@
+from typing import Callable
+
 from mage_ai.orchestration.queue.config import QueueConfig
 from mage_ai.orchestration.queue.queue import Queue
-from typing import Callable
 
 
 class CeleryQueue(Queue):
@@ -27,3 +28,9 @@ class CeleryQueue(Queue):
 
     def kill_job(self, job_id: str):
         pass
+
+    def start(self):
+        pass
+
+    def stop(self):
+        pass

mage_ai/orchestration/queue/process_queue.py

@@ -36,6 +36,11 @@ class JobStatus(str, Enum):
     CANCELLED = 'cancelled'
 
 
+class QueueStatus(str, Enum):
+    ACTIVE = 'active'
+    INACTIVE = 'inactive'
+
+
 class ProcessQueue(Queue):
     def __init__(self, queue_config: QueueConfig):
         """
@@ -53,6 +58,7 @@ class ProcessQueue(Queue):
                 jobs.
 
         """
+        self.status = QueueStatus.INACTIVE
         self.queue_config = queue_config
         self.process_queue_config = self.queue_config.process_queue_config
         self.queue = mp.Queue()
@@ -76,12 +82,16 @@ class ProcessQueue(Queue):
 
     def clean_up_jobs(self):
         """
-        Cleans up completed jobs from the job dictionary.
+        1. Cleans up completed jobs from the job dictionary.
+        2. Check whether there're jobs need to be killed.
         """
         job_ids = self.job_dict.keys()
         for job_id in job_ids:
-            if job_id in self.job_dict and not self.has_job(job_id):
-                del self.job_dict[job_id]
+            if job_id in self.job_dict:
+                if not self.has_job(job_id):
+                    del self.job_dict[job_id]
+                elif self.__should_kill_job(job_id):
+                    self.kill_job(job_id)
 
     def enqueue(self, job_id: str, target: Callable, *args, **kwargs):
         """
@@ -94,6 +104,9 @@ class ProcessQueue(Queue):
             **kwargs: Keyword arguments for the target function.
 
         """
+        if self.status != QueueStatus.ACTIVE:
+            self._print('Cannot enqueue a job to an inactive queue.')
+            return
         if self.has_job(job_id):
             self._print(f'Job {job_id} exists. Skip enqueue.')
             return
@@ -155,6 +168,7 @@ class ProcessQueue(Queue):
         print(f'Kill job {job_id}, job_dict {self.job_dict}')
         job = self.job_dict.get(job_id)
         if not job:
+            self.__set_kill_job(job_id)
             return
         if isinstance(job, int):
             if job == os.getpid():
@@ -165,6 +179,7 @@ class ProcessQueue(Queue):
             except Exception as err:
                 print(err)
         self.job_dict[job_id] = JobStatus.CANCELLED
+        self.__unset_kill_job(job_id)
 
     def start_worker_pool(self):
         """
@@ -182,6 +197,28 @@ class ProcessQueue(Queue):
         )
         self.worker_pool_proc.start()
 
+    def start(self):
+        self.status = QueueStatus.ACTIVE
+
+    def stop(self):
+        """
+        1. Stop enqueueing new jobs
+        2. Clear the queue
+        3. Kill all the running jobs
+        """
+        self.status = QueueStatus.INACTIVE
+        while not self.queue.empty():
+            try:
+                self.queue.get_nowait()
+            except self.queue.Empty:
+                break
+        job_ids = self.job_dict.keys()
+        for job_id in job_ids:
+            if job_id in self.job_dict:
+                if isinstance(self.job_dict.get(job_id), int):
+                    self.kill_job(job_id)
+                del self.job_dict[job_id]
+
     def is_worker_pool_alive(self) -> bool:
         """
         Checks if the worker pool process is alive.
@@ -197,6 +234,31 @@ class ProcessQueue(Queue):
     def __is_process_alive(self, pid: int) -> bool:
         return psutil.pid_exists(pid)
 
+    def __redis_key_kill_job(self, job_id):
+        return f'kill_job_{job_id}'
+
+    def __set_kill_job(self, job_id):
+        if not self.redis_client:
+            return
+        return self.redis_client.set(
+            self.__redis_key_kill_job(job_id),
+            '1',
+            ex=LIVENESS_TIMEOUT_SECONDS,
+        )
+
+    def __unset_kill_job(self, job_id):
+        if not self.redis_client:
+            return
+        key = self.__redis_key_kill_job(job_id)
+        if self.redis_client.get(key):
+            self.redis_client.delete(key)
+
+    def __should_kill_job(self, job_id):
+        if not self.redis_client:
+            return False
+        value = self.redis_client.get(self.__redis_key_kill_job(job_id))
+        return value is not None
+
 
 class Worker(mp.Process):
     def __init__(
@@ -275,10 +337,11 @@ def poll_job_and_execute(
         job_dict: The shared job dictionary.
 
     """
+    pid = os.getpid()
     workers = []
     while True:
         workers = [w for w in workers if w.is_alive()]
-        print(f'Worker pool size: {len(workers)}')
+        print(f'[Process {pid}] Worker pool size: {len(workers)}')
         if not workers and queue.empty():
             break
         while not queue.empty():

mage_ai/orchestration/queue/queue.py

@@ -19,5 +19,13 @@ class Queue(ABC):
     def kill_job(self, job_id: str):
         pass
 
+    @abstractmethod
+    def start(self):
+        pass
+
+    @abstractmethod
+    def stop(self):
+        pass
+
     def _print(self, msg):
         print(f'[{self.__class__.__name__}] {msg}')

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.69'