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

mage_ai/settings/server.py

@@ -1,14 +1,50 @@
+"""
+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
+"""
+
 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
+
+def get_bool_value(value: str) -> bool:
+    """
+    Helper method to get the boolean value from a string.
+
+    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')
+
+
+def get_int_value(value: str) -> int:
+    if not value:
+        return None
+    try:
+        int_value = int(value)
+    except Exception:
+        int_value = None
+    return int_value
+
+
+# ----------------------------------
+# Debugging and Environment Settings
+# ----------------------------------
 
 DEBUG = os.getenv('DEBUG', False)
+DEBUG_MEMORY = str(os.getenv('DEBUG_MEMORY', 0) or 0) in ['1', 'true', 'True']
+DEBUG_FILE_IO = str(os.getenv('DEBUG_FILE_IO', 0) or 0) in ['1', 'true', 'True']
 HIDE_ENV_VAR_VALUES = int(os.getenv('HIDE_ENV_VAR_VALUES', 1) or 1) == 1
-QUERY_API_KEY = 'api_key'
+
+
+# -------------------------
+# Server Authentication and Token Settings
+# -------------------------
 
 """
 import secrets
@@ -18,6 +54,18 @@ Make sure this value is the same in mage_ai/frontend/api/constants.ts
 """
 OAUTH2_APPLICATION_CLIENT_ID = 'zkWlN0PkIKSN0C11CfUHUj84OT5XOJ6tDZ6bDRO2'
 
+QUERY_API_KEY = 'api_key'
+
+# 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())
+
+
+# -------------------------
+# Notebook and Terminal Edit Access
+# -------------------------
+
 # valid values: 0, 1, 2
 try:
     DISABLE_NOTEBOOK_EDIT_ACCESS = int(os.getenv('DISABLE_NOTEBOOK_EDIT_ACCESS', 0))
@@ -25,6 +73,9 @@ except ValueError:
     DISABLE_NOTEBOOK_EDIT_ACCESS = 1 if os.getenv('DISABLE_NOTEBOOK_EDIT_ACCESS') else 0
 
 
+DISABLE_TERMINAL = get_bool_value(os.getenv('DISABLE_TERMINAL', '0').lower())
+
+
 def is_disable_pipeline_edit_access(
     disable_notebook_edit_access_override: int = None,
 ) -> bool:
@@ -34,33 +85,22 @@ def is_disable_pipeline_edit_access(
     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 ----------------------
+# Limit the image preview size to 5MBs by default. If a block returns an XBoost model
+# (e.g.) xgboost.Booster() instance, the preview in the UI is a tree plot of the model.
+# The image can be very large in byte size, so we limit the size to prevent the UI from freezing.
+MAX_OUTPUT_IMAGE_PREVIEW_SIZE = int(os.getenv('MAX_OUTPUT_IMAGE_PREVIEW_SIZE', 1024 * 1024 * 5))
 
-DISABLE_TERMINAL = os.getenv('DISABLE_TERMINAL', '0').lower() in ('true', '1', 't')
+# -------------------
+# User Authentication Settings
+# -------------------
 
-# ----------------- Authentication settings ----------------
-REQUIRE_USER_AUTHENTICATION = get_bool_value(
-    os.getenv('REQUIRE_USER_AUTHENTICATION', 'False')
-)
+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')
-    )
+    MAGE_ACCESS_TOKEN_EXPIRY_TIME = int(os.getenv('MAGE_ACCESS_TOKEN_EXPIRY_TIME', '2592000'))
 except ValueError:
     MAGE_ACCESS_TOKEN_EXPIRY_TIME = 2592000
 
@@ -68,9 +108,22 @@ except ValueError:
 # 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')
 
-# ----------------------------------------------------------
 
+# ---------------------
+# General Server Settings
+# ---------------------
+CONCURRENCY_CONFIG_BLOCK_RUN_LIMIT = get_int_value(os.getenv('CONCURRENCY_CONFIG_BLOCK_RUN_LIMIT'))
+CONCURRENCY_CONFIG_PIPELINE_RUN_LIMIT = get_int_value(
+    os.getenv('CONCURRENCY_CONFIG_PIPELINE_RUN_LIMIT')
+)
+DISABLE_AUTO_BROWSER_OPEN = get_bool_value(os.getenv('DISABLE_AUTO_BROWSER_OPEN', 'False'))
+DISABLE_AUTORELOAD = get_bool_value(os.getenv('DISABLE_AUTORELOAD', 'False'))
+# The hostname in Kubernetes or AWS ECS
 HOSTNAME = os.getenv('HOSTNAME')
+INITIAL_METADATA = os.getenv('INITIAL_METADATA')
+LOGS_DIR_PATH = os.getenv('LOGS_DIR_PATH')
+MAX_FILE_CACHE_SIZE = os.getenv('MAX_FILE_CACHE_SIZE') or (1024 * 1024)  # 1 MB
+ENABLE_USER_PROJECTS = get_bool_value(os.getenv('ENABLE_USER_PROJECTS'))
 REDIS_URL = os.getenv('REDIS_URL')
 SERVER_VERBOSITY = os.getenv('SERVER_VERBOSITY', 'info') or 'info'
 SERVER_LOGGING_FORMAT = os.getenv('SERVER_LOGGING_FORMAT', 'plaintext')
@@ -79,36 +132,49 @@ SERVER_LOGGING_TEMPLATE = os.getenv(
     '%(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 Settings
+# --------------------
 
 SHELL_COMMAND = os.getenv('SHELL_COMMAND', None)
 USE_UNIQUE_TERMINAL = os.getenv('USE_UNIQUE_TERMINAL', None)
 
-# sentry environment variables
+
+# -------------------------
+# Monitoring and Tracing Configuration
+# -------------------------
+
+# Sentry Configuration
 SENTRY_DSN = os.getenv('SENTRY_DSN')
 SENTRY_TRACES_SAMPLE_RATE = os.getenv('SENTRY_TRACES_SAMPLE_RATE', 1.0)
 
-# New relic enable environment variable
+# New Relic Configuration
 ENABLE_NEW_RELIC = os.getenv('ENABLE_NEW_RELIC', False)
 NEW_RELIC_CONFIG_PATH = os.getenv('NEW_RELIC_CONFIG_PATH', '')
 
+# Prometheus Configuration
 # If enabled, the /metrics route will expose Tornado server metrics
 ENABLE_PROMETHEUS = get_bool_value(os.getenv('ENABLE_PROMETHEUS', 'False'))
 
-# OpenTelemetry configuration
+# 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
+SYSTEM_LOGS_PARTITIONS = [
+    str(partition).strip() for partition in os.getenv('SYSTEM_LOGS_PARTITIONS', 'ds').split(',')
+]
+SYSTEM_LOGS_POLL_INTERVAL = float(os.getenv('SYSTEM_LOGS_POLL_INTERVAL', 0.1))
 
-MAX_FILE_CACHE_SIZE = os.getenv('MAX_FILE_CACHE_SIZE') or (1024 * 1024)  # 1 MB
+# -----------------------------
+# Mage URL and PATH Settings
+# -----------------------------
 
+DEFAULT_LOCALHOST_URL = 'http://localhost:6789'
+MAGE_PUBLIC_HOST = os.getenv('MAGE_PUBLIC_HOST') or DEFAULT_LOCALHOST_URL
 # 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')
@@ -118,10 +184,12 @@ 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())
+
+
+# -------------------------
+# Scheduler Settings
+# -------------------------
+
 # Sets the trigger interval of the scheduler to a numeric value, in seconds
 # Determines how often the scheduler gets invoked
 try:
@@ -129,11 +197,32 @@ try:
 except ValueError:
     SCHEDULER_TRIGGER_INTERVAL = 10
 
+# -------------------------
+# System level features
+# -------------------------
+# We need to use os.getenv again or else we can’t mock/patch the value in tests.
+DYNAMIC_BLOCKS_VERSION = int(os.getenv('DYNAMIC_BLOCKS_VERSION') or 1)
+DYNAMIC_BLOCKS_V2 = int(os.getenv('DYNAMIC_BLOCKS_VERSION') or 1) >= 2
+MEMORY_MANAGER_VERSION = int(os.getenv('MEMORY_MANAGER_VERSION') or 1)
+MEMORY_MANAGER_V2 = int(os.getenv('MEMORY_MANAGER_VERSION') or 1) >= 2
+MEMORY_MANAGER_PANDAS_VERSION = int(os.getenv('MEMORY_MANAGER_PANDAS_VERSION') or 1)
+MEMORY_MANAGER_POLARS_VERSION = int(os.getenv('MEMORY_MANAGER_POLARS_VERSION') or 1)
+MEMORY_MANAGER_PANDAS_V2 = int(os.getenv('MEMORY_MANAGER_PANDAS_VERSION') or 1) >= 2
+MEMORY_MANAGER_POLARS_V2 = int(os.getenv('MEMORY_MANAGER_POLARS_VERSION') or 1) >= 2
+VARIABLE_DATA_OUTPUT_META_CACHE = str(os.getenv('VARIABLE_DATA_OUTPUT_META_CACHE', 0) or 0) in [
+    '1',
+    'true',
+    'True',
+]
+
 # List of environment variables used to configure Mage. The value of these settings
 # will be copied between workspaces.
 MAGE_SETTINGS_ENVIRONMENT_VARIABLES = [
+    'CONCURRENCY_CONFIG_BLOCK_RUN_LIMIT',
+    'CONCURRENCY_CONFIG_PIPELINE_RUN_LIMIT',
     'DISABLE_NOTEBOOK_EDIT_ACCESS',
     'DISABLE_AUTO_BROWSER_OPEN',
+    'DISABLE_AUTORELOAD',
     'REQUIRE_USER_AUTHENTICATION',
     'AUTHENTICATION_MODE',
     'OAUTH_DEFAULT_ACCESS',
@@ -174,6 +263,7 @@ MAGE_SETTINGS_ENVIRONMENT_VARIABLES = [
     'OIDC_CLIENT_ID',
     'OIDC_CLIENT_SECRET',
     'OIDC_DISCOVERY_URL',
+    'OIDC_ROLES_MAPPING',
     'GHE_CLIENT_ID',
     'GHE_CLIENT_SECRET',
     'GHE_HOSTNAME',
@@ -184,4 +274,5 @@ MAGE_SETTINGS_ENVIRONMENT_VARIABLES = [
     'GITLAB_CLIENT_ID',
     'GITLAB_CLIENT_SECRET',
     'SERVER_LOGGING_TEMPLATE',
+    'MAX_OUTPUT_IMAGE_PREVIEW_SIZE',
 ]

mage_ai/shared/array.py

@@ -1,10 +1,12 @@
 import random
+from collections.abc import Iterable
+from typing import Any
 
 
 def batch(iterable, n=1):
     length = len(iterable)
     for ndx in range(0, length, n):
-        yield iterable[ndx:min(ndx + n, length)]
+        yield iterable[ndx: min(ndx + n, length)]
 
 
 def difference(li1, li2):
@@ -50,3 +52,24 @@ def unique_by(arr1, key):
         arr2.append(item)
         mapping[k] = True
     return arr2
+
+
+def is_iterable(
+    obj: Any, exclude_dicts: bool = True, exclude_strings: bool = True
+) -> bool:
+    """
+    Check if an object is iterable, optionally excluding strings.
+
+    Args:
+        obj: The object to check for iterability.
+        exclude_strings (bool): Whether to consider strings as non-iterable (to exclude them).
+
+    Returns:
+        bool: True if the object is iterable; False otherwise.
+    """
+    if (exclude_strings and isinstance(obj, str)) or (
+        exclude_dicts and isinstance(obj, dict)
+    ):
+        return False
+
+    return isinstance(obj, Iterable)

mage_ai/shared/complex.py

@@ -0,0 +1,45 @@
+import inspect
+from typing import Any
+
+
+def is_model_sklearn(data: Any) -> bool:
+    if inspect.isclass(data):
+        return False
+
+    try:
+        from sklearn.base import BaseEstimator, is_classifier, is_regressor
+
+        return (
+            is_classifier(data) or is_regressor(data) or isinstance(data, BaseEstimator)
+        )
+    except ImportError as err:
+        print(f"Error importing sklearn: {err}")
+        return False
+
+
+def is_model_xgboost(data: Any) -> bool:
+    """
+    Checks if the given data is an instance of an XGBoost model, either a Booster
+    or an object from XGBoost's scikit-learn API.
+    """
+    if inspect.isclass(data):
+        return False
+
+    # Check for direct instance of Booster
+    if inspect.isclass(data):
+        # Checking if it's a class reference rather than an instance
+        if "xgboost.core.Booster" == f"{data.__module__}.{data.__qualname__}":
+            return True
+    else:
+        # Check based on instance attributes and methods
+        if hasattr(data, "__class__"):
+            class_name = f"{data.__class__.__module__}.{data.__class__.__qualname__}"
+            if class_name.startswith("xgboost.core.Booster"):
+                return True
+
+    # Check for sklearn API models (like XGBClassifier, XGBRegressor)
+    # These models have a get_booster() method
+    if hasattr(data, "get_booster"):
+        return True
+
+    return False

mage_ai/shared/config.py

@@ -7,6 +7,7 @@ import yaml
 
 from mage_ai.shared.hash import merge_dict
 from mage_ai.shared.strings import camel_to_snake_case
+from mage_ai.shared.yaml import load_yaml
 
 
 @dataclass
@@ -77,7 +78,7 @@ class BaseConfig:
 
         try:
             with open(config_path, 'r') as stream:
-                file_config = yaml.load(stream, Loader=yaml.SafeLoader)
+                file_config = load_yaml(stream)
                 return merge_dict(config, file_config)
         except yaml.YAMLError as e:
             print(e)

mage_ai/shared/custom_logger.py

@@ -150,6 +150,17 @@ class ColorPrinter:
             __uuid='get_callers'
         )
 
+    def print_call_stack(self):
+        stack = inspect.stack()
+        stack_list = []
+        for frame_info in stack:
+            frame = frame_info.frame
+            filename = frame.f_code.co_filename
+            lineno = frame.f_lineno
+            function_name = frame.f_code.co_name
+            stack_list.append(f"File '{filename}', line {lineno}, in {function_name}")
+        self.info('Call stack:', stack=stack_list)
+
 
 class CustomFormatter(logging.Formatter):
     # https://stackoverflow.com/a/56944256/1084284

mage_ai/shared/dates.py

@@ -1,10 +1,12 @@
 import datetime
+from typing import List, Optional, Tuple
+
 import pytz
 
 from mage_ai.shared.array import find_index
 
 
-def compare(date1: datetime, date2: datetime) -> int:
+def compare(date1, date2) -> Optional[int]:
     if date1 is None or date2 is None:
         return None
 
@@ -27,8 +29,7 @@ def n_days_ago(n: int) -> str:
 def str_to_timedelta(period_str: str):
     unit = period_str[-1]
     if unit not in ['d', 'h', 'w']:
-        raise Exception(
-            'Please provide a valid period unit ("d", "h", or "w")')
+        raise Exception('Please provide a valid period unit ("d", "h", or "w")')
     if unit == 'd':
         return datetime.timedelta(days=int(period_str[:-1]))
     elif unit == 'h':
@@ -45,8 +46,8 @@ def week_of_month(day: datetime.datetime) -> int:
 
     # e.g. 7 - 5 = 2
     first_week_last_day = 7 - ((first_day_weekday + 1) % 7)
-    day_ranges_per_week = [
-        # If the 1st day is Friday, then Saturday is the 2nd, and that is the end of the 1st week
+    # If the 1st day is Friday, then Saturday is the 2nd, and that is the end of the 1st week
+    day_ranges_per_week: List[Tuple[int, int]] = [
         (1, first_week_last_day),
     ]
 
@@ -57,7 +58,10 @@ def week_of_month(day: datetime.datetime) -> int:
         # 3: (24, 30)
         # 4: (31, 37)
         day_ranges_per_week.append(
-            (first_week_last_day + 1 + (7 * i), first_week_last_day + (7 * (i + 1))),
+            (
+                first_week_last_day + 1 + (7 * i),
+                first_week_last_day + (7 * (i + 1)),
+            ),
         )
 
     # If the first day of the month is a Friday, here are the day ranges for each week:

mage_ai/shared/files.py

@@ -1,5 +1,8 @@
+import asyncio
 import glob
 import os
+import shutil
+import time
 from pathlib import Path
 from typing import Callable, Dict, List, Tuple
 
@@ -55,10 +58,15 @@ def read_last_line(filename: str) -> str:
 
 
 def get_full_file_paths_containing_item(root_full_path: str, comparator: Callable) -> List[str]:
-    configfiles = [os.path.join(
-        dirpath,
-        f,
-    ) for dirpath, dirnames, files in os.walk(root_full_path) for f in files if comparator(f)]
+    configfiles = [
+        os.path.join(
+            dirpath,
+            f,
+        )
+        for dirpath, dirnames, files in os.walk(root_full_path)
+        for f in files
+        if comparator(f)
+    ]
 
     return configfiles
 
@@ -101,10 +109,11 @@ def get_absolute_paths_from_all_files(
     for filename in glob.iglob(dir_path, recursive=True):
         absolute_path = os.path.abspath(filename)
 
-        if os.path.isfile(absolute_path) and \
-                (not include_hidden_files or not absolute_path.startswith('.')) and \
-                (not comparator or comparator(absolute_path)):
-
+        if (
+            os.path.isfile(absolute_path)
+            and (not include_hidden_files or not absolute_path.startswith('.'))
+            and (not comparator or comparator(absolute_path))
+        ):
             value = (absolute_path, os.path.getsize(filename), round(os.path.getmtime(filename)))
             arr.append(parse_values(value) if parse_values else value)
 
@@ -149,3 +158,49 @@ async def read_async(file_path: str) -> str:
         except Exception as err:
             if is_debug():
                 print(f'[ERROR] files.read_async: {err}.')
+
+
+def safe_delete_dir_sync(output_dir, verbose: bool = False):
+    max_attempts = 5
+    for attempt in range(max_attempts):
+        try:
+            shutil.rmtree(output_dir)
+            if verbose:
+                print(f'Successfully deleted {output_dir}')
+            break  # Successfully deleted; exit the loop
+        except OSError as err:
+            if attempt < max_attempts - 1:
+                time.sleep(2)  # Wait a bit for retry
+                continue
+            else:
+                print(f'Failed to delete {output_dir} after {max_attempts} attempts: {err}')
+                raise err
+
+
+async def safe_delete_dir_async(output_dir: str, verbose: bool = False):
+    loop = asyncio.get_event_loop()
+    max_attempts = 5
+
+    for attempt in range(max_attempts):
+        try:
+            await loop.run_in_executor(None, shutil.rmtree, output_dir)
+            if verbose:
+                print(f'Successfully deleted {output_dir}')
+            break  # Successfully deleted; exit the loop
+        except OSError as err:
+            # Handles the potential need for retry due to file locking/availability/etc.
+            if attempt < max_attempts - 1:
+                await asyncio.sleep(2)  # Non-blocking wait before retry
+                continue
+            else:
+                print(f'Failed to delete {output_dir} after {max_attempts} attempts: {err}')
+                raise
+
+
+def makedirs_sync(path: str):
+    os.makedirs(path, exist_ok=True)
+
+
+async def makedirs_async(path: str):
+    loop = asyncio.get_running_loop()
+    await loop.run_in_executor(None, os.makedirs, path, exist_ok=True)

mage_ai/shared/hash.py

@@ -23,11 +23,11 @@ def camel_case_keys_to_snake_case(d):
 
 def dig(obj_arg, arr_or_string):
     if type(arr_or_string) is str:
-        arr_or_string = arr_or_string.split('.')
+        arr_or_string = arr_or_string.split(".")
     arr = list(map(str.strip, arr_or_string))
 
     def _build(obj, key):
-        tup = re.split(r'\[(\d+)\]$', key)
+        tup = re.split(r"\[(\d+)\]$", key)
         if len(tup) >= 2:
             key, index = filter(lambda x: x, tup)
             if key and index:
@@ -56,7 +56,7 @@ def safe_dig(obj_arg, arr_or_string):
             key/index is missing or the object is None.
     """
     if isinstance(arr_or_string, str):
-        arr_or_string = arr_or_string.split('.')
+        arr_or_string = arr_or_string.split(".")
     arr = list(map(str.strip, arr_or_string))
 
     def _build(obj, key):
@@ -64,7 +64,7 @@ def safe_dig(obj_arg, arr_or_string):
         if obj is None or not isinstance(obj, dict) and not isinstance(obj, list):
             return None
 
-        tup = re.split(r'\[(\d+)\]$', key)
+        tup = re.split(r"\[(\d+)\]$", key)
         if len(tup) >= 2:
             key, index = filter(lambda x: x, tup)
             index = int(index) if index else None
@@ -96,15 +96,39 @@ def flatten(input_data):
             for k2, v2 in v1.items():
                 if type(v2) is dict:
                     for k3, v3 in v2.items():
-                        final_data[f'{k1}_{k2}_{k3}'] = v3
+                        final_data[f"{k1}_{k2}_{k3}"] = v3
                 else:
-                    final_data[f'{k1}_{k2}'] = v2
+                    final_data[f"{k1}_{k2}"] = v2
         else:
             final_data[k1] = v1
 
     return final_data
 
 
+def flatten_dict(d, parent_key="", sep="."):
+    items = []
+    for k, v in d.items():
+        new_key = f"{parent_key}{sep}{k}" if parent_key else k
+        if isinstance(v, dict):
+            items.extend(flatten_dict(v, new_key, sep=sep).items())
+        else:
+            items.append((new_key, v))
+    return dict(items)
+
+
+def unflatten_dict(d, sep="."):
+    result_dict = {}
+    for k, v in d.items():
+        parts = k.split(sep)
+        current_level = result_dict
+        for part in parts[:-1]:
+            if part not in current_level:
+                current_level[part] = {}
+            current_level = current_level[part]
+        current_level[parts[-1]] = v
+    return result_dict
+
+
 def get_json_value(str_value, arr_or_string):
     if not str_value:
         return str_value
@@ -207,11 +231,11 @@ def set_value(obj: Dict, keys: List[str], value) -> Dict:
 
     results = dict(__obj_to_set_value=obj, __value=value)
 
-    key = ''.join(f"['{key}']" for key in keys)
-    expression = f'__obj_to_set_value{key} = __value'
+    key = "".join(f"['{key}']" for key in keys)
+    expression = f"__obj_to_set_value{key} = __value"
     exec(expression, results)
 
-    return results['__obj_to_set_value']
+    return results["__obj_to_set_value"]
 
 
 def combine_into(child: Dict, parent: Dict) -> None:

mage_ai/shared/io.py

@@ -1,7 +1,7 @@
 import os
 import shutil
 import traceback
-from typing import Callable, Union
+from typing import Callable, Optional, Union
 
 import aiofiles
 
@@ -29,7 +29,9 @@ def chmod(path: Union[str, os.PathLike], mode: int, append: bool = True) -> None
         traceback.print_exc()
 
 
-def safe_write(filepath: str, content: str, write_func: Callable = None):
+def safe_write(
+    filepath: str, content: str, write_func: Optional[Callable] = None, write_mode: str = 'w'
+):
     temp_file_path = filepath + '.temp'
     if os.path.isfile(filepath):
         shutil.copy2(filepath, temp_file_path)
@@ -39,7 +41,7 @@ def safe_write(filepath: str, content: str, write_func: Callable = None):
 
     success = False
     try:
-        with open(filepath, 'w', encoding='utf-8') as fp:
+        with open(filepath, write_mode, encoding='utf-8') as fp:
             if write_func is not None:
                 write_func(fp, content)
             else:
@@ -57,7 +59,9 @@ def safe_write(filepath: str, content: str, write_func: Callable = None):
             traceback.print_exc()
 
 
-async def safe_write_async(filepath: str, content: str, write_func: Callable = None):
+async def safe_write_async(
+    filepath: str, content: str, write_func: Optional[Callable] = None, write_mode: str = 'w'
+):
     temp_file_path = filepath + '.temp'
     if os.path.isfile(filepath):
         shutil.copy2(filepath, temp_file_path)
@@ -67,7 +71,7 @@ async def safe_write_async(filepath: str, content: str, write_func: Callable = N
 
     success = False
     try:
-        async with aiofiles.open(filepath, mode='w', encoding='utf-8') as fp:
+        async with aiofiles.open(filepath, mode=write_mode, encoding='utf-8') as fp:
             if write_func is not None:
                 await write_func(fp, content)
             else: