mage-ai 0.9.73__py3-none-any.whl → 0.9.75__py3-none-any.whl

mage_ai/services/compute/models.py

@@ -1,6 +1,5 @@
 from abc import abstractmethod
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Dict, List
 
 from mage_ai.data_preparation.models.project import Project
@@ -10,6 +9,7 @@ from mage_ai.services.compute.constants import (
     ComputeManagementApplicationTab,
 )
 from mage_ai.services.spark.constants import ComputeServiceUUID
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.models import BaseDataClass
 
@@ -35,7 +35,7 @@ class ConnectionCredential(BaseDataClass):
             self.error = ErrorMessage.load(**self.error)
 
 
-class SetupStepStatus(str, Enum):
+class SetupStepStatus(StrEnum):
     COMPLETED = 'completed'
     ERROR = 'error'
     INCOMPLETE = 'incomplete'
@@ -137,7 +137,7 @@ class ComputeConnection(SetupStep):
 
 
 class ComputeService:
-    uuid = ComputeServiceUUID.STANDALONE_CLUSTER
+    service_uuid = ComputeServiceUUID.STANDALONE_CLUSTER
 
     def __init__(self, project: Project, with_clusters: bool = False):
         self.project = project
@@ -169,7 +169,7 @@ class ComputeService:
 
     @property
     def uuid(self) -> ComputeServiceUUID:
-        return self.__class__.uuid
+        return self.__class__.service_uuid
 
     @abstractmethod
     def active_cluster(self, **kwargs) -> Dict:

mage_ai/services/k8s/config.py

@@ -58,6 +58,7 @@ class K8sExecutorConfig(BaseConfig):
         service_account_name = DEFAULT_SERVICE_ACCOUNT_NAME
         affinity = None
         node_selector = None
+        scheduler_name = None
         tolerations = []
         volumes = []
         image_pull_secrets = {}
@@ -71,6 +72,8 @@ class K8sExecutorConfig(BaseConfig):
 
             if executor_config.pod.get('node_selector'):
                 node_selector = executor_config.pod['node_selector']
+            if executor_config.pod.get('scheduler_name'):
+                scheduler_name = executor_config.pod['scheduler_name']
 
             if executor_config.pod.get('tolerations'):
                 tolerations += [V1Toleration(**e) for e in executor_config.pod['tolerations']]
@@ -123,6 +126,7 @@ class K8sExecutorConfig(BaseConfig):
             containers=[container],
             image_pull_secrets=image_pull_secrets,
             node_selector=node_selector,
+            scheduler_name=scheduler_name,
             restart_policy='Never',
             service_account_name=service_account_name,
             tolerations=tolerations,

mage_ai/services/k8s/job_manager.py

@@ -172,6 +172,8 @@ class JobManager():
             pod_spec.tolerations = mage_server_pod_spec.tolerations
         if not pod_spec.node_selector:
             pod_spec.node_selector = mage_server_pod_spec.node_selector
+        if not pod_spec.scheduler_name:
+            pod_spec.scheduler_name = mage_server_pod_spec.scheduler_name
         pod_spec.image_pull_secrets = pod_spec.image_pull_secrets if pod_spec.image_pull_secrets \
             else mage_server_pod_spec.image_pull_secrets
         return pod_spec

mage_ai/services/spark/constants.py

@@ -1,12 +1,12 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 SPARK_DIRECTORY_NAME = '.spark'
 
 
-class ComputeServiceUUID(str, Enum):
+class ComputeServiceUUID(StrEnum):
     AWS_EMR = 'AWS_EMR'
     STANDALONE_CLUSTER = 'STANDALONE_CLUSTER'
 
 
-class SparkMaster(str, Enum):
+class SparkMaster(StrEnum):
     LOCAL = 'local'

mage_ai/services/spark/models/jobs.py

@@ -1,12 +1,12 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Dict, List
 
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.services.spark.models.base import BaseSparkModel
+from mage_ai.shared.enum import StrEnum
 
 
-class JobStatus(str, Enum):
+class JobStatus(StrEnum):
     FAILED = 'FAILED'
     RUNNING = 'RUNNING'
     SUCCEEDED = 'SUCCEEDED'

mage_ai/services/spark/models/sqls.py

@@ -1,14 +1,14 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import List
 
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.services.spark.models.base import BaseSparkModel
 from mage_ai.services.spark.models.jobs import Job
 from mage_ai.services.spark.models.stages import StageAttempt
+from mage_ai.shared.enum import StrEnum
 
 
-class SqlStatus(str, Enum):
+class SqlStatus(StrEnum):
     COMPLETED = 'COMPLETED'
 
 

mage_ai/services/spark/models/stages.py

@@ -1,25 +1,25 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Dict, List
 
 from mage_ai.services.spark.models.applications import Application
 from mage_ai.services.spark.models.base import BaseSparkModel
 from mage_ai.services.spark.models.metrics import Metrics
+from mage_ai.shared.enum import StrEnum
 
 
-class Locality(str, Enum):
+class Locality(StrEnum):
     NODE_LOCAL = 'NODE_LOCAL'
     PROCESS_LOCAL = 'PROCESS_LOCAL'
     RACK_LOCAL = 'RACK_LOCAL'
 
 
-class StageStatus(str, Enum):
+class StageStatus(StrEnum):
     COMPLETE = 'COMPLETE'
     PENDING = 'PENDING'
     SKIPPED = 'SKIPPED'
 
 
-class TaskStatus(str, Enum):
+class TaskStatus(StrEnum):
     SUCCESS = 'SUCCESS'
 
 

mage_ai/services/spark/models/threads.py

@@ -1,11 +1,11 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import List
 
 from mage_ai.services.spark.models.base import BaseSparkModel
+from mage_ai.shared.enum import StrEnum
 
 
-class ThreadState(str, Enum):
+class ThreadState(StrEnum):
     RUNNABLE = 'RUNNABLE'
     TIMED_WAITING = 'TIMED_WAITING'
     WAITING = 'WAITING'

mage_ai/services/teams/config.py

@@ -1,11 +1,22 @@
 from dataclasses import dataclass
+from typing import List
+
 from mage_ai.shared.config import BaseConfig
 
 
 @dataclass
 class TeamsConfig(BaseConfig):
-    webhook_url: str = None
+    webhook_url: List[str] = None
+
+    def __post_init__(self):
+        # Normalize webhook_url to a list if it's provided as a string
+        if isinstance(self.webhook_url, str):
+            self.webhook_url = [self.webhook_url]
+        elif self.webhook_url is None:
+            self.webhook_url = []
 
     @property
     def is_valid(self) -> bool:
-        return self.webhook_url is not None and self.webhook_url != 'None'
+        return bool(self.webhook_url) and all(
+            url and url != 'None' and url.strip() != '' for url in self.webhook_url
+        )

mage_ai/services/teams/teams.py

@@ -1,19 +1,21 @@
-from mage_ai.services.teams.config import TeamsConfig
 import requests
 
+from mage_ai.services.teams.config import TeamsConfig
+
 
 def send_teams_message(
     config: TeamsConfig,
     message: str,
     title: str = 'Mage pipeline run status logs',
 ) -> None:
-    requests.post(
-        url=config.webhook_url,
-        json={
-            'summary': title,
-            'sections': [{
-                'activityTitle': title,
-                'activitySubtitle': message,
-            }],
-        },
-    )
+    for url in config.webhook_url:
+        requests.post(
+            url=url,
+            json={
+                'summary': title,
+                'sections': [{
+                    'activityTitle': title,
+                    'activitySubtitle': message,
+                }],
+            },
+        )

mage_ai/settings/backends.py

@@ -1,13 +1,14 @@
 import base64
 import logging
 import os
-from enum import Enum
 from typing import Optional
 
+from mage_ai.shared.enum import StrEnum
+
 logger = logging.getLogger(__name__)
 
 
-class BackendType(str, Enum):
+class BackendType(StrEnum):
     """
     Enum for the different types of settings backends.
     """

mage_ai/settings/models/configuration_option.py

@@ -1,7 +1,6 @@
 import asyncio
 import os
 from dataclasses import dataclass
-from enum import Enum
 from pathlib import Path
 from typing import Dict, Union
 
@@ -15,6 +14,7 @@ from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.settings.platform import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.files import get_full_file_paths_containing_multi_items
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.models import BaseDataClass
@@ -37,11 +37,11 @@ async def read_file(full_path: str) -> str:
         return config
 
 
-class ConfigurationType(str, Enum):
+class ConfigurationType(StrEnum):
     DBT = 'dbt'
 
 
-class OptionType(str, Enum):
+class OptionType(StrEnum):
     PROFILES = 'profiles'
     PROJECTS = 'projects'
     TARGETS = 'targets'

mage_ai/settings/server.py

@@ -161,6 +161,11 @@ NEW_RELIC_CONFIG_PATH = os.getenv('NEW_RELIC_CONFIG_PATH', '')
 # If enabled, the /metrics route will expose Tornado server metrics
 ENABLE_PROMETHEUS = get_bool_value(os.getenv('ENABLE_PROMETHEUS', 'False'))
 
+# API switch configuration
+# If disabled, those AI APIs will not be used in backend.
+ENABLE_OPEN_AI = get_bool_value(os.getenv('ENABLE_OPEN_AI') or 'True')
+ENABLE_HUGGING_FACE = get_bool_value(os.getenv('ENABLE_HUGGING_FACE') or 'True')
+
 # OpenTelemetry Configuration
 OTEL_EXPORTER_OTLP_ENDPOINT = os.getenv('OTEL_EXPORTER_OTLP_ENDPOINT', None)
 OTEL_EXPORTER_OTLP_HTTP_ENDPOINT = os.getenv('OTEL_EXPORTER_HTTP_OTLP_ENDPOINT', None)

mage_ai/shared/constants.py

@@ -1,15 +1,17 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 ENV_DEV = 'dev'
 ENV_PROD = 'prod'
 ENV_STAGING = 'staging'
 ENV_TEST = 'test'
+ENV_TEST_MAGE = 'test_mage'
 
 VALID_ENVS = frozenset([
     ENV_DEV,
     ENV_PROD,
     ENV_STAGING,
     ENV_TEST,
+    ENV_TEST_MAGE,
 ])
 
 SAMPLE_SIZE = 1000
@@ -20,7 +22,7 @@ GCS_PREFIX = 'gs://'
 ENV_VAR_INSTANCE_TYPE = 'INSTANCE_TYPE'
 
 
-class InstanceType(str, Enum):
+class InstanceType(StrEnum):
     SERVER_AND_SCHEDULER = 'server_and_scheduler'
     SCHEDULER = 'scheduler'
     WEB_SERVER = 'web_server'

mage_ai/shared/custom_logger.py

@@ -1,15 +1,15 @@
 import inspect
 import logging
-from enum import Enum
 
 import simplejson
 
 from mage_ai.data_preparation.models.constants import BlockType
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.environments import is_deus_ex_machina
 from mage_ai.shared.parsers import encode_complex
 
 
-class Color(str, Enum):
+class Color(StrEnum):
     BLUE = "\x1b[1;34m"
     BOLD_RED = "\x1b[31;1m"
     GREEN = "\x1b[1;32m"
@@ -64,19 +64,19 @@ class ColorPrinter:
         self.label = None
 
     def debug(self, *args, **kwargs):
-        self.print(color=Color.PURPLE, *args, **kwargs)
+        self.print(*args, color=Color.PURPLE, **kwargs)
 
     def info(self, *args, **kwargs):
-        self.print(color=Color.BLUE, *args, **kwargs)
+        self.print(*args, color=Color.BLUE, **kwargs)
 
     def warning(self, *args, **kwargs):
-        self.print(color=Color.GREEN, *args, **kwargs)
+        self.print(*args, color=Color.GREEN, **kwargs)
 
     def error(self, *args, **kwargs):
-        self.print(color=Color.RED, *args, **kwargs)
+        self.print(*args, color=Color.RED, **kwargs)
 
     def critical(self, *args, **kwargs):
-        self.print(color=Color.BOLD_RED, *args, **kwargs)
+        self.print(*args, color=Color.BOLD_RED, **kwargs)
 
     def print(self, *args, **kwargs):
         if not is_deus_ex_machina():
@@ -183,39 +183,39 @@ class CustomFormatter(logging.Formatter):
     def debug(self, p: int = 0, *args, **kwargs):
         if is_deus_ex_machina():
             if p:
-                self.print(color=Color.PURPLE, *args, **kwargs)
+                self.print(*args, color=Color.PURPLE, **kwargs)
             else:
                 super().debug(*args, **kwargs)
 
     def info(self, p: int = 0, *args, **kwargs):
         if is_deus_ex_machina():
             if p:
-                self.print(color=Color.BLUE, *args, **kwargs)
+                self.print(*args, color=Color.BLUE, **kwargs)
             else:
                 super().info(*args, **kwargs)
 
     def warning(self, p: int = 0, *args, **kwargs):
         if is_deus_ex_machina():
             if p:
-                self.print(color=Color.YELLOW, *args, **kwargs)
+                self.print(*args, color=Color.YELLOW, **kwargs)
             else:
                 super().warning(*args, **kwargs)
 
     def error(self, p: int = 0, *args, **kwargs):
         if is_deus_ex_machina():
             if p:
-                self.print(color=Color.RED, *args, **kwargs)
+                self.print(*args, color=Color.RED, **kwargs)
             else:
                 super().error(*args, **kwargs)
 
     def critical(self, p: int = 0, *args, **kwargs):
         if is_deus_ex_machina():
             if p:
-                self.print(color=Color.BOLD_RED, *args, **kwargs)
+                self.print(*args, color=Color.BOLD_RED, **kwargs)
             else:
                 super().critical(*args, **kwargs)
 
-    def print(self, color: Color, *args, **kwargs):
+    def print(self, *args, color: Color, **kwargs):
         more = None
         if kwargs:
             more = simplejson.dumps(

mage_ai/shared/enum.py

@@ -0,0 +1,8 @@
+try:
+    # breaking change introduced in python 3.11
+    from enum import Enum, IntEnum, StrEnum  # noqa: F401
+except ImportError:  # pragma: no cover
+    from enum import Enum, IntEnum  # noqa: F401 # pragma: no cover
+
+    class StrEnum(str, Enum):  # pragma: no cover
+        pass  # pragma: no cover

mage_ai/shared/environments.py

@@ -1,7 +1,13 @@
 import os
 import sys
 
-from mage_ai.shared.constants import ENV_DEV, ENV_PROD, ENV_STAGING, ENV_TEST
+from mage_ai.shared.constants import (
+    ENV_DEV,
+    ENV_PROD,
+    ENV_STAGING,
+    ENV_TEST,
+    ENV_TEST_MAGE,
+)
 
 
 def is_deus_ex_machina():
@@ -16,8 +22,24 @@ def is_dev():
     return os.getenv('ENV', None) == 'dev' or os.getenv('ENV', None) == 'development'
 
 
+def is_test_mage():
+    """
+    The "test_mage" environment is used for running unit tests in the CI/CD pipeline.
+    When running unit tests, a "test.db" sqlite database is temporarily created for
+    storing test metadata. Refer to the mage_ai/orchestration/db/__init__.py file for
+    implementation details. The test_mage environment is separate from the "test"
+    environment so that Mage's unit tests don't affect a user's own testing metadata db.
+    """
+    return os.getenv('ENV', None) == 'test_mage' or any('unittest' in v for v in sys.argv)
+
+
 def is_test():
-    return os.getenv('ENV', None) == 'test' or any('unittest' in v for v in sys.argv)
+    """
+    The "test" environment is the default environment for testing where users may have
+    their own specific tests that they want to run with their own test metadata db at a
+    specific location.
+    """
+    return os.getenv('ENV', None) == 'test' or is_test_mage()
 
 
 def is_production():
@@ -29,7 +51,9 @@ def is_staging():
 
 
 def get_env():
-    if is_test():
+    if is_test_mage():
+        return ENV_TEST_MAGE
+    elif is_test():
         return ENV_TEST
     elif is_production():
         return ENV_PROD

mage_ai/shared/logger.py

@@ -2,10 +2,10 @@ import json
 import logging
 import time
 from contextlib import contextmanager, redirect_stdout
-from enum import Enum
 from typing import Callable, List
 
 from mage_ai.settings import SERVER_LOGGING_TEMPLATE
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import merge_dict
 
 logger = logging.getLogger(__name__)
@@ -107,7 +107,7 @@ def set_logging_format(logging_format: str = None, level: str = None) -> None:
             root_logger.exception('Invalid logging level %s', level)
 
 
-class LoggingLevel(str, Enum):
+class LoggingLevel(StrEnum):
     DEBUG = 'DEBUG'
     INFO = 'INFO'
     WARNING = 'WARNING'

mage_ai/shared/models.py

@@ -7,12 +7,13 @@ from typing import Any, Dict, List, Optional, Type, Union
 
 import inflection
 
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.environments import is_debug
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.parsers import encode_complex
 
 
-class BaseEnum(str, Enum):
+class BaseEnum(StrEnum):
     @classmethod
     def has_value(cls, value: Union[Any, str]) -> bool:
         if isinstance(value, cls):

mage_ai/streaming/constants.py

@@ -1,10 +1,10 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 DEFAULT_BATCH_SIZE = 100
 DEFAULT_TIMEOUT_MS = 500
 
 
-class SourceType(str, Enum):
+class SourceType(StrEnum):
     ACTIVEMQ = 'activemq'
     AMAZON_SQS = 'amazon_sqs'
     AZURE_EVENT_HUB = 'azure_event_hub'
@@ -17,7 +17,7 @@ class SourceType(str, Enum):
     MONGODB = 'mongodb'
 
 
-class SinkType(str, Enum):
+class SinkType(StrEnum):
     ACTIVEMQ = 'activemq'
     AMAZON_S3 = 'amazon_s3'
     AZURE_DATA_LAKE = 'azure_data_lake'

mage_ai/streaming/sinks/elasticsearch.py

@@ -16,6 +16,8 @@ class ElasticSearchConfig(BaseConfig):
     api_key: str = None
     ca_cert: str = None
     _id: str = None
+    _op_type: str = None
+    _source: str = None
 
 
 class ElasticSearchSink(BaseSink):
@@ -51,9 +53,17 @@ class ElasticSearchSink(BaseSink):
 
     def batch_write(self, messages: List[Dict]):
         self._print(f'Batch ingest data {messages}, time={time.time()}')
-        if self.config._id is not None:
-            docs = [{'_index': self.config.index_name, '_id': doc[self.config._id],
-                     'doc': doc} for doc in messages]
-        else:
-            docs = [{'_index': self.config.index_name, 'doc': doc} for doc in messages]
+        docs = []
+        for msg in messages:
+            doc = {'_index': self.config.index_name}
+            if self.config._id is not None:
+                doc['id'] = msg[self.config._id]
+            if self.config._op_type is not None:
+                doc['_op_type'] = self.config._op_type
+            if self.config._source is not None:
+                doc[self.config._source] = msg
+            else:
+                doc['_doc'] = msg
+            docs.append(doc)
+
         helpers.bulk(self.client, docs)

mage_ai/streaming/sinks/kafka.py

@@ -1,17 +1,17 @@
 import json
 import time
 from dataclasses import dataclass
-from enum import Enum
 from typing import Dict, List
 
 from kafka import KafkaProducer
 
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_MS
 from mage_ai.streaming.sinks.base import BaseSink
 
 
-class SecurityProtocol(str, Enum):
+class SecurityProtocol(StrEnum):
     SASL_SSL = 'SASL_SSL'
     SSL = 'SSL'
 

mage_ai/streaming/sinks/postgres.py

@@ -23,6 +23,9 @@ class PostgresConfig(BaseConfig):
     port: int = 5432
     unique_conflict_method: str = UNIQUE_CONFLICT_METHOD_IGNORE
     unique_constraints: List = field(default_factory=list)
+    allow_reserved_words: bool = False
+    auto_clean_name: bool = True
+    case_sensitive: bool = False
 
 
 class PostgresSink(BaseSink):
@@ -57,6 +60,9 @@ class PostgresSink(BaseSink):
             if_exists=ExportWritePolicy.APPEND,
             unique_conflict_method=self.config.unique_conflict_method,
             unique_constraints=self.config.unique_constraints,
+            allow_reserved_words=self.config.allow_reserved_words,
+            auto_clean_name=self.config.auto_clean_name,
+            case_sensitive=self.config.case_sensitive,
         )
 
     def destroy(self):

mage_ai/streaming/sources/amazon_sqs.py

@@ -2,12 +2,12 @@ import json
 import threading
 import time
 from dataclasses import dataclass
-from enum import Enum
 from typing import Callable, Dict
 
 import boto3
 
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE
 from mage_ai.streaming.sources.base import BaseSource
 from mage_ai.streaming.sources.shared import SerDeConfig, SerializationMethod
@@ -15,7 +15,7 @@ from mage_ai.streaming.sources.shared import SerDeConfig, SerializationMethod
 DEFAULT_WAIT_TIME_SECONDS = 1
 
 
-class MessageDeletionMethod(str, Enum):
+class MessageDeletionMethod(StrEnum):
     AFTER_RECEIVED = 'AFTER_RECEIVED'
     MANUAL = 'MANUAL'
 

mage_ai/streaming/sources/base.py

@@ -1,12 +1,12 @@
 import json
 from abc import ABC, abstractmethod
-from enum import Enum
 from typing import Callable, Dict
 
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.environments import is_test
 
 
-class SourceConsumeMethod(str, Enum):
+class SourceConsumeMethod(StrEnum):
     BATCH_READ = 'BATCH_READ'
     READ = 'READ'
     READ_ASYNC = 'READ_ASYNC'