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

mage_ai/authentication/providers/okta.py

@@ -9,14 +9,17 @@ from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.authentication.providers.oauth import OauthProvider
 from mage_ai.authentication.providers.sso import SsoProvider
 from mage_ai.authentication.providers.utils import get_base_url
-from mage_ai.settings.sso import OKTA_CLIENT_ID, OKTA_CLIENT_SECRET, OKTA_DOMAIN_URL
+from mage_ai.settings import get_settings_value
+from mage_ai.settings.keys import OKTA_CLIENT_ID, OKTA_CLIENT_SECRET, OKTA_DOMAIN_URL
 
 
 class OktaProvider(SsoProvider, OauthProvider):
     provider = ProviderName.OKTA
 
     def __init__(self):
-        self.hostname = OKTA_DOMAIN_URL
+        self.hostname = get_settings_value(OKTA_DOMAIN_URL)
+        self.client_id = get_settings_value(OKTA_CLIENT_ID)
+        self.client_secret = get_settings_value(OKTA_CLIENT_SECRET)
         self.__validate()
 
         if not self.hostname.startswith('https'):
@@ -27,11 +30,11 @@ class OktaProvider(SsoProvider, OauthProvider):
             raise Exception(
                 'Okta hostname is empty. '
                 'Make sure the OKTA_DOMAIN_URL environment variable is set.')
-        if not OKTA_CLIENT_ID:
+        if not self.client_id:
             raise Exception(
                 'Okta client id is empty. '
                 'Make sure the OKTA_CLIENT_ID environment variable is set.')
-        if not OKTA_CLIENT_SECRET:
+        if not self.client_secret:
             raise Exception(
                 'Okta client secret is empty. '
                 'Make sure the OKTA_CLIENT_SECRET environment variable is set.')
@@ -43,7 +46,7 @@ class OktaProvider(SsoProvider, OauthProvider):
             redirect_uri=redirect_uri,
         )
         query = dict(
-            client_id=OKTA_CLIENT_ID,
+            client_id=self.client_id,
             redirect_uri=urllib.parse.quote_plus(
                 f'{base_url}/oauth',
             ),
@@ -75,7 +78,7 @@ class OktaProvider(SsoProvider, OauthProvider):
                     code=code,
                     redirect_uri=f'{base_url}/oauth',
                 ),
-                auth=BasicAuth(OKTA_CLIENT_ID, OKTA_CLIENT_SECRET),
+                auth=BasicAuth(self.client_id, self.client_secret),
                 timeout=20,
             ) as response:
                 data = await response.json()

mage_ai/cluster_manager/kubernetes/workload_manager.py

@@ -199,6 +199,7 @@ class WorkloadManager:
         self,
         name: str,
         workspace_config: KubernetesWorkspaceConfig,
+        initial_metadata: Dict = None,
         project_type: str = ProjectType.STANDALONE,
     ):
         """
@@ -245,6 +246,7 @@ class WorkloadManager:
             project_uuid=workspace_config.project_uuid,
             container_config=container_config,
             set_base_path=ingress_name is not None,
+            initial_metadata=initial_metadata,
         )
         container_config['env'] = env_vars
 
@@ -713,6 +715,7 @@ class WorkloadManager:
         project_type: str = 'standalone',
         project_uuid: str = None,
         container_config: Dict = None,
+        initial_metadata: Dict = None,
         set_base_path: bool = False,
     ) -> List:
         env_vars = [
@@ -746,6 +749,13 @@ class WorkloadManager:
                     'value': project_uuid,
                 }
             )
+        if initial_metadata:
+            env_vars.append(
+                {
+                    'name': 'INITIAL_METADATA',
+                    'value': json.dumps(initial_metadata),
+                }
+            )
 
         connection_url_secrets_name = os.getenv(CONNECTION_URL_SECRETS_NAME)
         if connection_url_secrets_name:

mage_ai/cluster_manager/workspace/base.py

@@ -6,7 +6,11 @@ from typing import Dict, Optional
 from mage_ai.cluster_manager.config import LifecycleConfig, WorkspaceConfig
 from mage_ai.cluster_manager.constants import ClusterType
 from mage_ai.cluster_manager.errors import WorkspaceExistsError
-from mage_ai.data_preparation.repo_manager import ProjectType, get_project_type
+from mage_ai.data_preparation.repo_manager import (
+    ProjectType,
+    get_project_type,
+    get_repo_config,
+)
 from mage_ai.orchestration.constants import Entity
 from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.orchestration.db.models.oauth import Permission, User
@@ -115,6 +119,7 @@ class Workspace(abc.ABC):
 
             project_uuid = uuid.uuid4().hex
             data['project_uuid'] = project_uuid
+            data['workspace_initial_metadata'] = get_repo_config().workspace_initial_metadata
 
         workspace_class = cls.workspace_class_from_type(cluster_type)
 

mage_ai/cluster_manager/workspace/kubernetes.py

@@ -49,6 +49,7 @@ class KubernetesWorkspace(Workspace):
             Workspace: An initialized Workspace object.
         """
         namespace = kwargs.pop('namespace', os.getenv(KUBE_NAMESPACE))
+        workspace_initial_metadata = kwargs.pop('workspace_initial_metadata', {})
         workspace_config = KubernetesWorkspaceConfig.load(
             config=merge_dict(
                 kwargs,
@@ -70,7 +71,9 @@ class KubernetesWorkspace(Workspace):
         if get_project_type() == ProjectType.MAIN:
             extra_args = {
                 'project_type': ProjectType.SUB,
+                'initial_metadata': workspace_initial_metadata,
             }
+
         workload_manager.create_workload(
             name,
             workspace_config,

mage_ai/data_preparation/decorators.py

@@ -103,3 +103,18 @@ def data_integration_query(function):
 
 def preprocesser_functions(function):
     return function
+
+
+def streaming_source(cls):
+    return cls
+
+
+def collect_decorated_objs(decorated_objs):
+    """
+    Method to collect the decorated objects (function or class)
+    """
+    def custom_code(obj):
+        decorated_objs.append(obj)
+        return obj
+
+    return custom_code

mage_ai/data_preparation/executors/streaming_pipeline_executor.py

@@ -10,7 +10,7 @@ from jinja2 import Template
 
 from mage_ai.data_preparation.executors.pipeline_executor import PipelineExecutor
 from mage_ai.data_preparation.logging.logger import DictLogger
-from mage_ai.data_preparation.models.constants import BlockType
+from mage_ai.data_preparation.models.constants import BlockLanguage, BlockType
 from mage_ai.data_preparation.models.pipeline import Pipeline
 from mage_ai.data_preparation.shared.stream import StreamToLogger
 from mage_ai.data_preparation.shared.utils import get_template_vars
@@ -107,18 +107,28 @@ class StreamingPipelineExecutor(PipelineExecutor):
 
         if global_vars is None:
             global_vars = dict()
-        source_config = self.__interpolate_vars(
-            self.source_block.content,
-            global_vars=global_vars,
-        )
-        source = SourceFactory.get_source(
-            source_config,
-            checkpoint_path=os.path.join(
-                self.pipeline.pipeline_variables_dir,
-                'streaming_checkpoint',
-            ),
-        )
 
+        # Initialize source block
+        if self.source_block.language == BlockLanguage.PYTHON:
+            source = SourceFactory.get_python_source(
+                self.source_block.content,
+                global_vars=global_vars,
+            )
+        else:
+            # Default to YAML config
+            source_config = self.__interpolate_vars(
+                self.source_block.content,
+                global_vars=global_vars,
+            )
+            source = SourceFactory.get_source(
+                source_config,
+                checkpoint_path=os.path.join(
+                    self.pipeline.pipeline_variables_dir,
+                    'streaming_checkpoint',
+                ),
+            )
+
+        # Initialize destination blocks
         sinks_by_uuid = dict()
         for sink_block in self.sink_blocks:
             sinks_by_uuid[sink_block.uuid] = SinkFactory.get_sink(

mage_ai/data_preparation/git/__init__.py

@@ -8,6 +8,7 @@ from pathlib import Path
 from typing import Any, Callable, Dict, List
 from urllib.parse import urlparse, urlsplit, urlunsplit
 
+from mage_ai.authentication.oauth.constants import ProviderName
 from mage_ai.data_preparation.git.constants import DEFAULT_KNOWN_HOSTS_FILE
 from mage_ai.data_preparation.git.utils import (
     add_host_to_known_hosts as add_host_to_known_hosts_util,
@@ -18,6 +19,7 @@ from mage_ai.data_preparation.git.utils import (
     check_connection_async,
     create_ssh_keys,
     get_access_token,
+    get_provider_from_remote_url,
     poll_process_with_timeout,
     run_command,
 )
@@ -622,7 +624,14 @@ class Git:
             try:
                 for url in remote.urls:
                     if url.lower().startswith('https'):
-                        repository_names.append('/'.join(url.split('/')[-2:]).replace('.git', ''))
+                        provider = get_provider_from_remote_url(url)
+                        if provider == ProviderName.AZURE_DEVOPS:
+                            updated_url = url.replace('/_git', '')
+                            repository_names.append('/'.join(updated_url.split('/')[-2:]))
+                        else:
+                            repository_names.append(
+                                '/'.join(url.split('/')[-2:]).replace('.git', '')
+                            )
 
                         # Remove the token from the URL
                         # e.g. https://[user]:[token]@[netloc]

mage_ai/data_preparation/git/api.py

@@ -15,8 +15,10 @@ from mage_ai.data_preparation.git.utils import (
 )
 from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.db.models.oauth import Oauth2AccessToken, User
+from mage_ai.server.logger import Logger
 
 API_ENDPOINT = 'https://api.github.com'
+logger = Logger().new_server_logger(__name__)
 
 
 def get_oauth_client_id(provider: str) -> str:
@@ -262,4 +264,5 @@ def get_user(token: str, provider: str = ProviderName.GITHUB) -> Dict:
     try:
         return GitClient.get_client_for_provider(provider)(token).get_user()
     except Exception:
+        logger.exception('Error fetching user from git provider.')
         return dict()

mage_ai/data_preparation/git/clients/azure_devops.py

@@ -0,0 +1,106 @@
+import json
+
+import requests
+
+from mage_ai.data_preparation.git.clients.base import Client
+from mage_ai.settings import get_settings_value
+from mage_ai.settings.keys import AZURE_DEVOPS_ORGANIZATION
+
+
+class AzureDevopsClient(Client):
+    def __init__(self, access_token: str):
+        super().__init__(access_token)
+        self.organization = get_settings_value(AZURE_DEVOPS_ORGANIZATION)
+
+    def get_user(self):
+        return dict(username='oauth2')
+
+    def get_branches(self, repository_name: str):
+        project_name, repository_name = repository_name.split('/')
+        resp = requests.get(
+            f'https://dev.azure.com/{self.organization}/{project_name}/_apis/git/repositories/{repository_name}/refs?api-version=7.1-preview.1',  # noqa: E501
+            headers={
+                'Accept': 'application/json',
+                'Authorization': f'Bearer {self.access_token}',
+            },
+            timeout=10,
+        )
+        data = resp.json()
+
+        arr = []
+        for branch in data.get('value', []):
+            branch_name = branch.get('name')
+            if branch_name.startswith('refs/heads/'):
+                arr.append(branch_name)
+        return arr
+
+    def get_pull_requests(self, repository_name: str):
+        project_name, repository_name = repository_name.split('/')
+        resp = requests.get(
+            f'https://dev.azure.com/{self.organization}/{project_name}/_apis/git/repositories/{repository_name}/pullrequests?api-version=7.1-preview.1',  # noqa: E501
+            headers={
+                'Accept': 'application/json',
+                'Authorization': f'Bearer {self.access_token}',
+            },
+            timeout=10,
+        )
+        data = resp.json()
+
+        arr = []
+        for pr in data.get('value', []):
+            arr.append(
+                dict(
+                    body=pr.get('description'),
+                    created_at=pr.get('creationDate'),
+                    id=pr.get('pullRequestId'),
+                    is_merged=pr.get('status') == 'completed',
+                    merged=pr.get('status') == 'completed',
+                    state=pr.get('status'),
+                    title=pr.get('title'),
+                    url=f'https://dev.azure.com/{self.organization}/{project_name}/_git/{repository_name}/pullrequest/{pr.get("pullRequestId")}',  # noqa: E501
+                    user=pr.get('createdBy', {}).get('displayName'),
+                )
+            )
+
+        return arr
+
+    def create_pull_request(
+        self,
+        repository_name: str,
+        base_branch: str,
+        body: str,
+        compare_branch: str,
+        title: str,
+    ):
+        project_name, repository_name = repository_name.split('/')
+        data = {
+            'title': title,
+            'description': body,
+            'sourceRefName': compare_branch,
+            'targetRefName': base_branch,
+        }
+
+        resp = requests.post(
+            f'https://dev.azure.com/{self.organization}/{project_name}/_apis/git/repositories/{repository_name}/pullrequests?api-version=7.1-preview.1',  # noqa: E501
+            headers={
+                'Accept': 'application/json',
+                'Authorization': f'Bearer {self.access_token}',
+                'Content-Type': 'application/json',
+            },
+            data=json.dumps(data),
+            timeout=10,
+        )
+
+        resp.raise_for_status()
+        pr = resp.json()
+        return dict(
+            body=pr.get('description'),
+            created_at=pr.get('creationDate'),
+            id=pr.get('pullRequestId'),
+            is_merged=pr.get('status') == 'completed',
+            merged=pr.get('status') == 'completed',
+            state=pr.get('status'),
+            title=pr.get('title'),
+            url=pr.get('url'),
+            user=pr.get('createdBy', {}).get('displayName'),
+        )

mage_ai/data_preparation/git/clients/base.py

@@ -25,6 +25,12 @@ class Client(ABC):
             from mage_ai.data_preparation.git.clients.ghe import GHEClient
 
             return GHEClient
+        elif provider == ProviderName.AZURE_DEVOPS:
+            from mage_ai.data_preparation.git.clients.azure_devops import (
+                AzureDevopsClient,
+            )
+
+            return AzureDevopsClient
         else:
             raise NotImplementedError()
 

mage_ai/data_preparation/git/clients/gitlab.py

@@ -3,14 +3,15 @@ from urllib.parse import quote_plus
 
 import requests
 
-from mage_ai.authentication.oauth.constants import GITLAB_HOST
 from mage_ai.data_preparation.git.clients.base import Client
+from mage_ai.settings import get_settings_value
+from mage_ai.settings.keys import GITLAB_HOST
 
 
 class GitlabClient(Client):
     def __init__(self, access_token: str):
         super().__init__(access_token)
-        self.hostname = GITLAB_HOST or 'https://gitlab.com'
+        self.hostname = get_settings_value(GITLAB_HOST, 'https://gitlab.com')
 
     def get_user(self):
         # For GitLab, we don't need to make a request to the API to get the username

mage_ai/data_preparation/git/utils.py

@@ -5,21 +5,21 @@ import subprocess
 from typing import Callable
 from urllib.parse import urlparse, urlsplit, urlunsplit
 
-from mage_ai.authentication.oauth.constants import (
-    BITBUCKET_HOST,
-    GITLAB_HOST,
-    ProviderName,
-    get_ghe_hostname,
-)
+from mage_ai.authentication.oauth.constants import ProviderName, get_ghe_hostname
 from mage_ai.data_preparation.git.constants import (
     DEFAULT_KNOWN_HOSTS_FILE,
     DEFAULT_SSH_KEY_DIRECTORY,
-    GIT_ACCESS_TOKEN_VAR,
-    GIT_SSH_PRIVATE_KEY_VAR,
-    GIT_SSH_PUBLIC_KEY_VAR,
 )
 from mage_ai.data_preparation.shared.secrets import get_secret_value
 from mage_ai.data_preparation.sync import AuthType, GitConfig
+from mage_ai.settings import get_settings_value
+from mage_ai.settings.keys import (
+    BITBUCKET_HOST,
+    GIT_ACCESS_TOKEN,
+    GIT_SSH_PRIVATE_KEY,
+    GIT_SSH_PUBLIC_KEY,
+    GITLAB_HOST,
+)
 from mage_ai.settings.repo import get_repo_path
 
 
@@ -38,10 +38,14 @@ def get_provider_from_remote_url(remote_url: str) -> str:
     if not remote_url:
         return ProviderName.GITHUB
 
-    if BITBUCKET_HOST and BITBUCKET_HOST in remote_url or 'bitbucket.org' in remote_url:
+    bitbucket_host = get_settings_value(BITBUCKET_HOST)
+    gitlab_host = get_settings_value(GITLAB_HOST)
+    if bitbucket_host and bitbucket_host in remote_url or 'bitbucket.org' in remote_url:
         return ProviderName.BITBUCKET
-    elif GITLAB_HOST and GITLAB_HOST in remote_url or 'gitlab.com' in remote_url:
+    elif gitlab_host and gitlab_host in remote_url or 'gitlab.com' in remote_url:
         return ProviderName.GITLAB
+    elif 'dev.azure.com' in remote_url:
+        return ProviderName.AZURE_DEVOPS
     elif ghe_hostname and ghe_hostname in remote_url:
         return ProviderName.GHE
     else:
@@ -59,13 +63,13 @@ def create_ssh_keys(git_config: GitConfig, repo_path: str, overwrite: bool = Fal
         )
         if not os.path.exists(public_key_file) or overwrite:
             try:
-                public_key = get_secret_value(
-                    pubk_secret_name,
-                    repo_name=repo_path,
-                    suppress_warning=True,
-                )
-                if os.getenv(GIT_SSH_PUBLIC_KEY_VAR):
-                    public_key = os.getenv(GIT_SSH_PUBLIC_KEY_VAR)
+                public_key = get_settings_value(GIT_SSH_PUBLIC_KEY)
+                if not public_key:
+                    public_key = get_secret_value(
+                        pubk_secret_name,
+                        repo_name=repo_path,
+                        suppress_warning=True,
+                    )
                 if public_key:
                     with open(public_key_file, 'w') as f:
                         f.write(base64.b64decode(public_key).decode('utf-8'))
@@ -81,13 +85,13 @@ def create_ssh_keys(git_config: GitConfig, repo_path: str, overwrite: bool = Fal
         )
         if not os.path.exists(custom_private_key_file) or overwrite:
             try:
-                private_key = get_secret_value(
-                    pk_secret_name,
-                    repo_name=repo_path,
-                    suppress_warning=True,
-                )
-                if os.getenv(GIT_SSH_PRIVATE_KEY_VAR):
-                    private_key = os.getenv(GIT_SSH_PRIVATE_KEY_VAR)
+                private_key = get_settings_value(GIT_SSH_PRIVATE_KEY)
+                if not private_key:
+                    private_key = get_secret_value(
+                        pk_secret_name,
+                        repo_name=repo_path,
+                        suppress_warning=True,
+                    )
                 if private_key:
                     with open(custom_private_key_file, 'w') as f:
                         f.write(base64.b64decode(private_key).decode('utf-8'))
@@ -137,10 +141,8 @@ def add_host_to_known_hosts(remote_repo_link: str):
 
 
 def get_access_token(git_config, repo_path: str = None) -> str:
-    token = None
-    if os.getenv(GIT_ACCESS_TOKEN_VAR):
-        token = os.getenv(GIT_ACCESS_TOKEN_VAR)
-    elif git_config and git_config.access_token_secret_name:
+    token = get_settings_value(GIT_ACCESS_TOKEN)
+    if not token and git_config and git_config.access_token_secret_name:
         token = get_secret_value(
             git_config.access_token_secret_name,
             repo_name=repo_path or get_repo_path(),

mage_ai/data_preparation/models/block/__init__.py

@@ -339,7 +339,6 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
         self.upstream_blocks = []
         self.downstream_blocks = []
         self.test_functions = []
-        self.global_vars = {}
         self.template_runtime_configuration = {}
 
         self.dynamic_block_index = None
@@ -1198,7 +1197,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                         f'Please run upstream blocks {upstream_block_uuids} '
                         'before running the current block.'
                     )
-            global_vars = self.__enrich_global_vars(
+            global_vars = self.enrich_global_vars(
                 global_vars,
                 dynamic_block_index=dynamic_block_index,
             )
@@ -2161,6 +2160,9 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                 )
                 pairs.append(tup)
 
+            if len(pairs) > 10:
+                # Limit the number of dynamic block children we display output for in the UI
+                pairs = pairs[:DATAFRAME_SAMPLE_COUNT_PREVIEW]
             for pair in pairs:
                 child_data = None
                 metadata = None
@@ -2201,6 +2203,7 @@ class Block(DataIntegrationMixin, SparkBlock, ProjectPlatformAccessible):
                 self.pipeline.uuid,
                 block_uuid,
                 partition=execution_partition,
+                max_results=DATAFRAME_SAMPLE_COUNT_PREVIEW,
             )
 
             if not include_print_outputs:
@@ -3034,7 +3037,7 @@ df = get_variable('{self.pipeline.uuid}', '{self.uuid}', 'df')
         variable_mapping = merge_dict(output_variables, consolidated_print_variables)
         return variable_mapping
 
-    def __enrich_global_vars(
+    def enrich_global_vars(
         self,
         global_vars: Dict = None,
         dynamic_block_index: int = None,
@@ -3655,23 +3658,29 @@ class AddonBlock(Block):
             global_vars['dynamic_block_index'] = dynamic_block_index
 
         if parent_block:
+            if parent_block.global_vars is None:
+                global_vars_copy = global_vars.copy()
+                parent_block.enrich_global_vars(
+                    global_vars=global_vars_copy,
+                    dynamic_block_index=dynamic_block_index,
+                )
+            global_vars = merge_dict(parent_block.global_vars, global_vars)
             global_vars['parent_block_uuid'] = parent_block.uuid
 
-        if parent_block and \
-                parent_block.pipeline and \
-                PipelineType.INTEGRATION == parent_block.pipeline.type:
-
-            template_runtime_configuration = parent_block.template_runtime_configuration
-            index = template_runtime_configuration.get('index', None)
-            is_last_block_run = template_runtime_configuration.get('is_last_block_run', False)
-            selected_streams = template_runtime_configuration.get('selected_streams', [])
-            stream = selected_streams[0] if len(selected_streams) >= 1 else None
-            destination_table = template_runtime_configuration.get('destination_table', stream)
-
-            global_vars['index'] = index
-            global_vars['is_last_block_run'] = is_last_block_run
-            global_vars['stream'] = stream
-            global_vars['destination_table'] = destination_table
+            if parent_block.pipeline and \
+                    PipelineType.INTEGRATION == parent_block.pipeline.type:
+
+                template_runtime_configuration = parent_block.template_runtime_configuration
+                index = template_runtime_configuration.get('index', None)
+                is_last_block_run = template_runtime_configuration.get('is_last_block_run', False)
+                selected_streams = template_runtime_configuration.get('selected_streams', [])
+                stream = selected_streams[0] if len(selected_streams) >= 1 else None
+                destination_table = template_runtime_configuration.get('destination_table', stream)
+
+                global_vars['index'] = index
+                global_vars['is_last_block_run'] = is_last_block_run
+                global_vars['stream'] = stream
+                global_vars['destination_table'] = destination_table
 
         return global_vars