mage-ai 0.8.80__py3-none-any.whl → 0.8.82__py3-none-any.whl

mage_ai/api/presenters/SyncPresenter.py

@@ -14,4 +14,5 @@ class SyncPresenter(BasePresenter):
         'type',
         'access_token',
         'auth_type',
+        'user_git_settings',
     ]

mage_ai/api/resources/GitBranchResource.py

@@ -26,7 +26,7 @@ class GitBranchResource(GenericResource):
     async def member(self, pk, user, **kwargs):
         branch = None
         preferences = get_preferences(user=user)
-        if preferences.has_valid_git_config():
+        if preferences.is_git_integration_enabled():
             git_manager = Git.get_manager(user=user)
             branch = git_manager.current_branch
         return self(dict(name=branch), user, **kwargs)

mage_ai/api/resources/SyncResource.py

@@ -1,20 +1,21 @@
+import os
+from typing import Dict
+
 from mage_ai.api.resources.GenericResource import GenericResource
-from mage_ai.data_preparation.preferences import (
-    get_preferences,
-    Preferences,
-)
+from mage_ai.data_preparation.preferences import get_preferences
+from mage_ai.data_preparation.repo_manager import get_repo_path
 from mage_ai.data_preparation.shared.secrets import create_secret
 from mage_ai.data_preparation.sync import (
-    GitConfig,
     GIT_ACCESS_TOKEN_SECRET_NAME,
     GIT_SSH_PRIVATE_KEY_SECRET_NAME,
     GIT_SSH_PUBLIC_KEY_SECRET_NAME,
+    GitConfig,
+    UserGitConfig,
 )
 from mage_ai.data_preparation.sync.git_sync import GitSync
 from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import User
 from mage_ai.orchestration.db.models.secrets import Secret
-import os
 
 
 def get_ssh_public_key_secret_name(user: User = None) -> str:
@@ -35,8 +36,7 @@ def get_access_token_secret_name(user: User = None) -> str:
 class SyncResource(GenericResource):
     @classmethod
     def collection(self, query, meta, user, **kwargs):
-        preferences = get_preferences(user=user)
-        sync_config = preferences.sync_config
+        sync_config = self.get_project_sync_config(user)
         return self.build_result_set(
             [sync_config],
             user,
@@ -46,59 +46,101 @@ class SyncResource(GenericResource):
     @classmethod
     @safe_db_query
     def create(self, payload, user, **kwargs):
-        ssh_public_key = payload.pop('ssh_public_key', None)
-        ssh_private_key = payload.pop('ssh_private_key', None)
-
-        if ssh_public_key:
-            secret_name = get_ssh_public_key_secret_name(user=user)
-            secret = Secret.query.filter(
-                Secret.name == secret_name).one_or_none()
-            if secret:
-                secret.delete()
-            create_secret(secret_name, ssh_public_key)
-            payload['ssh_public_key_secret_name'] = secret_name
-        if ssh_private_key:
-            secret_name = get_ssh_private_key_secret_name(user=user)
-            secret = Secret.query.filter(
-                Secret.name == secret_name).one_or_none()
-            if secret:
-                secret.delete()
-            create_secret(secret_name, ssh_private_key)
-            payload['ssh_private_key_secret_name'] = secret_name
-
-        access_token = payload.pop('access_token', None)
-        if access_token:
-            secret_name = get_access_token_secret_name(user=user)
-            secret = Secret.query.filter(
-                Secret.name == secret_name).one_or_none()
-            if secret:
-                secret.delete()
-            create_secret(secret_name, access_token)
-            payload['access_token_secret_name'] = secret_name
+        user_settings = payload.pop('user_git_settings', dict())
 
-        preferences = Preferences(user=user) if user else get_preferences()
+        payload = self.update_user_settings(payload)
+        preferences = get_preferences()
         updated_config = dict(preferences.sync_config, **payload)
         # default repo_path to os.getcwd()
         if not updated_config.get('repo_path', None):
             updated_config['repo_path'] = os.getcwd()
-        # Validate payload
+
+        # Validate payloads
+        user_payload = self.update_user_settings(user_settings, user=user)
+        UserGitConfig.load(config=user_payload)
         sync_config = GitConfig.load(config=updated_config)
 
+        # Update user git settings if they are included
+        if user:
+            repo_path = get_repo_path()
+            user_preferences = user.preferences or {}
+            user_git_settings = user.git_settings or {}
+            user_preferences[repo_path] = {
+                **user_preferences.get(repo_path, {}),
+                'git_settings': {
+                    **user_git_settings,
+                    **user_payload,
+                }
+            }
+            user.refresh()
+            user.update(preferences=user_preferences)
+        else:
+            updated_config.update(user_payload)
+
         preferences.update_preferences(dict(sync_config=updated_config))
 
         GitSync(sync_config)
 
-        return self(get_preferences(user=user).sync_config, user, **kwargs)
+        return self(get_preferences().sync_config, user, **kwargs)
 
     @classmethod
     def member(self, pk, user, **kwargs):
-        return self(get_preferences(user=user).sync_config, user, **kwargs)
+        sync_config = self.get_project_sync_config(user)
+        return self(sync_config, user, **kwargs)
 
     def update(self, payload, **kwargs):
+        self.model.pop('user_git_settings')
         config = GitConfig.load(config=self.model)
         sync = GitSync(config)
         action_type = payload.get('action_type')
         if action_type == 'sync_data':
             sync.sync_data()
+        elif action_type == 'reset':
+            sync.reset()
 
         return self
+
+    @classmethod
+    def get_project_sync_config(self, user):
+        sync_config = get_preferences().sync_config
+        # Make it backwards compatible with storing all of the git settings in the user
+        # preferences field.
+        if user and user.git_settings:
+            sync_config['user_git_settings'] = user.git_settings
+        else:
+            sync_config['user_git_settings'] = UserGitConfig.from_dict(sync_config).to_dict()
+        return sync_config
+
+    @classmethod
+    def update_user_settings(self, payload, user=None) -> Dict:
+        user_payload = payload.copy()
+        ssh_public_key = user_payload.pop('ssh_public_key', None)
+        ssh_private_key = user_payload.pop('ssh_private_key', None)
+
+        if ssh_public_key:
+            secret_name = get_ssh_public_key_secret_name(user=user)
+            secret = Secret.query.filter(
+                Secret.name == secret_name).one_or_none()
+            if secret:
+                secret.delete()
+            create_secret(secret_name, ssh_public_key)
+            user_payload['ssh_public_key_secret_name'] = secret_name
+        if ssh_private_key:
+            secret_name = get_ssh_private_key_secret_name(user=user)
+            secret = Secret.query.filter(
+                Secret.name == secret_name).one_or_none()
+            if secret:
+                secret.delete()
+            create_secret(secret_name, ssh_private_key)
+            user_payload['ssh_private_key_secret_name'] = secret_name
+
+        access_token = user_payload.pop('access_token', None)
+        if access_token:
+            secret_name = get_access_token_secret_name(user=user)
+            secret = Secret.query.filter(
+                Secret.name == secret_name).one_or_none()
+            if secret:
+                secret.delete()
+            create_secret(secret_name, access_token)
+            user_payload['access_token_secret_name'] = secret_name
+        return user_payload

mage_ai/data_preparation/executors/block_executor.py

@@ -112,6 +112,8 @@ class BlockExecutor:
                     )
                 self._execute_callback(
                     'on_failure',
+                    dynamic_block_index=dynamic_block_index,
+                    dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                     global_vars=global_vars,
                     logging_tags=tags,
                     pipeline_run=pipeline_run,
@@ -129,6 +131,8 @@ class BlockExecutor:
                 )
             self._execute_callback(
                 'on_success',
+                dynamic_block_index=dynamic_block_index,
+                dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
                 global_vars=global_vars,
                 logging_tags=tags,
                 pipeline_run=pipeline_run,
@@ -194,6 +198,8 @@ class BlockExecutor:
         global_vars,
         logging_tags,
         pipeline_run,
+        dynamic_block_index: Union[int, None] = None,
+        dynamic_upstream_block_uuids: Union[List[str], None] = None,
     ):
         arr = []
         if self.block.callback_block:
@@ -206,6 +212,9 @@ class BlockExecutor:
             try:
                 callback_block.execute_callback(
                     callback,
+                    dynamic_block_index=dynamic_block_index,
+                    dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+                    execution_partition=self.execution_partition,
                     global_vars=global_vars,
                     logger=self.logger,
                     logging_tags=logging_tags,

mage_ai/data_preparation/models/block/__init__.py

@@ -10,7 +10,7 @@ from datetime import datetime
 from inspect import Parameter, isfunction, signature
 from logging import Logger
 from queue import Queue
-from typing import Any, Callable, Dict, List, Set, Tuple
+from typing import Any, Callable, Dict, List, Set, Tuple, Union
 
 import pandas as pd
 import simplejson
@@ -2132,6 +2132,9 @@ class CallbackBlock(Block):
     def execute_callback(
         self,
         callback: str,
+        dynamic_block_index: Union[int, None] = None,
+        dynamic_upstream_block_uuids: Union[List[str], None] = None,
+        execution_partition: str = None,
         global_vars: Dict = None,
         logger: Logger = None,
         logging_tags: Dict = None,
@@ -2194,11 +2197,45 @@ class CallbackBlock(Block):
                 callback_functions_legacy = success_functions
                 callback_status = CallbackStatus.SUCCESS
 
+            # Fetch input variables
+            input_vars, kwargs_vars, upstream_block_uuids = self.fetch_input_variables(
+                None,
+                execution_partition,
+                global_vars,
+                dynamic_block_index=dynamic_block_index,
+                dynamic_upstream_block_uuids=dynamic_upstream_block_uuids,
+            )
+
+            # Copied logic from the method self.execute_block
+            outputs_from_input_vars = {}
+            upstream_block_uuids_length = len(upstream_block_uuids)
+            for idx, input_var in enumerate(input_vars):
+                if idx < upstream_block_uuids_length:
+                    upstream_block_uuid = upstream_block_uuids[idx]
+                    outputs_from_input_vars[upstream_block_uuid] = input_var
+                    outputs_from_input_vars[f'df_{idx + 1}'] = input_var
+
+            global_vars_copy = global_vars.copy()
+            for kwargs_var in kwargs_vars:
+                global_vars_copy.update(kwargs_var)
+
             for callback_function in callback_functions_legacy:
-                callback_function(**global_vars)
+                callback_function(**global_vars_copy)
 
             for callback_function in callback_functions:
-                callback_function(callback_status, **global_vars)
+                try:
+                    # As of version 0.8.81, callback functions have access to the parent block’s
+                    # data output.
+                    callback_function(callback_status, *input_vars, **global_vars_copy)
+                except TypeError:
+                    # This try except block will make the above code backwards compatible in case
+                    # a user has already written callback functions with only keyword arguments.
+                    callback_function(
+                        callback_status,
+                        **merge_dict(global_vars_copy, dict(
+                            __input=outputs_from_input_vars,
+                        )),
+                    )
 
     def update_content(self, content, widget=False):
         if not self.file.exists():

mage_ai/data_preparation/models/pipelines/integration_pipeline.py

@@ -204,11 +204,14 @@ class IntegrationPipeline(Pipeline):
             stderr = e.stderr.decode('utf-8').split('\n')
 
             json_object = {}
+            error = ''
             for line in stderr:
                 if line.startswith('ERROR'):
-                    json_object = next(extract_json_objects(line))
-
-            error = dig(json_object, 'tags.error')
+                    try:
+                        json_object = next(extract_json_objects(line))
+                        error = dig(json_object, 'tags.error')
+                    except Exception:
+                        error = line
             raise Exception(error)
 
     def preview_data(self, block_type: BlockType, streams: List[str] = None) -> List[str]:
@@ -274,11 +277,14 @@ class IntegrationPipeline(Pipeline):
             stderr = e.stderr.decode('utf-8').split('\n')
 
             json_object = {}
+            error = None
             for line in stderr:
                 if line.startswith('ERROR'):
-                    json_object = next(extract_json_objects(line))
-
-            error = dig(json_object, 'tags.error')
+                    try:
+                        json_object = next(extract_json_objects(line))
+                        error = dig(json_object, 'tags.error')
+                    except Exception:
+                        error = line
             if not error:
                 raise Exception('The sample data was not able to be loaded. Please check \
                                 if the stream still exists. If it does not, click the "View and \
@@ -366,17 +372,23 @@ class IntegrationPipeline(Pipeline):
                 message = e.stderr.decode('utf-8')
                 raise Exception(message)
 
-    def streams(self, variables: Dict = {}) -> List[Dict]:
+    def streams(self, variables: Dict = None) -> List[Dict]:
+        if variables is None:
+            variables = {}
         return self.__catalog(variables)['streams']
 
-    def __catalog(self, variables: Dict = {}) -> Dict:
+    def __catalog(self, variables: Dict = None) -> Dict:
+        if variables is None:
+            variables = {}
         return get_catalog(
             self.data_loader,
             self.__global_variables(variables),
             pipeline=self,
         )
 
-    def __global_variables(self, variables: Dict = {}) -> Dict:
+    def __global_variables(self, variables: Dict = None) -> Dict:
+        if variables is None:
+            variables = {}
         d = get_global_variables(self.uuid) or dict()
         d.update(variables)
         return d

mage_ai/data_preparation/models/utils.py

@@ -1,9 +1,11 @@
-from mage_ai.shared.parsers import encode_complex
+import traceback
 from typing import Dict
+
 import dask.dataframe as dd
 import pandas as pd
 import simplejson
-import traceback
+
+from mage_ai.shared.parsers import encode_complex
 
 MAX_PARTITION_BYTE_SIZE = 100 * 1024 * 1024
 JSON_SERIALIZABLE_COLUMN_TYPES = [
@@ -56,7 +58,7 @@ def deserialize_columns(row: pd.Series, column_types: Dict) -> pd.Series:
             continue
 
         val = row[column]
-        if val is not None:
+        if val is not None and type(val) is str:
             row[column] = simplejson.loads(val)
 
     return row

mage_ai/data_preparation/models/variable.py

@@ -1,17 +1,25 @@
+import json
+import os
+import traceback
 from enum import Enum
-from mage_ai.data_cleaner.shared.utils import (
-    is_geo_dataframe,
-    is_spark_dataframe,
-)
+from typing import Any, Dict, List
+
+import numpy as np
+import pandas as pd
+import polars as pl
+from pandas.api.types import is_object_dtype
+from pandas.core.indexes.range import RangeIndex
+
+from mage_ai.data_cleaner.shared.utils import is_geo_dataframe, is_spark_dataframe
 from mage_ai.data_preparation.models.constants import (
     DATAFRAME_ANALYSIS_KEYS,
     DATAFRAME_SAMPLE_COUNT,
     DATAFRAME_SAMPLE_MAX_COLUMNS,
     VARIABLE_DIR,
 )
-from mage_ai.data_preparation.models.utils import (
+from mage_ai.data_preparation.models.utils import (  # dask_from_pandas,
+    STRING_SERIALIZABLE_COLUMN_TYPES,
     apply_transform_pandas,
-    # dask_from_pandas,
     cast_column_types,
     deserialize_columns,
     serialize_columns,
@@ -20,14 +28,6 @@ from mage_ai.data_preparation.storage.base_storage import BaseStorage
 from mage_ai.data_preparation.storage.local_storage import LocalStorage
 from mage_ai.shared.parsers import sample_output
 from mage_ai.shared.utils import clean_name
-from pandas.api.types import is_object_dtype
-from typing import Any, Dict, List
-import json
-import numpy as np
-import os
-import pandas as pd
-import polars as pl
-import traceback
 
 DATAFRAME_COLUMN_TYPES_FILE = 'data_column_types.json'
 DATAFRAME_PARQUET_FILE = 'data.parquet'
@@ -54,11 +54,14 @@ class Variable:
         block_uuid: str,
         partition: str = None,
         spark=None,
-        storage: BaseStorage = LocalStorage(),
+        storage: BaseStorage = None,
         variable_type: VariableType = None
     ) -> None:
         self.uuid = uuid
-        self.storage = storage
+        if storage is None:
+            self.storage = LocalStorage()
+        else:
+            self.storage = storage
         # if not self.storage.path_exists(pipeline_path):
         #     raise Exception(f'Pipeline path {pipeline_path} does not exist.')
         self.pipeline_path = pipeline_path
@@ -269,7 +272,9 @@ class Variable:
             self.storage.remove(file_path)
             self.storage.remove_dir(self.variable_path)
 
-    def __read_json(self, default_value={}, sample: bool = False) -> Dict:
+    def __read_json(self, default_value: Dict = None, sample: bool = False) -> Dict:
+        if default_value is None:
+            default_value = {}
         # For backward compatibility
         old_file_path = os.path.join(self.variable_dir_path, f'{self.uuid}.json')
         file_path = os.path.join(self.variable_path, JSON_FILE)
@@ -291,7 +296,9 @@ class Variable:
             data = sample_output(data)[0]
         return data
 
-    async def __read_json_async(self, default_value={}, sample: bool = False) -> Dict:
+    async def __read_json_async(self, default_value: Dict = None, sample: bool = False) -> Dict:
+        if default_value is None:
+            default_value = {}
         # For backward compatibility
         old_file_path = os.path.join(self.variable_dir_path, f'{self.uuid}.json')
         file_path = os.path.join(self.variable_path, JSON_FILE)
@@ -352,7 +359,12 @@ class Variable:
                 df = df.iloc[:sample_count]
         return df
 
-    def __read_parquet(self, sample: bool = False, sample_count: int = None) -> pd.DataFrame:
+    def __read_parquet(
+        self,
+        sample: bool = False,
+        sample_count: int = None,
+        raise_exception: bool = False,
+    ) -> pd.DataFrame:
         file_path = os.path.join(self.variable_path, DATAFRAME_PARQUET_FILE)
         sample_file_path = os.path.join(self.variable_path, DATAFRAME_PARQUET_SAMPLE_FILE)
 
@@ -361,12 +373,16 @@ class Variable:
             try:
                 df = self.storage.read_parquet(sample_file_path, engine='pyarrow')
                 read_sample_success = True
-            except Exception:
+            except Exception as e:
+                if raise_exception:
+                    raise e
                 pass
         if not read_sample_success:
             try:
                 df = self.storage.read_parquet(file_path, engine='pyarrow')
-            except Exception:
+            except Exception as e:
+                if raise_exception:
+                    raise e
                 df = pd.DataFrame()
         if sample:
             sample_count = sample_count or DATAFRAME_SAMPLE_COUNT
@@ -415,10 +431,14 @@ class Variable:
                 series_non_null = df_output[c].dropna()
                 if len(series_non_null) > 0:
                     coltype = type(series_non_null.iloc[0])
-
                     if is_object_dtype(series_non_null.dtype):
+                        if coltype.__name__ in STRING_SERIALIZABLE_COLUMN_TYPES:
+                            cast_coltype = str
+                        else:
+                            cast_coltype = coltype
                         try:
-                            df_output[c] = series_non_null.astype(coltype)
+                            df_output[c] = series_non_null.astype(cast_coltype)
+                            coltype = str
                         except Exception:
                             # Fall back to convert to string
                             # df_output[c] = series_non_null.astype(str)
@@ -432,6 +452,22 @@ class Variable:
                         column_types[c] = type(series_non_null.iloc[0].item()).__name__
 
         self.storage.makedirs(self.variable_path, exist_ok=True)
+        with open(os.path.join(self.variable_path, DATAFRAME_COLUMN_TYPES_FILE), 'w') as f:
+            f.write(json.dumps(column_types))
+
+        # Try using Polars to write the dataframe to improve performance
+        if type(df_output.index) is RangeIndex and df_output.index.start == 0 \
+                and df_output.index.stop == df_output.shape[0] and df_output.index.step == 1:
+            # Polars ignores any index
+            try:
+                pl_df = pl.from_pandas(df_output)
+                self.__write_polars_dataframe(pl_df)
+                # Test read dataframe from parquet
+                self. __read_parquet(sample=True, raise_exception=True)
+
+                return
+            except Exception:
+                pass
 
         # ddf = dask_from_pandas(df_output)
         df_output_serialized = apply_transform_pandas(
@@ -444,9 +480,6 @@ class Variable:
             os.path.join(self.variable_path, DATAFRAME_PARQUET_FILE),
         )
 
-        with open(os.path.join(self.variable_path, DATAFRAME_COLUMN_TYPES_FILE), 'w') as f:
-            f.write(json.dumps(column_types))
-
         try:
             df_sample_output = df_output_serialized.iloc[
                 :DATAFRAME_SAMPLE_COUNT,

mage_ai/data_preparation/preferences.py

@@ -1,10 +1,13 @@
+import os
+import traceback
 from typing import Dict
+
+import yaml
+
 from mage_ai.data_preparation.models.constants import PREFERENCES_FILE
 from mage_ai.data_preparation.repo_manager import get_repo_path
 from mage_ai.orchestration.db.models.oauth import User
-import os
-import traceback
-import yaml
+from mage_ai.shared.hash import merge_dict
 
 # Git environment variables
 GIT_REPO_LINK_VAR = 'GIT_REPO_LINK'
@@ -27,19 +30,20 @@ class Preferences:
         self.preferences_file_path = \
             os.path.join(self.repo_path, PREFERENCES_FILE)
         self.user = user
-        preferences = dict()
+        project_preferences = dict()
         try:
-            if user:
-                preferences = user.preferences or {}
+            if user and user.preferences and user.git_settings is None:
+                project_preferences = user.preferences
             elif config_dict:
-                preferences = config_dict
+                project_preferences = config_dict
             elif os.path.exists(self.preferences_file_path):
                 with open(self.preferences_file_path) as f:
-                    preferences = yaml.full_load(f.read()) or {}
+                    project_preferences = yaml.full_load(f.read()) or {}
         except Exception:
             traceback.print_exc()
             pass
 
+        # Git settings
         if os.getenv(GIT_REPO_LINK_VAR):
             self.sync_config = dict(
                 remote_repo_link=os.getenv(GIT_REPO_LINK_VAR),
@@ -51,19 +55,23 @@ class Preferences:
                 sync_on_pipeline_run=bool(int(os.getenv(GIT_SYNC_ON_PIPELINE_RUN_TYPE) or 0)),
             )
         else:
-            self.sync_config = preferences.get('sync_config', dict())
+            project_sync_config = project_preferences.get('sync_config', dict())
+            if user:
+                user_git_settings = user.git_settings or {}
+                self.sync_config = merge_dict(project_sync_config, user_git_settings)
+            else:
+                self.sync_config = project_sync_config
 
-    def has_valid_git_config(self) -> bool:
-        return 'remote_repo_link' in self.sync_config and 'repo_path' in self.sync_config
+    def is_git_integration_enabled(self) -> bool:
+        return 'remote_repo_link' in self.sync_config and \
+            'repo_path' in self.sync_config and \
+            self.sync_config.get('branch') is None
 
     def update_preferences(self, updates: Dict):
         preferences = self.to_dict()
         preferences.update(updates)
-        if self.user:
-            self.user.update(preferences=preferences)
-        else:
-            with open(self.preferences_file_path, 'w') as f:
-                yaml.dump(preferences, f)
+        with open(self.preferences_file_path, 'w') as f:
+            yaml.dump(preferences, f)
 
     def to_dict(self) -> Dict:
         return dict(

mage_ai/data_preparation/sync/__init__.py

@@ -1,7 +1,9 @@
+import inspect
+import os
 from dataclasses import dataclass
 from enum import Enum
+
 from mage_ai.shared.config import BaseConfig
-import os
 
 GIT_ACCESS_TOKEN_SECRET_NAME = 'mage_git_access_token'
 GIT_SSH_PRIVATE_KEY_SECRET_NAME = 'mage_git_ssh_private_key_b64'
@@ -17,13 +19,30 @@ class AuthType(str, Enum):
 class GitConfig(BaseConfig):
     remote_repo_link: str
     repo_path: str = os.getcwd()
-    username: str = ''
-    email: str = ''
     branch: str = 'main'
     sync_on_pipeline_run: bool = False
+    auth_type: AuthType = AuthType.SSH
+    # User settings moved to UserGitConfig, these will be used for Git syncs
+    username: str = ''
+    email: str = ''
     ssh_private_key_secret_name: str = GIT_SSH_PRIVATE_KEY_SECRET_NAME
     ssh_public_key_secret_name: str = GIT_SSH_PUBLIC_KEY_SECRET_NAME
-    auth_type: AuthType = AuthType.SSH
     access_token_secret_name: str = GIT_ACCESS_TOKEN_SECRET_NAME
     # This is not necessary anymore, but leaving it for backwards compatibility
     type: str = 'git'
+
+
+@dataclass
+class UserGitConfig(BaseConfig):
+    username: str = ''
+    email: str = ''
+    ssh_private_key_secret_name: str = GIT_SSH_PRIVATE_KEY_SECRET_NAME
+    ssh_public_key_secret_name: str = GIT_SSH_PUBLIC_KEY_SECRET_NAME
+    access_token_secret_name: str = GIT_ACCESS_TOKEN_SECRET_NAME
+
+    @classmethod
+    def from_dict(cls, config):
+        return cls(**{
+            k: v for k, v in config.items()
+            if k in inspect.signature(cls).parameters
+        })

mage_ai/data_preparation/sync/git_sync.py

@@ -16,3 +16,10 @@ class GitSync(BaseSync):
             verbose=True,
         ):
             self.git_manager.reset(self.branch)
+
+    def reset(self):
+        with VerboseFunctionExec(
+            f'Attempting to clone from remote repo {self.remote_repo_link}',
+            verbose=True,
+        ):
+            self.git_manager.clone()