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

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/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,9 +1,9 @@
-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 = 'test_mage'
 
 VALID_ENVS = frozenset([
     ENV_DEV,
@@ -20,7 +20,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,11 @@
+try:
+    # breaking change introduced in python 3.11
+    from enum import Enum, IntEnum, StrEnum
+except ImportError:  # pragma: no cover
+    from enum import Enum  # pragma: no cover
+
+    class IntEnum(int, Enum):  # pragma: no cover
+        pass  # pragma: no cover
+
+    class StrEnum(str, Enum):  # pragma: no cover
+        pass  # pragma: no cover

mage_ai/shared/environments.py

@@ -17,7 +17,7 @@ def is_dev():
 
 
 def is_test():
-    return os.getenv('ENV', None) == 'test' or any('unittest' in v for v in sys.argv)
+    return os.getenv('ENV', None) == 'test_mage' or any('unittest' in v for v in sys.argv)
 
 
 def is_production():

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/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'

mage_ai/streaming/sources/kafka.py

@@ -1,19 +1,25 @@
 import importlib
 import json
+import sys
+
+if sys.version_info >= (3, 12, 0):
+    import six
+    sys.modules['kafka.vendor.six.moves'] = six.moves
+
 import time
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Callable, Dict, List
 
 from kafka import KafkaConsumer, TopicPartition
 
 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.sources.base import BaseSource
 from mage_ai.streaming.sources.shared import SerDeConfig, SerializationMethod
 
 
-class SecurityProtocol(str, Enum):
+class SecurityProtocol(StrEnum):
     SASL_PLAINTEXT = 'SASL_PLAINTEXT'
     SASL_SSL = 'SASL_SSL'
     SSL = 'SSL'
@@ -102,8 +108,18 @@ class KafkaSource(BaseSource):
             consumer_kwargs['sasl_plain_username'] = self.config.sasl_config.username
             consumer_kwargs['sasl_plain_password'] = self.config.sasl_config.password
 
-            if self.config.ssl_config is not None and self.config.ssl_config.cafile:
-                consumer_kwargs['ssl_cafile'] = self.config.ssl_config.cafile
+            if self.config.ssl_config:
+                consumer_kwargs[
+                    'ssl_check_hostname'] = self.config.ssl_config.check_hostname
+                if self.config.ssl_config.cafile:
+                    consumer_kwargs['ssl_cafile'] = self.config.ssl_config.cafile
+                if self.config.ssl_config.certfile:
+                    consumer_kwargs['ssl_certfile'] = self.config.ssl_config.certfile
+                if self.config.ssl_config.keyfile:
+                    consumer_kwargs['ssl_keyfile'] = self.config.ssl_config.keyfile
+                if self.config.ssl_config.password:
+                    consumer_kwargs['ssl_password'] = self.config.ssl_config.password
+
         elif self.config.security_protocol == SecurityProtocol.SASL_PLAINTEXT:
             consumer_kwargs['security_protocol'] = SecurityProtocol.SASL_PLAINTEXT
             consumer_kwargs['sasl_mechanism'] = self.config.sasl_config.mechanism

mage_ai/streaming/sources/nats_js.py

@@ -1,18 +1,20 @@
 import asyncio
-import enum
 import json
 import ssl
 import threading
 from dataclasses import dataclass
-from typing import Callable, Dict, Optional
+from typing import Callable, Dict, Optional, Tuple, Union
 
 import nats
 from nats.errors import NoServersError, TimeoutError
 
 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.sources.base import BaseSource, SourceConsumeMethod
 
+Credentials = Union[str, Tuple[str, str]]
+
 
 @dataclass
 class SSLConfig:
@@ -22,7 +24,7 @@ class SSLConfig:
     check_hostname: bool = False
 
 
-class ConsumerType(str, enum.Enum):
+class ConsumerType(StrEnum):
     PULL = "PULL"
     PUSH = "PUSH"
 
@@ -33,6 +35,7 @@ class NATSConfig(BaseConfig):
     stream_name: str
     subject: str = None
     nkeys_seed_str: Optional[str] = None
+    user_credentials: Optional[Credentials] = None
     use_tls: bool = False
     ssl_config: Optional[SSLConfig] = None
     consumer_name: Optional[str] = None
@@ -99,6 +102,10 @@ class NATSSource(BaseSource):
             if self.config.nkeys_seed_str:
                 connect_opts["nkeys_seed_str"] = self.config.nkeys_seed_str
 
+            # Use credentials if provided
+            if self.config.user_credentials:
+                connect_opts["user_credentials"] = self.config.user_credentials
+
             # Establish connection with the configured options
             self.nc = await nats.connect(**connect_opts)
             self.js = self.nc.jetstream()

mage_ai/streaming/sources/shared.py

@@ -1,8 +1,9 @@
 from dataclasses import dataclass
-from enum import Enum
 
+from mage_ai.shared.enum import StrEnum
 
-class SerializationMethod(str, Enum):
+
+class SerializationMethod(StrEnum):
     AVRO = 'AVRO'
     JSON = 'JSON'
     PROTOBUF = 'PROTOBUF'

mage_ai/system/constants.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 
-class LogType(str, Enum):
+class LogType(StrEnum):
     END = 'END'
     GENERIC = 'GENERIC'
     MEMORY = 'MEMORY'

mage_ai/tests/api/policies/permissions/test_base_policy_with_permissions.py

@@ -1,5 +1,4 @@
 import secrets
-from enum import Enum
 from typing import Any, Callable, Dict, List, Union
 from unittest.mock import patch
 
@@ -19,13 +18,14 @@ from mage_ai.authentication.permissions.constants import (
 )
 from mage_ai.orchestration.db.models.oauth import Permission, Role, User, UserRole
 from mage_ai.shared.array import find
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import index_by, merge_dict
 from mage_ai.tests.api.mixins import BootstrapMixin
 from mage_ai.tests.api.operations.test_base import BaseApiTestCase
 from mage_ai.tests.api.policies.permissions.mixins import PermissionsMixin
 
 
-class TestSuite(str, Enum):
+class TestSuite(StrEnum):
     AUTHORIZED = 'AUTHORIZED'
     DISABLED = 'DISABLED'
     INVERSE = 'INVERSE'