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

mage_ai/settings/backends.py

@@ -0,0 +1,95 @@
+import base64
+import logging
+import os
+from enum import Enum
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class BackendType(str, Enum):
+    """
+    Enum for the different types of settings backends.
+    """
+    AWS_SECRETS_MANAGER = 'aws_secrets_manager'
+
+
+class SettingsBackend:
+    """
+    Base settings backend class. A settings backend is a read-only storage of Mage settings.
+    The source of the configuration settings is dependent on the specific loader. The default source
+    are environment variables.
+    """
+    backend_type = None
+
+    def get(self, key: str, **kwargs) -> Optional[str]:
+        """
+        Loads the configuration setting stored under `key`.
+
+        Args:
+            key (str): Name of the setting to load
+
+        Returns:
+            Any: The setting value stored under `key` in the configuration manager. If key
+                 doesn't exist, returns None.
+        """
+        value = self._get(key, **kwargs)  # noqa: E1128
+        if value is None:
+            value = os.getenv(key)
+        if value is None:
+            value = kwargs.get('default')
+        return value
+
+    def _get(self, key: str, **kwargs) -> Optional[str]:
+        """
+        Loads the configuration setting stored under `key`.
+
+        Args:
+            key (str): Name of the setting to load
+
+        Returns:
+            Any: The setting value stored under `key` in the configuration manager. If key
+                 doesn't exist, returns None.
+        """
+        return None
+
+
+class AWSSecretsManagerBackend(SettingsBackend):
+    """
+    Settings backend that loads configuration settings from AWS Secrets Manager.
+    """
+    backend_type = BackendType.AWS_SECRETS_MANAGER
+
+    def __init__(self, **kwargs) -> None:
+        import boto3
+
+        self.client = boto3.client('secretsmanager')
+        self.prefix = kwargs.get('prefix', '')
+        self.use_cache = kwargs.get('use_cache', False)
+        self.cache = None
+        if self.use_cache:
+            from aws_secretsmanager_caching import SecretCache, SecretCacheConfig
+            cache_config_arg = kwargs.get('cache_config', {})
+            cache_config = SecretCacheConfig(**cache_config_arg)
+            self.cache = SecretCache(config=cache_config, client=self.client)
+
+    def _get(self, key: str, **kwargs) -> Optional[str]:
+        from botocore.exceptions import ClientError
+        if self.prefix:
+            key = f'{self.prefix}{key}'
+        if self.cache is not None:
+            return self.cache.get_secret_string(key)
+        else:
+            try:
+                secret_response = self.client.get_secret_value(
+                    SecretId=key,
+                )
+                if 'SecretBinary' in secret_response:
+                    binary = secret_response['SecretBinary']
+                    return base64.b64decode(binary)
+                else:
+                    return secret_response['SecretString']
+            except ClientError as error:
+                if error.response['Error']['Code'] != 'ResourceNotFoundException':
+                    logger.exception('Failed to get secret %s from AWS Secrets Manager.', key)
+                return None

mage_ai/settings/keys/__init__.py

@@ -0,0 +1 @@
+from mage_ai.settings.keys.auth import *  # noqa: F401, F403

mage_ai/settings/keys/auth.py

@@ -0,0 +1,76 @@
+# LDAP settings
+
+LDAP_SERVER = 'LDAP_SERVER'
+LDAP_BIND_DN = 'LDAP_BIND_DN'
+LDAP_BIND_PASSWORD = 'LDAP_BIND_PASSWORD'
+LDAP_BASE_DN = 'LDAP_BASE_DN'
+LDAP_AUTHENTICATION_FILTER = 'LDAP_AUTHENTICATION_FILTER'
+LDAP_AUTHORIZATION_FILTER = 'LDAP_AUTHORIZATION_FILTER'
+LDAP_ADMIN_USERNAME = 'LDAP_ADMIN_USERNAME'
+# Default access level to give to users created when authenticated through LDAP
+# for the first time. value should be the name of a Mage role (e.g. Viewer, Editor, Admin)
+LDAP_DEFAULT_ACCESS = 'LDAP_DEFAULT_ACCESS'
+LDAP_GROUP_FIELD = 'LDAP_GROUP_FIELD'
+LDAP_ROLES_MAPPING = 'LDAP_ROLES_MAPPING'
+
+# Git settings
+
+GIT_REPO_LINK = 'GIT_REPO_LINK'
+GIT_REPO_PATH = 'GIT_REPO_PATH'
+GIT_USERNAME = 'GIT_USERNAME'
+GIT_EMAIL = 'GIT_EMAIL'
+GIT_AUTH_TYPE = 'GIT_AUTH_TYPE'
+GIT_BRANCH = 'GIT_BRANCH'
+GIT_SYNC_ON_PIPELINE_RUN = 'GIT_SYNC_ON_PIPELINE_RUN'
+GIT_SYNC_ON_START = 'GIT_SYNC_ON_START'
+GIT_SYNC_ON_EXECUTOR_START = 'GIT_SYNC_ON_EXECUTOR_START'
+GIT_SYNC_SUBMODULES = 'GIT_SYNC_SUBMODULES'
+
+GIT_ENABLE_GIT_INTEGRATION = 'GIT_ENABLE_GIT_INTEGRATION'
+GIT_OVERWRITE_WITH_PROJECT_SETTINGS = 'GIT_OVERWRITE_WITH_PROJECT_SETTINGS'
+
+GIT_SSH_PUBLIC_KEY = 'GIT_SSH_PUBLIC_KEY'
+GIT_SSH_PRIVATE_KEY = 'GIT_SSH_PRIVATE_KEY'
+GIT_ACCESS_TOKEN = 'GIT_ACCESS_TOKEN'
+
+# Git Provider OAuth settings
+
+GHE_CLIENT_ID = 'GHE_CLIENT_ID'
+GHE_CLIENT_SECRET = 'GHE_CLIENT_SECRET'
+GHE_HOSTNAME = 'GHE_HOSTNAME'
+
+GITLAB_HOST = 'GITLAB_HOST'
+GITLAB_CLIENT_ID = 'GITLAB_CLIENT_ID'
+GITLAB_CLIENT_SECRET = 'GITLAB_CLIENT_SECRET'
+
+BITBUCKET_HOST = 'BITBUCKET_HOST'
+BITBUCKET_OAUTH_KEY = 'BITBUCKET_OAUTH_KEY'
+BITBUCKET_OAUTH_SECRET = 'BITBUCKET_OAUTH_SECRET'
+
+# Okta settings
+
+OKTA_DOMAIN_URL = 'OKTA_DOMAIN_URL'
+OKTA_CLIENT_ID = 'OKTA_CLIENT_ID'
+OKTA_CLIENT_SECRET = 'OKTA_CLIENT_SECRET'
+
+# Google settings
+
+GOOGLE_CLIENT_ID = 'GOOGLE_CLIENT_ID'
+GOOGLE_CLIENT_SECRET = 'GOOGLE_CLIENT_SECRET'
+
+# Active directory settings
+
+ACTIVE_DIRECTORY_CLIENT_ID = 'ACTIVE_DIRECTORY_CLIENT_ID'
+ACTIVE_DIRECTORY_CLIENT_SECRET = 'ACTIVE_DIRECTORY_CLIENT_SECRET'
+ACTIVE_DIRECTORY_DIRECTORY_ID = 'ACTIVE_DIRECTORY_DIRECTORY_ID'
+ACTIVE_DIRECTORY_ROLES_MAPPING = 'ACTIVE_DIRECTORY_ROLES_MAPPING'
+
+# Azure DevOps settings
+
+AZURE_DEVOPS_ORGANIZATION = 'AZURE_DEVOPS_ORGANIZATION'
+
+# OIDC
+
+OIDC_CLIENT_ID = 'OIDC_CLIENT_ID'
+OIDC_CLIENT_SECRET = 'OIDC_CLIENT_SECRET'
+OIDC_DISCOVERY_URL = 'OIDC_DISCOVERY_URL'

mage_ai/settings/server.py

@@ -0,0 +1,187 @@
+import os
+
+from .secret_generation import generate_jwt_secret
+
+# If you add a new environment variable, make sure to check if it should be added to
+# the `MAGE_SETTINGS_ENVIRONMENT_VARIABLES` list at the bottom of this file. Also, update
+# the environment variable documentation at docs/development/variables/environment-variables.mdx
+
+DEBUG = os.getenv('DEBUG', False)
+HIDE_ENV_VAR_VALUES = int(os.getenv('HIDE_ENV_VAR_VALUES', 1) or 1) == 1
+QUERY_API_KEY = 'api_key'
+
+"""
+import secrets
+secrets.token_urlsafe()
+
+Make sure this value is the same in mage_ai/frontend/api/constants.ts
+"""
+OAUTH2_APPLICATION_CLIENT_ID = 'zkWlN0PkIKSN0C11CfUHUj84OT5XOJ6tDZ6bDRO2'
+
+# valid values: 0, 1, 2
+try:
+    DISABLE_NOTEBOOK_EDIT_ACCESS = int(os.getenv('DISABLE_NOTEBOOK_EDIT_ACCESS', 0))
+except ValueError:
+    DISABLE_NOTEBOOK_EDIT_ACCESS = 1 if os.getenv('DISABLE_NOTEBOOK_EDIT_ACCESS') else 0
+
+
+def is_disable_pipeline_edit_access(
+    disable_notebook_edit_access_override: int = None,
+) -> bool:
+    value = DISABLE_NOTEBOOK_EDIT_ACCESS
+    if disable_notebook_edit_access_override is not None:
+        value = disable_notebook_edit_access_override
+    return value >= 1
+
+
+def get_bool_value(value: str) -> bool:
+    """
+    Converts a string environment variable to a bool value. Returns True if the value
+    is 'true', '1', or 't' (case insensitive). Otherwise, False
+    """
+    if value is None:
+        return False
+
+    return value.lower() in ('true', '1', 't')
+
+
+# ------------------------- DISABLE TERMINAL ----------------------
+
+DISABLE_TERMINAL = os.getenv('DISABLE_TERMINAL', '0').lower() in ('true', '1', 't')
+
+# ----------------- Authentication settings ----------------
+REQUIRE_USER_AUTHENTICATION = get_bool_value(
+    os.getenv('REQUIRE_USER_AUTHENTICATION', 'False')
+)
+REQUIRE_USER_PERMISSIONS = REQUIRE_USER_AUTHENTICATION and get_bool_value(
+    os.getenv('REQUIRE_USER_PERMISSIONS', 'False')
+)
+AUTHENTICATION_MODE = os.getenv('AUTHENTICATION_MODE', 'LOCAL')
+try:
+    MAGE_ACCESS_TOKEN_EXPIRY_TIME = int(
+        os.getenv('MAGE_ACCESS_TOKEN_EXPIRY_TIME', '2592000')
+    )
+except ValueError:
+    MAGE_ACCESS_TOKEN_EXPIRY_TIME = 2592000
+
+# Default access level to give to users created when authenticated through OAuth
+# for the first time. value should be the name of a Mage role (e.g. Viewer, Editor, Admin)
+OAUTH_DEFAULT_ACCESS = os.getenv('OAUTH_DEFAULT_ACCESS')
+
+# ----------------------------------------------------------
+
+HOSTNAME = os.getenv('HOSTNAME')
+REDIS_URL = os.getenv('REDIS_URL')
+SERVER_VERBOSITY = os.getenv('SERVER_VERBOSITY', 'info') or 'info'
+SERVER_LOGGING_FORMAT = os.getenv('SERVER_LOGGING_FORMAT', 'plaintext')
+SERVER_LOGGING_TEMPLATE = os.getenv(
+    'SERVER_LOGGING_TEMPLATE',
+    '%(levelname)s:%(name)s:%(message)s',
+)
+
+INITIAL_METADATA = os.getenv('INITIAL_METADATA', '{{}}')
+
+DISABLE_AUTO_BROWSER_OPEN = get_bool_value(os.getenv('DISABLE_AUTO_BROWSER_OPEN', 'False'))
+
+SHELL_COMMAND = os.getenv('SHELL_COMMAND', None)
+USE_UNIQUE_TERMINAL = os.getenv('USE_UNIQUE_TERMINAL', None)
+
+# sentry environment variables
+SENTRY_DSN = os.getenv('SENTRY_DSN')
+SENTRY_TRACES_SAMPLE_RATE = os.getenv('SENTRY_TRACES_SAMPLE_RATE', 1.0)
+
+# New relic enable environment variable
+ENABLE_NEW_RELIC = os.getenv('ENABLE_NEW_RELIC', False)
+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'))
+
+# 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)
+OTEL_PYTHON_TORNADO_EXCLUDED_URLS = (
+    os.getenv('OTEL_PYTHON_TORNADO_EXCLUDED_URLS') or '/api/statuses'
+)
+
+DEFAULT_LOCALHOST_URL = 'http://localhost:6789'
+MAGE_PUBLIC_HOST = os.getenv('MAGE_PUBLIC_HOST') or DEFAULT_LOCALHOST_URL
+
+MAX_FILE_CACHE_SIZE = os.getenv('MAX_FILE_CACHE_SIZE') or (1024 * 1024)  # 1 MB
+
+# All base path variables should not include a leading forward slash
+# e.g. MAGE_BASE_PATH = 'test_prefix' -> localhost:6789/test_prefix/pipelines
+BASE_PATH = os.getenv('MAGE_BASE_PATH')
+# Requests base path is used to configure the base path for the frontend requests. Defaults
+# to the MAGE_BASE_PATH environment variable.
+REQUESTS_BASE_PATH = os.getenv('MAGE_REQUESTS_BASE_PATH', BASE_PATH)
+# Routes base path is used to configure the base path for the backend routes. Defaults
+# to the MAGE_BASE_PATH environment variable.
+ROUTES_BASE_PATH = os.getenv('MAGE_ROUTES_BASE_PATH', BASE_PATH)
+# Used for OAUTH
+JWT_SECRET = os.getenv('JWT_SECRET', 'materia')
+# Used for generating download tokens
+JWT_DOWNLOAD_SECRET = os.getenv('JWT_DOWNLOAD_SECRET', generate_jwt_secret())
+# Sets the trigger interval of the scheduler to a numeric value, in seconds
+# Determines how often the scheduler gets invoked
+try:
+    SCHEDULER_TRIGGER_INTERVAL = float(os.getenv('SCHEDULER_TRIGGER_INTERVAL', '10'))
+except ValueError:
+    SCHEDULER_TRIGGER_INTERVAL = 10
+
+# List of environment variables used to configure Mage. The value of these settings
+# will be copied between workspaces.
+MAGE_SETTINGS_ENVIRONMENT_VARIABLES = [
+    'DISABLE_NOTEBOOK_EDIT_ACCESS',
+    'DISABLE_AUTO_BROWSER_OPEN',
+    'REQUIRE_USER_AUTHENTICATION',
+    'AUTHENTICATION_MODE',
+    'OAUTH_DEFAULT_ACCESS',
+    'LDAP_SERVER',
+    'LDAP_BIND_DN',
+    'LDAP_BIND_PASSWORD',
+    'LDAP_BASE_DN',
+    'LDAP_AUTHENTICATION_FILTER',
+    'LDAP_AUTHORIZATION_FILTER',
+    'LDAP_ADMIN_USERNAME',
+    'LDAP_DEFAULT_ACCESS',
+    'LDAP_GROUP_FIELD',
+    'LDAP_ROLES_MAPPING',
+    'SERVER_VERBOSITY',
+    'SERVER_LOGGING_FORMAT',
+    'SHELL_COMMAND',
+    'USE_UNIQUE_TERMINAL',
+    'SENTRY_DSN',
+    'SENTRY_TRACES_SAMPLE_RATE',
+    'MAGE_PUBLIC_HOST',
+    'SCHEDULER_TRIGGER_INTERVAL',
+    'REQUIRE_USER_PERMISSIONS',
+    'ENABLE_PROMETHEUS',
+    'OTEL_EXPORTER_OTLP_ENDPOINT',
+    'OTEL_EXPORTER_OTLP_HTTP_ENDPOINT',
+    'MAX_FILE_CACHE_SIZE',
+    'REDIS_URL',
+    # Oauth variables
+    'ACTIVE_DIRECTORY_DIRECTORY_ID',
+    'ACTIVE_DIRECTORY_CLIENT_ID',
+    'ACTIVE_DIRECTORY_CLIENT_SECRET',
+    'ACTIVE_DIRECTORY_ROLES_MAPPING',
+    'OKTA_DOMAIN_URL',
+    'OKTA_CLIENT_ID',
+    'OKTA_CLIENT_SECRET',
+    'GOOGLE_CLIENT_ID',
+    'GOOGLE_CLIENT_SECRET',
+    'OIDC_CLIENT_ID',
+    'OIDC_CLIENT_SECRET',
+    'OIDC_DISCOVERY_URL',
+    'GHE_CLIENT_ID',
+    'GHE_CLIENT_SECRET',
+    'GHE_HOSTNAME',
+    'BITBUCKET_HOST',
+    'BITBUCKET_OAUTH_KEY',
+    'BITBUCKET_OAUTH_SECRET',
+    'GITLAB_HOST',
+    'GITLAB_CLIENT_ID',
+    'GITLAB_CLIENT_SECRET',
+    'SERVER_LOGGING_TEMPLATE',
+]

mage_ai/shared/io.py

@@ -39,7 +39,7 @@ def safe_write(filepath: str, content: str, write_func: Callable = None):
 
     success = False
     try:
-        with open(filepath, 'w') as fp:
+        with open(filepath, 'w', encoding='utf-8') as fp:
             if write_func is not None:
                 write_func(fp, content)
             else:
@@ -67,7 +67,7 @@ async def safe_write_async(filepath: str, content: str, write_func: Callable = N
 
     success = False
     try:
-        async with aiofiles.open(filepath, mode='w') as fp:
+        async with aiofiles.open(filepath, mode='w', encoding='utf-8') as fp:
             if write_func is not None:
                 await write_func(fp, content)
             else:

mage_ai/shared/logger.py

@@ -5,6 +5,7 @@ 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.hash import merge_dict
 
 logger = logging.getLogger(__name__)
@@ -84,16 +85,21 @@ class JSONFormatter(logging.Formatter):
 
 
 def set_logging_format(logging_format: str = None, level: str = None) -> None:
-    handler = logging.StreamHandler()
-    handler.setFormatter(JSONFormatter())
-
-    root_logger = logging.getLogger()
     if isinstance(logging_format, str):
         logging_format = logging_format.lower()
+
+    root_logger = logging.getLogger()
+    if len(root_logger.handlers) > 0:
+        root_logger.removeHandler(root_logger.handlers[0])
+
+    handler = logging.StreamHandler()
     if logging_format == 'json':
-        if len(root_logger.handlers) > 0:
-            root_logger.removeHandler(root_logger.handlers[0])
+        handler.setFormatter(JSONFormatter())
         root_logger.addHandler(handler)
+    else:
+        handler.setFormatter(logging.Formatter(SERVER_LOGGING_TEMPLATE))
+        root_logger.addHandler(handler)
+
     if level:
         try:
             root_logger.setLevel(level.upper())

mage_ai/streaming/sources/base_python.py

@@ -0,0 +1,30 @@
+from typing import Callable
+
+from mage_ai.streaming.sources.base import BaseSource
+
+
+class BasePythonSource(BaseSource):
+    def __init__(self, **kwargs):
+        """
+        Not require config in the python source
+        """
+        super().__init__(None, **kwargs)
+
+    def init_client(self):
+        """
+        Intialize the client for the source.
+        """
+
+    def read(self, handler: Callable):
+        """
+        Read the message from the source and use handler to process the message.
+
+        This method only needs to be implemented when consume_method is 'READ'.
+        """
+
+    def batch_read(self, handler: Callable):
+        """
+        Batch read the messages from the source and use handler to process the messages.
+
+        This method only needs to be implemented when consume_method is 'BATCH_READ'.
+        """

mage_ai/streaming/sources/source_factory.py

@@ -1,5 +1,6 @@
 from typing import Dict
 
+from mage_ai.data_preparation.decorators import collect_decorated_objs
 from mage_ai.streaming.constants import SourceType
 
 
@@ -51,3 +52,27 @@ class SourceFactory:
         raise Exception(
             f'Consuming data from {connector_type} is not supported in streaming pipelines yet.',
         )
+
+    @classmethod
+    def get_python_source(self, content: str, **kwargs):
+        """
+        Find the class that's decorated with streaming_source from the source code.
+
+        Args:
+            content (str): The python code that contains the streaming source implementation.
+            **kwargs: {'global_vars': {...}}
+
+        Returns:
+            The initialized class object.
+
+        Raises:
+            Exception: Description
+        """
+        decorated_sources = []
+
+        exec(content, {'streaming_source': collect_decorated_objs(decorated_sources)})
+
+        if not decorated_sources:
+            raise Exception('Not find the class that has streaming_source decorator.')
+
+        return decorated_sources[0](**kwargs)

mage_ai/tests/api/endpoints/test_oauths.py

@@ -1,3 +1,5 @@
+import os
+
 from mage_ai.tests.api.endpoints.mixins import (
     BaseAPIEndpointTest,
     build_list_endpoint_tests,
@@ -5,7 +7,17 @@ from mage_ai.tests.api.endpoints.mixins import (
 
 
 class OauthAPIEndpointTest(BaseAPIEndpointTest):
-    pass
+    @classmethod
+    def setUpClass(self):
+        super().setUpClass()
+        os.environ['GOOGLE_CLIENT_ID'] = 'test'
+        os.environ['GOOGLE_CLIENT_SECRET'] = 'test secret'
+
+    @classmethod
+    def tearDownClass(self):
+        super().tearDownClass()
+        del os.environ['GOOGLE_CLIENT_ID']
+        del os.environ['GOOGLE_CLIENT_SECRET']
 
 
 build_list_endpoint_tests(
@@ -18,8 +30,4 @@ build_list_endpoint_tests(
         'url',
         'redirect_query_params',
     ],
-    patch_function_settings=[
-        ('mage_ai.settings.sso.GOOGLE_CLIENT_ID', lambda: 'test'),
-        ('mage_ai.settings.sso.GOOGLE_CLIENT_SECRET', lambda: 'test secret'),
-    ]
 )

mage_ai/tests/api/operations/test_operations.py

@@ -7,7 +7,7 @@ from mage_ai.tests.api.operations.test_base import BaseApiTestCase
 
 
 class OperationTests(BaseApiTestCase):
-    @patch('mage_ai.settings.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
+    @patch('mage_ai.settings.server.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
     @patch('mage_ai.api.utils.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
     @patch('mage_ai.api.policies.BasePolicy.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
     @patch('mage_ai.api.policies.BasePolicy.REQUIRE_USER_AUTHENTICATION', 0)
@@ -25,7 +25,7 @@ class OperationTests(BaseApiTestCase):
 
         self.assertIsNotNone(response['error'])
 
-    @patch('mage_ai.settings.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
+    @patch('mage_ai.settings.server.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
     @patch('mage_ai.api.utils.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
     @patch('mage_ai.api.policies.BasePolicy.DISABLE_NOTEBOOK_EDIT_ACCESS', 1)
     @patch('mage_ai.api.policies.BasePolicy.REQUIRE_USER_AUTHENTICATION', 1)