mage-ai 0.9.69__py3-none-any.whl → 0.9.71__py3-none-any.whl

mage_ai/data_preparation/git/utils.py

@@ -6,21 +6,24 @@ from typing import Callable
 from urllib.parse import urlparse, urlsplit, urlunsplit
 
 from mage_ai.authentication.oauth.constants import ProviderName, get_ghe_hostname
+from mage_ai.authentication.oauth.utils import access_tokens_for_client
 from mage_ai.data_preparation.git.constants import (
     DEFAULT_KNOWN_HOSTS_FILE,
     DEFAULT_SSH_KEY_DIRECTORY,
 )
+from mage_ai.data_preparation.repo_manager import get_project_uuid
 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.orchestration.db.models.oauth import Oauth2AccessToken, User
+from mage_ai.settings import get_bool_value, get_settings_value
 from mage_ai.settings.keys import (
     BITBUCKET_HOST,
     GIT_ACCESS_TOKEN,
+    GIT_OVERWRITE_WITH_PROJECT_SETTINGS,
     GIT_SSH_PRIVATE_KEY,
     GIT_SSH_PUBLIC_KEY,
     GITLAB_HOST,
 )
-from mage_ai.settings.repo import get_repo_path
 
 
 def get_auth_type_from_url(remote_url: str) -> AuthType:
@@ -52,24 +55,31 @@ def get_provider_from_remote_url(remote_url: str) -> str:
         return ProviderName.GITHUB
 
 
-def create_ssh_keys(git_config: GitConfig, repo_path: str, overwrite: bool = False) -> str:
+def create_ssh_keys(
+    git_config: GitConfig, repo_path: str, overwrite: bool = False
+) -> str:
     if not os.path.exists(DEFAULT_SSH_KEY_DIRECTORY):
         os.mkdir(DEFAULT_SSH_KEY_DIRECTORY, 0o700)
     pubk_secret_name = git_config.ssh_public_key_secret_name
+    overwrite_with_project_settings = get_bool_value(
+        get_settings_value(GIT_OVERWRITE_WITH_PROJECT_SETTINGS)
+    )
     if pubk_secret_name:
         public_key_file = os.path.join(
-            DEFAULT_SSH_KEY_DIRECTORY,
-            f'id_rsa_{pubk_secret_name}.pub'
+            DEFAULT_SSH_KEY_DIRECTORY, f'id_rsa_{pubk_secret_name}.pub'
         )
         if not os.path.exists(public_key_file) or overwrite:
             try:
                 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,
-                    )
+                public_key_from_secrets = get_secret_value(
+                    pubk_secret_name,
+                    repo_name=repo_path,
+                    suppress_warning=True,
+                )
+                if public_key_from_secrets and (
+                    not public_key or overwrite_with_project_settings
+                ):
+                    public_key = public_key_from_secrets
                 if public_key:
                     with open(public_key_file, 'w') as f:
                         f.write(base64.b64decode(public_key).decode('utf-8'))
@@ -80,18 +90,20 @@ def create_ssh_keys(git_config: GitConfig, repo_path: str, overwrite: bool = Fal
     private_key_file = os.path.join(DEFAULT_SSH_KEY_DIRECTORY, 'id_rsa')
     if pk_secret_name:
         custom_private_key_file = os.path.join(
-            DEFAULT_SSH_KEY_DIRECTORY,
-            f'id_rsa_{pk_secret_name}'
+            DEFAULT_SSH_KEY_DIRECTORY, f'id_rsa_{pk_secret_name}'
         )
         if not os.path.exists(custom_private_key_file) or overwrite:
             try:
                 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,
-                    )
+                private_key_from_secrets = get_secret_value(
+                    pk_secret_name,
+                    repo_name=repo_path,
+                    suppress_warning=True,
+                )
+                if private_key_from_secrets and (
+                    not private_key or overwrite_with_project_settings
+                ):
+                    private_key = private_key_from_secrets
                 if private_key:
                     with open(custom_private_key_file, 'w') as f:
                         f.write(base64.b64decode(private_key).decode('utf-8'))
@@ -140,13 +152,19 @@ def add_host_to_known_hosts(remote_repo_link: str):
     return False
 
 
-def get_access_token(git_config, repo_path: str = None) -> str:
+def get_access_token(git_config, repo_path: str) -> str:
     token = get_settings_value(GIT_ACCESS_TOKEN)
-    if not token and git_config and git_config.access_token_secret_name:
-        token = get_secret_value(
+    overwrite_with_project_settings = get_bool_value(
+        get_settings_value(GIT_OVERWRITE_WITH_PROJECT_SETTINGS)
+    )
+    if git_config and git_config.access_token_secret_name:
+        token_from_secrets = get_secret_value(
             git_config.access_token_secret_name,
-            repo_name=repo_path or get_repo_path(),
+            repo_name=repo_path,
+            suppress_warning=True,
         )
+        if token_from_secrets and (not token or overwrite_with_project_settings):
+            token = token_from_secrets
 
     return token
 
@@ -179,10 +197,7 @@ async def poll_process_with_timeout(
     if return_code is not None:
         out, err = proc.communicate()
         if return_code != 0:
-            message = (
-                err.decode('UTF-8') if err
-                else error_message
-            )
+            message = err.decode('UTF-8') if err else error_message
             raise ChildProcessError(message)
         else:
             return out.decode('UTF-8') if out else None
@@ -205,17 +220,16 @@ async def check_connection_async(git, remote_name: str) -> None:
             'Error connecting to remote, make sure your access token or SSH key is '
             'set up properly.'
         ),
+        timeout=5,
     )
 
 
 async def validate_authentication_for_remote_url(git, remote_url: str) -> None:
     proc = git.ls_remote(remote_url, as_process=True)
 
-    asyncio.run(
-        poll_process_with_timeout(
-            proc,
-            error_message='Error connecting to remote, make sure your access is valid.',
-        )
+    await poll_process_with_timeout(
+        proc,
+        error_message='Error connecting to remote, make sure your access is valid.',
     )
 
 
@@ -258,14 +272,30 @@ def execute_on_remote_branch(func: Callable, branch) -> None:
                     if hostname:
                         await __add_host_to_known_hosts(remote_name)
                     else:
-                        raise TimeoutError("""
+                        raise TimeoutError(
+                            """
 Connecting to remote timed out, make sure your SSH key is set up properly
 and your repository host is added as a known host. More information
 here: https://docs.mage.ai/developing-in-the-cloud/setting-up-git#5-add-github-com-to-known-hosts
-""")
+"""
+                        )
             return await func(*args, **kwargs)
         else:
             await check_connection_async(git, remote_name)
             return await func(*args, **kwargs)
 
     return wrapper
+
+
+def get_oauth_client_id(provider: str) -> str:
+    return f'{provider}_{get_project_uuid()}'
+
+
+def get_oauth_access_token_for_user(
+    user: User, provider: str = None
+) -> Oauth2AccessToken:
+    if not provider:
+        provider = ProviderName.GHE if get_ghe_hostname() else ProviderName.GITHUB
+    access_tokens = access_tokens_for_client(get_oauth_client_id(provider), user=user)
+    if access_tokens:
+        return access_tokens[0]

mage_ai/data_preparation/logging/logger_manager.py

@@ -7,11 +7,12 @@ from datetime import datetime
 from typing import Callable, Dict, List
 
 from mage_ai.data_preparation.logging import LoggingConfig
-from mage_ai.data_preparation.models.constants import LOGS_DIR
+from mage_ai.data_preparation.models.constants import LOGS_SUBDIR
 from mage_ai.data_preparation.models.file import File
 from mage_ai.data_preparation.repo_manager import RepoConfig, get_repo_config
 from mage_ai.data_preparation.storage.local_storage import LocalStorage
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.settings.server import LOGS_DIR_PATH
 from mage_ai.shared.array import find
 from mage_ai.shared.dates import str_to_timedelta
 
@@ -175,13 +176,13 @@ class LoggerManager:
             str: The log file path prefix based on pipeline_uuid, logs_dir, partition, and
                  subpartition.
         """
-        logs_dir = self.logs_dir or self.repo_config.variables_dir
+        logs_dir = self.logs_dir or LOGS_DIR_PATH or self.repo_config.variables_dir
 
         return os.path.join(
             logs_dir,
             'pipelines',
             pipeline_uuid or self.pipeline_uuid or 'all_pipelines',
-            LOGS_DIR,
+            LOGS_SUBDIR,
             self.partition or '',
             self.subpartition or '',
         )