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

mage_ai/shared/enum.py

@@ -1,11 +1,8 @@
 try:
     # breaking change introduced in python 3.11
-    from enum import Enum, IntEnum, StrEnum
+    from enum import Enum, IntEnum, StrEnum  # noqa: F401
 except ImportError:  # pragma: no cover
-    from enum import Enum  # pragma: no cover
-
-    class IntEnum(int, Enum):  # pragma: no cover
-        pass  # 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,10 +22,26 @@ def is_dev():
     return os.getenv('ENV', None) == 'dev' or os.getenv('ENV', None) == 'development'
 
 
-def is_test():
+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():
+    """
+    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():
     return os.getenv('ENV', None) == 'production' or os.getenv('ENV', None) == 'prod'
 
@@ -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/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

@@ -4,7 +4,6 @@ from dataclasses import dataclass
 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
@@ -12,6 +11,7 @@ from mage_ai.streaming.sinks.base import BaseSink
 
 
 class SecurityProtocol(StrEnum):
+    SASL_PLAINTEXT = 'SASL_PLAINTEXT'
     SASL_SSL = 'SASL_SSL'
     SSL = 'SSL'
 
@@ -21,6 +21,11 @@ class SASLConfig:
     mechanism: str = 'PLAIN'
     username: str = None
     password: str = None
+    # OAuth configuration for OAUTHBEARER mechanism
+    oauth_token_url: str = None
+    oauth_client_id: str = None
+    oauth_client_secret: str = None
+    oauth_scope: str = None
 
 
 @dataclass
@@ -84,10 +89,23 @@ class KafkaSink(BaseSink):
             kwargs['ssl_password'] = self.config.ssl_config.password
             kwargs['ssl_check_hostname'] = self.config.ssl_config.check_hostname
         elif self.config.security_protocol == SecurityProtocol.SASL_SSL:
+            if not self.config.sasl_config:
+                raise Exception('sasl_config is required when security_protocol is SASL_SSL')
             kwargs['security_protocol'] = SecurityProtocol.SASL_SSL
             kwargs['sasl_mechanism'] = self.config.sasl_config.mechanism
-            kwargs['sasl_plain_username'] = self.config.sasl_config.username
-            kwargs['sasl_plain_password'] = self.config.sasl_config.password
+
+            # Handle OAUTHBEARER mechanism
+            if self.config.sasl_config.mechanism == 'OAUTHBEARER':
+                from mage_ai.streaming.sources.kafka_oauth import (
+                    create_oauth_token_provider,
+                )
+
+                token_provider = create_oauth_token_provider(self.config.sasl_config)
+                kwargs['sasl_oauth_token_provider'] = token_provider
+            else:
+                # Handle PLAIN, SCRAM mechanisms
+                kwargs['sasl_plain_username'] = self.config.sasl_config.username
+                kwargs['sasl_plain_password'] = self.config.sasl_config.password
 
             if self.config.ssl_config is not None and self.config.ssl_config.cafile:
                 kwargs['ssl_cafile'] = self.config.ssl_config.cafile

mage_ai/streaming/sources/kafka.py

@@ -11,7 +11,6 @@ from dataclasses import dataclass, field
 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
@@ -30,6 +29,11 @@ class SASLConfig:
     mechanism: str = 'PLAIN'
     username: str = None
     password: str = None
+    # OAuth configuration for OAUTHBEARER mechanism
+    oauth_token_url: str = None
+    oauth_client_id: str = None
+    oauth_client_secret: str = None
+    oauth_scope: str = None
 
 
 @dataclass
@@ -91,9 +95,11 @@ class KafkaSource(BaseSource):
             api_version=self.config.api_version,
             auto_offset_reset=self.config.auto_offset_reset,
             max_partition_fetch_bytes=self.config.max_partition_fetch_bytes,
-            enable_auto_commit=True,
+            enable_auto_commit=False,
         )
         if self.config.security_protocol == SecurityProtocol.SSL:
+            if not self.config.ssl_config:
+                raise Exception('ssl_config is required when security_protocol is SSL')
             consumer_kwargs['security_protocol'] = SecurityProtocol.SSL
             consumer_kwargs['ssl_cafile'] = self.config.ssl_config.cafile
             consumer_kwargs['ssl_certfile'] = self.config.ssl_config.certfile
@@ -103,10 +109,23 @@ class KafkaSource(BaseSource):
                 'ssl_check_hostname'
             ] = self.config.ssl_config.check_hostname
         elif self.config.security_protocol == SecurityProtocol.SASL_SSL:
+            if not self.config.sasl_config:
+                raise Exception('sasl_config is required when security_protocol is SASL_SSL')
             consumer_kwargs['security_protocol'] = SecurityProtocol.SASL_SSL
             consumer_kwargs['sasl_mechanism'] = self.config.sasl_config.mechanism
-            consumer_kwargs['sasl_plain_username'] = self.config.sasl_config.username
-            consumer_kwargs['sasl_plain_password'] = self.config.sasl_config.password
+
+            # Handle OAUTHBEARER mechanism
+            if self.config.sasl_config.mechanism == 'OAUTHBEARER':
+                from mage_ai.streaming.sources.kafka_oauth import (
+                    create_oauth_token_provider,
+                )
+
+                token_provider = create_oauth_token_provider(self.config.sasl_config)
+                consumer_kwargs['sasl_oauth_token_provider'] = token_provider
+            else:
+                # Handle PLAIN, SCRAM mechanisms
+                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:
                 consumer_kwargs[
@@ -121,10 +140,27 @@ class KafkaSource(BaseSource):
                     consumer_kwargs['ssl_password'] = self.config.ssl_config.password
 
         elif self.config.security_protocol == SecurityProtocol.SASL_PLAINTEXT:
+            if not self.config.sasl_config:
+                raise Exception('sasl_config is required when security_protocol is SASL_PLAINTEXT')
             consumer_kwargs['security_protocol'] = SecurityProtocol.SASL_PLAINTEXT
             consumer_kwargs['sasl_mechanism'] = self.config.sasl_config.mechanism
-            consumer_kwargs['sasl_plain_username'] = self.config.sasl_config.username
-            consumer_kwargs['sasl_plain_password'] = self.config.sasl_config.password
+
+            # Handle OAUTHBEARER mechanism
+            if self.config.sasl_config.mechanism == 'OAUTHBEARER':
+                self._print(
+                    'WARNING: Using OAUTHBEARER with SASL_PLAINTEXT is not recommended. '
+                    'OAuth tokens should be used with SASL_SSL for secure transmission.'
+                )
+                from mage_ai.streaming.sources.kafka_oauth import (
+                    create_oauth_token_provider,
+                )
+
+                token_provider = create_oauth_token_provider(self.config.sasl_config)
+                consumer_kwargs['sasl_oauth_token_provider'] = token_provider
+            else:
+                # Handle PLAIN, SCRAM mechanisms
+                consumer_kwargs['sasl_plain_username'] = self.config.sasl_config.username
+                consumer_kwargs['sasl_plain_password'] = self.config.sasl_config.password
 
         if self.config.topic:
             topics = [self.config.topic]
@@ -171,12 +207,32 @@ class KafkaSource(BaseSource):
                     self.registry_client, self.config.topic
                 )
 
+    def _decode_key(self, key):
+        """Decode message key based on serialization method.
+
+        When using RAW_VALUE, keep key as raw bytes.
+        Otherwise, attempt UTF-8 decode with fallback to raw bytes on error.
+        """
+        if not key:
+            return None
+
+        # Keep as raw bytes for RAW_VALUE serialization
+        if (self.config.serde_config and
+                self.config.serde_config.serialization_method == SerializationMethod.RAW_VALUE):
+            return key
+
+        # Attempt UTF-8 decode, fallback to raw bytes on error
+        try:
+            return key.decode('utf-8')
+        except UnicodeDecodeError:
+            return key
+
     def _convert_message(self, message):
         if self.config.include_metadata:
             message = {
                 'data': self.__deserialize_message(message.value),
                 'metadata': {
-                    'key': message.key.decode() if message.key else None,
+                    'key': self._decode_key(message.key),
                     'partition': message.partition,
                     'offset': message.offset,
                     'time': int(message.timestamp),
@@ -270,6 +326,7 @@ class KafkaSource(BaseSource):
                     message_values.append(message)
             if len(message_values) > 0:
                 handler(message_values)
+            self.consumer.commit()
 
     def test_connection(self):
         self.consumer._client.check_version(timeout=5)

mage_ai/streaming/sources/kafka_oauth.py

@@ -0,0 +1,182 @@
+"""
+OAuth token provider for Kafka SASL OAUTHBEARER authentication.
+
+This module provides a token provider implementation for OAuth 2.0 client credentials flow,
+which can be used with kafka-python's OAUTHBEARER SASL mechanism.
+"""
+import logging
+import time
+from typing import Dict, Optional
+
+import requests
+
+try:
+    from kafka.sasl.oauth import AbstractTokenProvider
+except ImportError:
+    # Fallback for when kafka is not installed
+    class AbstractTokenProvider:
+        """Fallback token provider interface when kafka-python is not installed."""
+        
+        def __init__(self, **config):
+            pass
+        
+        def token(self):
+            raise NotImplementedError("token() must be implemented by subclass")
+        
+        def extensions(self):
+            return {}
+
+
+logger = logging.getLogger(__name__)
+
+
+def create_oauth_token_provider(sasl_config):
+    """
+    Create an OAuth token provider for OAUTHBEARER mechanism.
+    
+    Args:
+        sasl_config: SASL configuration containing OAuth parameters
+        
+    Returns:
+        ClientCredentialsTokenProvider instance
+        
+    Raises:
+        Exception: If required OAuth parameters are missing
+    """
+    if not sasl_config.oauth_token_url:
+        raise Exception(
+            'oauth_token_url is required in sasl_config for OAUTHBEARER mechanism'
+        )
+    if not sasl_config.oauth_client_id:
+        raise Exception(
+            'oauth_client_id is required in sasl_config for OAUTHBEARER mechanism'
+        )
+    if not sasl_config.oauth_client_secret:
+        raise Exception(
+            'oauth_client_secret is required in sasl_config for OAUTHBEARER mechanism'
+        )
+    
+    return ClientCredentialsTokenProvider(
+        token_url=sasl_config.oauth_token_url,
+        client_id=sasl_config.oauth_client_id,
+        client_secret=sasl_config.oauth_client_secret,
+        scope=sasl_config.oauth_scope,
+    )
+
+
+class ClientCredentialsTokenProvider(AbstractTokenProvider):
+    """
+    Token provider that implements OAuth 2.0 client credentials flow.
+    
+    This provider fetches access tokens from an OAuth 2.0 authorization server
+    (e.g., Keycloak) using client credentials and manages token refresh.
+    
+    Args:
+        token_url: The OAuth 2.0 token endpoint URL
+        client_id: The OAuth 2.0 client ID
+        client_secret: The OAuth 2.0 client secret
+        scope: Optional scope for the token request
+        token_expiry_buffer: Buffer time in seconds before token expiry to refresh (default: 60)
+    """
+    
+    def __init__(
+        self,
+        token_url: str,
+        client_id: str,
+        client_secret: str,
+        scope: Optional[str] = None,
+        token_expiry_buffer: int = 60,
+        **config
+    ):
+        super().__init__(**config)
+        self.token_url = token_url
+        self.client_id = client_id
+        self.client_secret = client_secret
+        self.scope = scope
+        self.token_expiry_buffer = token_expiry_buffer
+        
+        self._access_token: Optional[str] = None
+        self._token_expiry_time: Optional[float] = None
+    
+    def token(self) -> str:
+        """
+        Returns a valid access token, fetching a new one if necessary.
+        
+        Returns:
+            str: The access token to use for authentication
+        """
+        if self._is_token_expired():
+            self._fetch_token()
+        
+        return self._access_token
+    
+    def _is_token_expired(self) -> bool:
+        """
+        Check if the current token is expired or will expire soon.
+        
+        Returns:
+            bool: True if token is expired or will expire within buffer time
+        """
+        if self._access_token is None or self._token_expiry_time is None:
+            return True
+        
+        current_time = time.time()
+        return current_time >= (self._token_expiry_time - self.token_expiry_buffer)
+    
+    def _fetch_token(self):
+        """
+        Fetch a new access token from the OAuth 2.0 authorization server.
+        
+        Raises:
+            RuntimeError: If token fetch fails
+        """
+        try:
+            data = {
+                'grant_type': 'client_credentials',
+                'client_id': self.client_id,
+                'client_secret': self.client_secret,
+            }
+            
+            if self.scope:
+                data['scope'] = self.scope
+            
+            logger.debug(f"Fetching OAuth token from {self.token_url}")
+            
+            response = requests.post(
+                self.token_url,
+                data=data,
+                headers={'Content-Type': 'application/x-www-form-urlencoded'},
+                timeout=10,
+            )
+            
+            response.raise_for_status()
+            token_data = response.json()
+            
+            # Validate that access_token is present in the response
+            if 'access_token' not in token_data:
+                logger.error("OAuth server response missing 'access_token' field")
+                raise RuntimeError(
+                    "Invalid token response format: missing 'access_token' field"
+                )
+            
+            self._access_token = token_data['access_token']
+            expires_in = token_data.get('expires_in', 3600)
+            self._token_expiry_time = time.time() + expires_in
+            
+            logger.debug(f"Successfully fetched OAuth token, expires in {expires_in} seconds")
+            
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Failed to fetch OAuth token: {e}")
+            raise RuntimeError(f"Failed to fetch OAuth token: {e}")
+        except (KeyError, ValueError) as e:
+            logger.error(f"Invalid token response format: {e}")
+            raise RuntimeError(f"Invalid token response format: {e}")
+    
+    def extensions(self) -> Dict[str, str]:
+        """
+        Return optional extensions to be sent with the SASL/OAUTHBEARER request.
+        
+        Returns:
+            dict: Empty dict as no extensions are needed for basic client credentials flow
+        """
+        return {}

mage_ai/tests/api/endpoints/test_blocks.py

@@ -1,6 +1,4 @@
-import os
-
-from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
+from mage_ai.data_preparation.models.constants import BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.tests.api.endpoints.mixins import (
@@ -11,7 +9,6 @@ from mage_ai.tests.api.endpoints.mixins import (
     build_list_endpoint_tests,
     build_update_endpoint_tests,
 )
-from mage_ai.tests.shared.mixins import DBTMixin
 
 
 def get_pipeline(self):
@@ -109,100 +106,3 @@ build_delete_endpoint_tests(
     resource='block',
     resource_parent='pipelines',
 )
-
-
-class DBTBlockAPIEndpointTest(BaseAPIEndpointTest, DBTMixin):
-    pass
-
-
-def __assert_after_create(self):
-    pipeline = get_pipeline(self)
-    blocks = pipeline.blocks_by_uuid.values()
-    blocks_dbt = list(filter(
-        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.SQL,
-        blocks,
-    ))
-
-    self.assertEqual(len(blocks), 5)
-    self.assertEqual(len(blocks_dbt), 1)
-
-    block = blocks_dbt[0]
-    self.assertEqual({
-      'dbt_project_name': 'dbt/test_project',
-      'file_path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
-      'file_source': {
-        'path': 'dbt/test_project/models/example/my_first_dbt_model.sql',
-        'project_path': 'dbt/test_project',
-      },
-    }, block.configuration)
-
-    return True
-
-
-def __assert_after_create_yaml(self):
-    pipeline = get_pipeline(self)
-    blocks = pipeline.blocks_by_uuid.values()
-    blocks_dbt = list(filter(
-        lambda x: x.type == BlockType.DBT and x.language == BlockLanguage.YAML,
-        blocks,
-    ))
-
-    self.assertEqual(len(blocks), 5)
-    self.assertEqual(len(blocks_dbt), 1)
-
-    block = blocks_dbt[0]
-    self.assertEqual({
-      'dbt_project_name': 'dbt/test_project',
-      'file_source': {
-        'path': 'dbts/dbt_yaml.yaml',
-        'project_path': 'dbt/test_project',
-      },
-    }, block.configuration)
-
-    return True
-
-
-build_create_endpoint_tests(
-    DBTBlockAPIEndpointTest,
-    assert_after_create_count=__assert_after_create,
-    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
-    build_payload=lambda self: dict(
-        name=self.faker.unique.name(),
-        type=BlockType.DBT,
-        language=BlockLanguage.SQL,
-        configuration=dict(
-            file_path=os.path.join(
-                os.path.basename(self.dbt_directory),
-                'test_project',
-                'models',
-                'example',
-                'my_first_dbt_model.sql',
-            ),
-        ),
-    ),
-    get_resource_parent_id=lambda self: self.pipeline.uuid,
-    resource='block',
-    resource_parent='pipelines',
-)
-
-
-build_create_endpoint_tests(
-    DBTBlockAPIEndpointTest,
-    test_uuid='dbt_yaml',
-    assert_after_create_count=__assert_after_create_yaml,
-    assert_before_create_count=lambda self: len(get_pipeline(self).blocks_by_uuid.values()) == 4,
-    build_payload=lambda self: dict(
-        name='dbt_yaml',
-        type=BlockType.DBT,
-        language=BlockLanguage.YAML,
-        configuration=dict(
-            dbt_project_name=os.path.join(
-                os.path.basename(self.dbt_directory),
-                'test_project',
-            ),
-        ),
-    ),
-    get_resource_parent_id=lambda self: self.pipeline.uuid,
-    resource='block',
-    resource_parent='pipelines',
-)

mage_ai/tests/api/endpoints/test_configuration_options.py

@@ -1,6 +1,5 @@
 import json
 import os
-from unittest.mock import patch
 
 from mage_ai.authentication.permissions.constants import EntityName
 from mage_ai.settings.models.configuration_option import ConfigurationType, OptionType
@@ -8,7 +7,7 @@ from mage_ai.tests.api.endpoints.mixins import (
     BaseAPIEndpointTest,
     build_list_endpoint_tests,
 )
-from mage_ai.tests.shared.mixins import DBTMixin, ProjectPlatformMixin
+from mage_ai.tests.shared.mixins import DBTMixin
 
 CURRENT_FILE_PATH = os.path.dirname(os.path.realpath(__file__))
 
@@ -51,49 +50,3 @@ build_list_endpoint_tests(
     assert_after=__assert_after_list,
     build_query=__build_query,
 )
-
-
-@patch('mage_ai.settings.platform.utils.project_platform_activated', lambda: True)
-@patch('mage_ai.settings.repo.project_platform_activated', lambda: True)
-class ConfigurationOptionProjectPlatformAPIEndpointTest(
-    DBTMixin,
-    ProjectPlatformMixin,
-    BaseAPIEndpointTest,
-):
-    pass
-
-
-def __assert_after_list2(self, result, **kwargs):
-    with open(os.path.join(
-        CURRENT_FILE_PATH,
-        'mocks',
-        'mock_dbt_configuration_options_project_platform.json',
-    )) as f:
-        self.assertEqual(
-            result,
-            json.loads(f.read()),
-        )
-
-
-build_list_endpoint_tests(
-    ConfigurationOptionProjectPlatformAPIEndpointTest,
-    list_count=1,
-    get_resource_parent_id=lambda test_case: test_case.pipeline.uuid,
-    resource='configuration_option',
-    resource_parent='pipeline',
-    result_keys_to_compare=[
-        'configuration_type',
-        'metadata',
-        'option',
-        'option_type',
-        'resource_type',
-        'uuid',
-    ],
-    assert_after=__assert_after_list2,
-    build_query=__build_query,
-    patch_function_settings=[
-        ('mage_ai.settings.models.configuration_option.project_platform_activated', lambda: True),
-        ('mage_ai.settings.platform.utils.project_platform_activated', lambda: True),
-        ('mage_ai.settings.repo.project_platform_activated', lambda: True),
-    ],
-)