mage-ai 0.9.73__py3-none-any.whl → 0.9.75__py3-none-any.whl

mage_ai/data_preparation/preferences.py

@@ -37,15 +37,13 @@ ENV_VAR_TO_CONFIG_KEY = {
     GIT_SYNC_SUBMODULES: 'sync_submodules',
 }
 
-BOOLEAN_ENV_VARS = set(
-    [
-        GIT_SYNC_ON_PIPELINE_RUN,
-        GIT_SYNC_ON_START,
-        GIT_SYNC_ON_EXECUTOR_START,
-        GIT_SYNC_SUBMODULES,
-        GIT_ENABLE_GIT_INTEGRATION,
-    ]
-)
+BOOLEAN_ENV_VARS = set([
+    GIT_SYNC_ON_PIPELINE_RUN,
+    GIT_SYNC_ON_START,
+    GIT_SYNC_ON_EXECUTOR_START,
+    GIT_SYNC_SUBMODULES,
+    GIT_ENABLE_GIT_INTEGRATION,
+])
 
 
 def get_value_for_sync_config(env_var, value) -> bool:
@@ -67,10 +65,7 @@ def build_sync_config(project_sync_config: Dict) -> Dict:
     }
 
     # Read settings from settings backend
-    config_from_settings = {
-        k: get_settings_value(k)
-        for k in ENV_VAR_TO_CONFIG_KEY.keys()
-    }
+    config_from_settings = {k: get_settings_value(k) for k in ENV_VAR_TO_CONFIG_KEY.keys()}
 
     if any([v is not None for v in config_from_settings.values()]):
         # Map sync config keys to settings values
@@ -153,9 +148,7 @@ class Preferences:
 def get_preferences(repo_path=None, user: User = None) -> Preferences:
     default_preferences = Preferences(repo_path=repo_path)
     if user:
-        if user.preferences is None and os.path.exists(
-            default_preferences.preferences_file_path
-        ):
+        if user.preferences is None and os.path.exists(default_preferences.preferences_file_path):
             return default_preferences
         else:
             return Preferences(user=user, repo_path=repo_path)

mage_ai/data_preparation/repo_manager.py

@@ -3,7 +3,6 @@ import logging
 import os
 import traceback
 import uuid
-from enum import Enum
 from typing import Dict, Optional
 from warnings import warn
 
@@ -25,6 +24,7 @@ from mage_ai.settings.repo import get_repo_path as get_repo_path_new
 from mage_ai.settings.repo import get_variables_dir
 from mage_ai.settings.repo import set_repo_path as set_repo_path_new
 from mage_ai.settings.utils import base_repo_path
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.environments import is_debug
 from mage_ai.shared.yaml import load_yaml, trim_strings
 
@@ -36,7 +36,7 @@ yml.indent(mapping=2, sequence=2, offset=0)
 logger = logging.getLogger(__name__)
 
 
-class ProjectType(str, Enum):
+class ProjectType(StrEnum):
     MAIN = 'main'
     SUB = 'sub'
     STANDALONE = 'standalone'

mage_ai/data_preparation/sync/__init__.py

@@ -1,16 +1,16 @@
 import inspect
 import os
 from dataclasses import dataclass
-from enum import Enum
 
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 
 GIT_ACCESS_TOKEN_SECRET_NAME = 'mage_git_access_token'
 GIT_SSH_PRIVATE_KEY_SECRET_NAME = 'mage_git_ssh_private_key_b64'
 GIT_SSH_PUBLIC_KEY_SECRET_NAME = 'mage_git_ssh_public_key_b64'
 
 
-class AuthType(str, Enum):
+class AuthType(StrEnum):
     SSH = 'ssh'
     HTTPS = 'https'
     OAUTH = 'oauth'

mage_ai/data_preparation/templates/data_exporters/streaming/elasticsearch.yaml

@@ -2,6 +2,9 @@ connector_type: elasticsearch
 host: http://localhost
 index_name: elastic_index
 api_key: api_key_for_elastic_search
+# # When using Data streams this must be set to create as they are Append-only indexes
+# _op_type: None / index, create or update
+# _source: None (Default '_doc') / Specify the source field
 # # Whether to verify SSL certificates to authenticate.
 # verify_certs: true
 # ca_cert: "certificate.cert"

mage_ai/data_preparation/templates/data_exporters/streaming/generic_python.py

@@ -1,5 +1,5 @@
 from mage_ai.streaming.sinks.base_python import BasePythonSink
-from typing import Callable
+from typing import Callable, Dict, List
 
 if 'streaming_sink' not in globals():
     from mage_ai.data_preparation.decorators import streaming_sink

mage_ai/data_preparation/templates/data_loaders/streaming/nats.yaml

@@ -17,12 +17,18 @@ timeout: 5
 
 # Optional: TLS configuration
 use_tls: false
-# cafile: "path/to/ca.pem"
-# certfile: "path/to/client-cert.pem"
-# keyfile: "path/to/client-key.pem"
+
+# Optional: SSL configuration
+# ssl_config:
+#   cafile: "path/to/ca.pem"
+#   certfile: "path/to/client-cert.pem"
+#   keyfile: "path/to/client-key.pem"
 
 # Optional: NKEY seed string
 # nkeys_seed_str: "SUAPAEYJQOQHJ4"
 
+# Optional: user credentials
+# user_credentials: "path/to/credentials.creds"
+
 # Optional: Consumer name for a durable consumer identifier, will use stream name if not provided
 consumer_name: "my_consumer"

mage_ai/errors/constants.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import IntEnum
 
 """
 - **400 Bad Request**: The server could not understand the request due to invalid syntax.
@@ -13,7 +13,7 @@ from enum import Enum
 """
 
 
-class ErrorCode(int, Enum):
+class ErrorCode(IntEnum):
     CODE_400 = 400
     CODE_401 = 401
     CODE_402 = 402

mage_ai/io/base.py

@@ -1,12 +1,13 @@
 import os
 from abc import ABC, abstractmethod
-from enum import Enum
 from typing import IO, Any, Callable, Dict, Union
 
 import pandas as pd
+import polars as pl
 from pandas import DataFrame
 
 from mage_ai.io.constants import SQL_RESERVED_WORDS
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.logger import VerbosePrintHandler
 from mage_ai.shared.models import BaseEnum
 from mage_ai.shared.utils import clean_name
@@ -14,7 +15,7 @@ from mage_ai.shared.utils import clean_name
 QUERY_ROW_LIMIT = 10_000_000
 
 
-class DataSource(str, Enum):
+class DataSource(StrEnum):
     ALGOLIA = 'algolia'
     API = 'api'
     BIGQUERY = 'bigquery'
@@ -40,12 +41,32 @@ class DataSource(str, Enum):
     WEAVIATE = 'weaviate'
 
 
-class FileFormat(str, Enum):
+class FileFormat(StrEnum):
     CSV = 'csv'
     JSON = 'json'
     PARQUET = 'parquet'
     HDF5 = 'hdf5'
     XML = 'xml'
+    EXCEL = 'excel'
+
+    @classmethod
+    def from_extension(cls, ext: str):
+        if ext == 'csv':
+            return cls.CSV
+        elif ext == 'json':
+            return cls.JSON
+        elif ext == 'parquet':
+            return cls.PARQUET
+        elif ext == 'hdf5':
+            return cls.HDF5
+        elif ext == 'xml':
+            return cls.XML
+        elif ext in ('xls', 'xlsx'):
+            return cls.EXCEL
+        else:
+            raise ValueError(
+                f'None file format found for this file extension: {ext}'
+            )
 
 
 class ExportWritePolicy(BaseEnum):
@@ -125,7 +146,9 @@ class BaseFile(BaseIO):
         pass
 
     def _get_file_format(self, filepath: Union[os.PathLike, str]) -> str:
-        return os.path.splitext(os.path.basename(filepath))[-1][1:]
+        return FileFormat.from_extension(
+            os.path.splitext(os.path.basename(filepath))[-1][1:]
+        )
 
     def __get_reader(self, format: Union[FileFormat, str, None]) -> Callable:
         """
@@ -150,6 +173,8 @@ class BaseFile(BaseIO):
             return pd.read_hdf
         elif format == FileFormat.XML:
             return pd.read_xml
+        elif format == FileFormat.EXCEL:
+            return pd.read_excel
         else:
             raise ValueError(f"Invalid format '{format}' specified.")
 
@@ -190,7 +215,7 @@ class BaseFile(BaseIO):
 
     def _write(
         self,
-        df: DataFrame,
+        df,
         format: Union[FileFormat, str, None],
         output: Union[IO, os.PathLike, str],
         **kwargs,
@@ -218,7 +243,7 @@ class BaseFile(BaseIO):
                 raise ValueError('Cannot write HDF5 file to buffer of any type.')
             name = os.path.splitext(os.path.basename(output))[0]
             kwargs.setdefault('key', name)
-        elif format == FileFormat.PARQUET:
+        elif format == FileFormat.PARQUET and isinstance(df, DataFrame):
             if 'coerce_timestamps' not in kwargs:
                 kwargs['coerce_timestamps'] = 'ms'
                 kwargs['allow_truncated_timestamps'] = True
@@ -226,7 +251,7 @@ class BaseFile(BaseIO):
 
     def __get_writer(
         self,
-        df: DataFrame,
+        df,
         format: Union[FileFormat, str, None],
     ) -> Callable:
         """
@@ -239,15 +264,31 @@ class BaseFile(BaseIO):
         Returns:
             Callable: File writer method
         """
-        if format == FileFormat.CSV:
-            return df.to_csv
-        elif format == FileFormat.JSON:
-            return df.to_json
-        elif format == FileFormat.HDF5:
-            return df.to_hdf
-        elif format == FileFormat.XML:
-            return df.to_xml
-        return df.to_parquet
+        if isinstance(df, pl.DataFrame):  # polars DataFrame
+            if format == FileFormat.CSV:
+                return df.write_csv
+            elif format == FileFormat.JSON:
+                return df.write_json
+            elif format == FileFormat.HDF5:
+                return df.write_hdf5
+            elif format == FileFormat.XML:
+                return df.write_xml
+            elif format == FileFormat.EXCEL:
+                return df.write_excel
+            return df.write_parquet
+
+        elif isinstance(df, DataFrame):  # pandas DataFrame
+            if format == FileFormat.CSV:
+                return df.to_csv
+            elif format == FileFormat.JSON:
+                return df.to_json
+            elif format == FileFormat.HDF5:
+                return df.to_hdf
+            elif format == FileFormat.XML:
+                return df.to_xml
+            elif format == FileFormat.EXCEL:
+                return df.to_excel
+            return df.to_parquet
 
     def __del__(self):
         if self.verbose and self.printer.exists_previous_message:

mage_ai/io/config.py

@@ -1,6 +1,5 @@
 import os
 from abc import ABC, abstractmethod
-from enum import Enum
 from pathlib import Path
 from typing import Any, Dict, Union
 
@@ -9,9 +8,10 @@ from jinja2 import Template
 
 from mage_ai.data_preparation.shared.utils import get_template_vars
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.enum import StrEnum
 
 
-class ConfigKey(str, Enum):
+class ConfigKey(StrEnum):
     """
     List of configuration settings for use with data IO clients.
     """
@@ -333,7 +333,7 @@ class EnvironmentVariableLoader(BaseConfigLoader):
         return os.getenv(env_var)
 
 
-class VerboseConfigKey(str, Enum):
+class VerboseConfigKey(StrEnum):
     """
     Config key headers for the verbose configuration file format.
     """

mage_ai/io/export_utils.py

@@ -1,9 +1,9 @@
-from enum import Enum
 from typing import Callable, Dict, List, Mapping
 
 from pandas import DataFrame, Series
 from pandas.api.types import infer_dtype
 
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.utils import clean_name
 
 """
@@ -19,7 +19,7 @@ class BadConversionError(Exception):
     pass
 
 
-class PandasTypes(str, Enum):
+class PandasTypes(StrEnum):
     """
     Internal datatypes defined by the pandas Public API
     """

mage_ai/io/google_cloud_storage.py

@@ -1,6 +1,7 @@
 from io import BytesIO
 from typing import Union
 
+import polars as pl
 from google.cloud import storage
 from google.oauth2 import service_account
 from pandas import DataFrame
@@ -103,7 +104,7 @@ class GoogleCloudStorage(BaseFile):
 
     def export(
         self,
-        data: Union[DataFrame, str],
+        data: Union[DataFrame, pl.DataFrame, str],
         bucket_name: str,
         object_key: str,
         format: Union[FileFormat, str, None] = None,
@@ -127,7 +128,7 @@ class GoogleCloudStorage(BaseFile):
         ):
             bucket = self.client.get_bucket(bucket_name)
             blob = bucket.blob(object_key)
-            if isinstance(data, DataFrame):
+            if isinstance(data, DataFrame) or isinstance(data, pl.DataFrame):
                 buffer = BytesIO()
                 self._write(data, format, buffer, **kwargs)
                 buffer.seek(0)

mage_ai/io/io_config.py

@@ -1,12 +1,13 @@
 import os
-from enum import Enum
 from pathlib import Path
 from typing import Any, Mapping, Optional, Union
 
 import yaml
 
+from mage_ai.shared.enum import StrEnum
 
-class IOConfigKeys(str, Enum):
+
+class IOConfigKeys(StrEnum):
     AWS = 'AWS'
     BIGQUERY = 'BigQuery'
     CLICKHOUSE = 'ClickHouse'

mage_ai/io/oracledb.py

@@ -244,10 +244,8 @@ FETCH FIRST {limit} ROWS ONLY
             # Remove extraneous surrounding double quotes
             # that get added while performing conversion to string.
             df_[col] = df_[col].apply(lambda x: x.strip('"') if x and isinstance(x, str) else x)
-        df_.replace({np.NaN: None}, inplace=True)
-        values = []
-        for i in range(0, len(df_)):
-            values.append(tuple(df_.fillna('').values[i]))
+        df_.fillna('', inplace=True)
+        values = list(df_.itertuples(index=False, name=None))
 
         # Create values placeholder
         colmn_names = df.columns.tolist()

mage_ai/io/postgres.py

@@ -99,6 +99,8 @@ class Postgres(BaseSQL):
             password = self.settings['password']
             port = self.settings['port']
             user = self.settings['user']
+            keepalives = self.settings.get('keepalives', 1)
+            keepalives_idle = self.settings.get('keepalives_idle', 300)
             if self.settings['connection_method'] == 'ssh_tunnel':
                 ssh_setting = dict(ssh_username=self.settings['ssh_username'])
                 if self.settings['ssh_pkey'] is not None:
@@ -128,19 +130,35 @@ class Postgres(BaseSQL):
                 host = '127.0.0.1'
                 port = self.ssh_tunnel.local_bind_port
 
-            connect_opts = dict(
-                database=database,
-                host=host,
-                password=password,
-                port=port,
-                user=user,
-                keepalives=1,
-                keepalives_idle=300,
+            connect_opts = self.settings.copy()
+            # See recognized keyword parameters for psycopg2.connect()
+            # https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PARAMKEYWORDS
+            unrecognized_keys = [
+                'dbname',
+                'schema',
+                'connection_method',
+                'ssh_host',
+                'ssh_port',
+                'ssh_username',
+                'ssh_password',
+                'ssh_pkey',
+                "verbose"
+            ]
+            for key in unrecognized_keys:
+                connect_opts.pop(key, None)
+
+            connect_opts.update(
+                {
+                    'database': database,
+                    'host': host,
+                    'password': password,
+                    'port': port,
+                    'user': user,
+                    'keepalives': keepalives,
+                    'keepalives_idle': keepalives_idle,
+                }
             )
 
-            if self.settings.get('connect_timeout'):
-                connect_opts['connect_timeout'] = self.settings['connect_timeout']
-
             try:
                 self._ctx = connect(**connect_opts)
             except Exception:
@@ -166,13 +184,14 @@ class Postgres(BaseSQL):
         self,
         schema_name: str
     ) -> str:
+        clean_schema_name = schema_name.strip('"')
         return f"""
         DO $$
         BEGIN
             IF NOT EXISTS (
                 SELECT schema_name
                 FROM information_schema.schemata
-                WHERE schema_name = '{schema_name}'
+                WHERE schema_name = '{clean_schema_name}'
             ) THEN
                 EXECUTE 'CREATE SCHEMA {schema_name}';
             END IF;
@@ -276,6 +295,14 @@ class Postgres(BaseSQL):
 
         return 'text'
 
+    @staticmethod
+    def _clean_array_value(val: str) -> str:
+        if val is None or type(val) is not str or len(val) < 2:
+            return val
+        if val[0] == '[' and val[-1] == ']':
+            return str(val).replace('[', '{').replace(']', '}')
+        return val
+
     def upload_dataframe(
         self,
         cursor: _psycopg.cursor,
@@ -297,13 +324,6 @@ class Postgres(BaseSQL):
             # Use COPY command
             use_insert_command = False
 
-        def clean_array_value(val):
-            if val is None or type(val) is not str or len(val) < 2:
-                return val
-            if val[0] == '[' and val[-1] == ']':
-                return '{' + val[1:-1] + '}'
-            return val
-
         def serialize_obj(val):
             if type(val) is dict or type(val) is np.ndarray:
                 return simplejson.dumps(
@@ -318,7 +338,7 @@ class Postgres(BaseSQL):
                     ignore_nan=True,
                 )
             elif not use_insert_command and type(val) is list:
-                return clean_array_value(simplejson.dumps(
+                return self._clean_array_value(simplejson.dumps(
                     val,
                     default=encode_complex,
                     ignore_nan=True,

mage_ai/io/utils.py

@@ -29,3 +29,21 @@ def escape_quotes(line: str, single: bool = True, double: bool = True) -> str:
     if double:
         new_line = new_line.replace('\"', '\\"')
     return new_line
+
+
+def map_json_to_airtable(data_types):
+    # Extract the non-null type (ignoring 'null')
+    data_type = next((t for t in data_types if t != 'null'), 'string')
+
+    # Mapping from JSON types to Airtable types
+    type_mapping = {
+        'string': 'multilineText',
+        'integer': 'number',
+        'boolean': 'checkbox',
+        'array': 'multipleSelects',
+        'object': 'singleCollaborator',
+        'number': 'number',
+        'date-time': 'dateTime'
+    }
+
+    return type_mapping.get(data_type, 'str')

mage_ai/kernels/magic/constants.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 
-class ExecutionStatus(str, Enum):
+class ExecutionStatus(StrEnum):
     CANCELLED = 'cancelled'
     ERROR = 'error'
     FAILURE = 'failure'
@@ -10,12 +10,12 @@ class ExecutionStatus(str, Enum):
     SUCCESS = 'success'
 
 
-class EventStreamType(str, Enum):
+class EventStreamType(StrEnum):
     EXECUTION = 'execution'
     TASK = 'task'
 
 
-class ResultType(str, Enum):
+class ResultType(StrEnum):
     DATA = 'data'
     STATUS = 'status'
     STDOUT = 'stdout'

mage_ai/orchestration/constants.py

@@ -1,4 +1,4 @@
-import enum
+from mage_ai.shared.enum import StrEnum
 
 DATABASE_CONNECTION_URL_ENV_VAR = 'MAGE_DATABASE_CONNECTION_URL'
 
@@ -18,7 +18,7 @@ PG_DB_NAME = 'DB_NAME'
 PIPELINE_RUN_MAGE_VARIABLES_KEY = '__mage_variables'
 
 
-class Entity(str, enum.Enum):
+class Entity(StrEnum):
     # Permissions saved to the DB should not have the "ANY" entity. It should only be used
     # when evaluating permissions.
     ANY = 'any'

mage_ai/orchestration/db/__init__.py

@@ -13,7 +13,7 @@ from mage_ai.orchestration.db.setup import get_postgres_connection_url
 from mage_ai.orchestration.db.utils import get_user_info_from_db_connection_url
 from mage_ai.settings import OTEL_EXPORTER_OTLP_ENDPOINT
 from mage_ai.settings.repo import get_variables_dir
-from mage_ai.shared.environments import is_debug, is_test
+from mage_ai.shared.environments import is_debug, is_test_mage
 
 DB_RETRY_COUNT = 2
 TEST_DB = 'test.db'
@@ -28,7 +28,7 @@ db_kwargs = dict(
 if OTEL_EXPORTER_OTLP_ENDPOINT:
     from opentelemetry.instrumentation.sqlalchemy import SQLAlchemyInstrumentor
 
-if is_test():
+if is_test_mage():
     db_connection_url = f'sqlite:///{TEST_DB}'
     db_kwargs['connect_args']['check_same_thread'] = False
 elif not db_connection_url:
@@ -37,7 +37,7 @@ elif not db_connection_url:
     if pg_db_connection_url:
         db_connection_url = pg_db_connection_url
     else:
-        if is_test():
+        if is_test_mage():
             db_connection_url = f'sqlite:///{TEST_DB}'
         elif os.path.exists(os.path.join('mage_ai', 'orchestration', 'db')):
             # For local dev environment

mage_ai/orchestration/db/constants.py

@@ -1,6 +1,6 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 
-class DatabaseType(str, Enum):
+class DatabaseType(StrEnum):
     POSTGRESQL = 'postgresql'
     SQLITE = 'sqlite'

mage_ai/orchestration/db/models/oauth.py

@@ -1,4 +1,3 @@
-import enum
 import re
 from datetime import datetime
 from typing import Callable, List, Union
@@ -31,6 +30,7 @@ from mage_ai.orchestration.db import db_connection, safe_db_query
 from mage_ai.orchestration.db.errors import ValidationError
 from mage_ai.orchestration.db.models.base import BaseModel
 from mage_ai.shared.array import find
+from mage_ai.shared.enum import IntEnum, StrEnum
 from mage_ai.shared.environments import is_test
 from mage_ai.shared.hash import group_by, merge_dict
 
@@ -299,7 +299,7 @@ class Role(BaseModel):
     user = relationship(User, back_populates='created_roles')
 
     # Default global roles created by Mage
-    class DefaultRole(str, enum.Enum):
+    class DefaultRole(StrEnum):
         OWNER = 'Owner'
         ADMIN = 'Admin'
         EDITOR = 'Editor'
@@ -561,7 +561,7 @@ class UserRole(BaseModel):
 
 
 class Permission(BaseModel):
-    class Access(int, enum.Enum):
+    class Access(IntEnum):
         OWNER = PermissionAccess.OWNER.value
         ADMIN = PermissionAccess.ADMIN.value
         EDITOR = PermissionAccess.EDITOR.value
@@ -841,11 +841,11 @@ class RolePermission(BaseModel):
 
 
 class Oauth2Application(BaseModel):
-    class AuthorizationGrantType(str, enum.Enum):
+    class AuthorizationGrantType(StrEnum):
         AUTHORIZATION_CODE = 'authorization-code'
         CLIENT_CREDENTIALS = 'client-credentials'
 
-    class ClientType(str, enum.Enum):
+    class ClientType(StrEnum):
         PRIVATE = 'private'
         PUBLIC = 'public'