mage-ai 0.9.74__py3-none-any.whl → 0.9.76__py3-none-any.whl

mage_ai/io/airtable.py

@@ -0,0 +1,104 @@
+from typing import Union
+
+import pandas as pd
+import polars as pl
+from pyairtable.api import Api
+
+from mage_ai.io.base import BaseIO
+from mage_ai.io.config import BaseConfigLoader, ConfigKey
+
+
+class Airtable(BaseIO):
+    """
+    Handles data transfer between Airtable tables and the Mage app.
+    """
+
+    def __init__(
+            self,
+            token: str,
+            verbose: bool = True,
+            **kwargs) -> None:
+        """
+        Initializes a connection to Airtable.
+
+        Args:
+            token (str): Airtable API token.
+            verbose (bool, optional): Whether to print verbose output. Defaults to True.
+            **kwargs: Additional keyword arguments to pass to the pyairtable Api constructor.
+        """
+        super().__init__(verbose=verbose)
+        self.client = Api(token, **kwargs)  # Create the Airtable API client
+
+    @classmethod
+    def with_config(
+            cls,
+            config: BaseConfigLoader
+    ) -> 'Airtable':
+        """
+        Initializes an Airtable client from a configuration loader.
+
+        Args:
+            config (BaseConfigLoader): Configuration loader object.
+
+        Raises:
+            ValueError: If no valid Airtable API token is found in the configuration.
+
+        Returns:
+            Airtable: An instance of the Airtable class.
+        """
+        if ConfigKey.AIRTABLE_ACCESS_TOKEN not in config:
+            raise ValueError(
+                'No valid API token found for Airtable.'
+                'You must specify your access token in your config.'
+            )
+
+        return cls(
+            token=config[ConfigKey.AIRTABLE_ACCESS_TOKEN]
+        )
+
+    def load(
+            self,
+            base_id: str,
+            table_name: str,
+            **kwargs,
+    ) -> pd.DataFrame:
+        """
+        Loads data from an Airtable table into a Pandas DataFrame.
+
+        Args:
+            base_id (str): The ID of the Airtable base (e.g., 'app*****').
+            table_name (str): The name or ID of the Airtable table (e.g., 'tbl*****').
+            **kwargs: Additional keyword arguments to pass to the pyairtable table.all() method.
+
+        Returns:
+            DataFrame: A Pandas DataFrame containing the data from the Airtable table.
+        """
+        with self.printer.print_msg(
+                f'Loading data frame from table \'{table_name}\' at airtable app \'{base_id}\''
+        ):
+            table = self.client.table(base_id, table_name)  # Get the Airtable table
+            data = table.all(**kwargs)  # Fetch all records from the table
+
+            # Flatten the Airtable data structure into a list of dictionaries
+            flattened_data = []
+            for record in data:
+                flattened_record = {
+                    'id': record['id'],
+                    'createdTime': record['createdTime']
+                }
+                fields = record['fields']
+                flattened_record.update(fields)
+                flattened_data.append(flattened_record)
+
+            return pd.DataFrame(flattened_data)  # Create and return a DataFrame
+
+    def export(
+            self,
+            df: Union[pd.DataFrame, pl.DataFrame],
+            *args,
+            **kwargs
+    ) -> None:
+        """
+        Not implemented yet. This method is intended to export data to Airtable.
+        """
+        pass

mage_ai/io/base.py

@@ -16,6 +16,7 @@ QUERY_ROW_LIMIT = 10_000_000
 
 
 class DataSource(StrEnum):
+    AIRTABLE = 'airtable'
     ALGOLIA = 'algolia'
     API = 'api'
     BIGQUERY = 'bigquery'
@@ -47,6 +48,26 @@ class FileFormat(StrEnum):
     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):
@@ -126,7 +147,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:
         """
@@ -151,6 +174,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.")
 
@@ -249,6 +274,8 @@ class BaseFile(BaseIO):
                 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
@@ -260,6 +287,8 @@ class BaseFile(BaseIO):
                 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):

mage_ai/io/bigquery.py

@@ -1,7 +1,9 @@
+import traceback
 import uuid
 from typing import Dict, List, Mapping, Union
 
 import numpy as np
+import pandas as pd
 from google.cloud.bigquery import (
     Client,
     LoadJobConfig,
@@ -399,7 +401,41 @@ WHERE table_id = '{table_name}'
             # Clean column names
             if type(df) is DataFrame:
                 df.columns = df.columns.str.replace(' ', '_')
+            table = None
+            try:
+                # Cast column types
+                table = self.client.get_table(table_id)
+            except Exception:
+                print(f'Table {table_id} does not exist.')
+                pass
 
+            if table is not None:
+                try:
+                    timestamp_columns = [field.name for field in table.schema
+                                         if field.field_type == 'TIMESTAMP']
+
+                    # Convert TIMESTAMP columns in DataFrame
+                    for col in timestamp_columns:
+                        if col in df.columns:
+                            df[col] = pd.to_datetime(df[col])
+                except Exception:
+                    print('Fail to cast column types in dataframe.')
+                    traceback.print_exc()
+                if (
+                    not config.schema and
+                    config.write_disposition == WriteDisposition.WRITE_TRUNCATE
+                ):
+                    df_columns = df.columns.tolist()
+                    config.schema = [
+                        SchemaField(
+                            field.name,
+                            field.field_type,
+                            mode=field.mode,
+                            fields=field.fields,
+                        )
+                        for field in table.schema
+                        if field.name in df_columns
+                    ]
             return self.client.load_table_from_dataframe(df, table_id, job_config=config).result()
 
     def execute(self, query_string: str, **kwargs) -> None:

mage_ai/io/config.py

@@ -16,6 +16,8 @@ class ConfigKey(StrEnum):
     List of configuration settings for use with data IO clients.
     """
 
+    AIRTABLE_ACCESS_TOKEN = "AIRTABLE_ACCESS_TOKEN"
+
     ALGOLIA_APP_ID = 'ALGOLIA_APP_ID'
     ALGOLIA_API_KEY = 'ALGOLIA_API_KEY'
     ALGOLIA_INDEX_NAME = 'ALGOLIA_INDEX_NAME'
@@ -78,6 +80,7 @@ class ConfigKey(StrEnum):
     MYSQL_PASSWORD = 'MYSQL_PASSWORD'
     MYSQL_PORT = 'MYSQL_PORT'
     MYSQL_USER = 'MYSQL_USER'
+    MYSQL_ALLOW_LOCAL_INFILE = 'MYSQL_ALLOW_LOCAL_INFILE'
 
     ORACLEDB_USER = 'ORACLEDB_USER'
     ORACLEDB_PASSWORD = 'ORACLEDB_PASSWORD'
@@ -338,6 +341,7 @@ class VerboseConfigKey(StrEnum):
     Config key headers for the verbose configuration file format.
     """
 
+    AIRTABLE = 'Airtable'
     ALGOLIA = 'Algolia'
     AWS = 'AWS'
     BIGQUERY = 'BigQuery'
@@ -356,6 +360,7 @@ class VerboseConfigKey(StrEnum):
 
 class ConfigFileLoader(BaseConfigLoader):
     KEY_MAP = {
+        ConfigKey.AIRTABLE_ACCESS_TOKEN: VerboseConfigKey.AIRTABLE,
         ConfigKey.ALGOLIA_APP_ID: (
             VerboseConfigKey.ALGOLIA, 'app_id'),
         ConfigKey.ALGOLIA_API_KEY: (

mage_ai/io/mssql.py

@@ -14,6 +14,7 @@ from mage_ai.io.config import BaseConfigLoader, ConfigKey
 from mage_ai.io.constants import UNIQUE_CONFLICT_METHOD_UPDATE
 from mage_ai.io.export_utils import PandasTypes
 from mage_ai.io.sql import BaseSQL
+from mage_ai.shared.hash import extract
 from mage_ai.shared.parsers import encode_complex
 
 MERGE_TABLE_SQL = '''MERGE {table_name} AS t
@@ -197,9 +198,11 @@ class MSSQL(BaseSQL):
                 TrustServerCertificate='yes',
             ),
         )
+        conn_kwargs = extract(kwargs, ['pool_size', 'max_overflow'])
         engine = create_engine(
             connection_url,
             fast_executemany=True,
+            **conn_kwargs,
         )
 
         unique_conflict_method = kwargs.get('unique_conflict_method')
@@ -242,6 +245,7 @@ class MSSQL(BaseSQL):
                     method=merge_table,
                 )
                 return
+        sql_kwargs = extract(kwargs, ['chunksize', 'method'])
 
         df.to_sql(
             table_name,
@@ -249,6 +253,7 @@ class MSSQL(BaseSQL):
             schema=schema_name,
             if_exists=if_exists or ExportWritePolicy.REPLACE,
             index=False,
+            **sql_kwargs,
         )
 
     def get_type(self, column: Series, dtype: str) -> str:

mage_ai/io/mysql.py

@@ -25,6 +25,7 @@ class MySQL(BaseSQL):
         password: str,
         user: str,
         port: int = 3306,
+        allow_local_infile: bool = False,
         verbose: bool = True,
         **kwargs,
     ) -> None:
@@ -35,18 +36,22 @@ class MySQL(BaseSQL):
             port=port or 3306,
             user=user,
             verbose=verbose,
+            allow_local_infile=allow_local_infile,
             **kwargs,
         )
 
     @classmethod
     def with_config(cls, config: BaseConfigLoader) -> 'MySQL':
-        return cls(
+        conn_kwargs = dict(
             database=config[ConfigKey.MYSQL_DATABASE],
             host=config[ConfigKey.MYSQL_HOST],
             password=config[ConfigKey.MYSQL_PASSWORD],
             port=config[ConfigKey.MYSQL_PORT],
             user=config[ConfigKey.MYSQL_USER],
         )
+        if config[ConfigKey.MYSQL_ALLOW_LOCAL_INFILE] is not None:
+            conn_kwargs['allow_local_infile'] = config[ConfigKey.MYSQL_ALLOW_LOCAL_INFILE]
+        return cls(**conn_kwargs)
 
     def build_create_table_command(
         self,

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:
@@ -277,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,
@@ -298,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(
@@ -319,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/redshift.py

@@ -5,6 +5,7 @@ from typing import Dict, Union
 from pandas import DataFrame
 from redshift_connector import connect
 
+from mage_ai.data_preparation.models.block.sql.utils.shared import split_query_string
 from mage_ai.io.base import QUERY_ROW_LIMIT, ExportWritePolicy
 from mage_ai.io.config import BaseConfigLoader, ConfigKey
 from mage_ai.io.export_utils import clean_df_for_export, infer_dtypes
@@ -67,6 +68,18 @@ class Redshift(BaseSQL):
             with self.conn.cursor() as cur:
                 cur.execute(query_string, **kwargs)
 
+    def execute_query_raw(self, query: str, **kwargs) -> None:
+        """
+        Overwrite execute query to process multiple queries in one string.
+        """
+        results = []
+        with self.conn.cursor() as cursor:
+            for query_string in split_query_string(query):
+                result = cursor.execute(query_string)
+                results.append(result)
+        self.conn.commit()
+        return results
+
     def load(
         self,
         query_string: str,

mage_ai/io/sql.py

@@ -301,6 +301,7 @@ class BaseSQL(BaseSQLConnection):
                     if_exists=if_exists,
                     unique_conflict_method=unique_conflict_method,
                     unique_constraints=unique_constraints,
+                    **kwargs,
                 )
                 return
 

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/orchestration/db/__init__.py

@@ -1,5 +1,6 @@
 import logging
 import os
+from functools import wraps
 from urllib.parse import parse_qs, quote_plus, urlparse
 
 import sqlalchemy
@@ -13,7 +14,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 +29,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 +38,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
@@ -169,6 +170,25 @@ def safe_db_query(func):
     return func_with_rollback
 
 
+def safe_db_query_async(func):
+    @wraps(func)
+    async def func_with_rollback_async(*args, **kwargs):
+        retry_count = 0
+        while True:
+            try:
+                return await func(*args, **kwargs)
+            except (
+                sqlalchemy.exc.OperationalError,
+                sqlalchemy.exc.PendingRollbackError,
+                sqlalchemy.exc.InternalError,
+            ) as e:
+                db_connection.session.rollback()
+                if retry_count >= DB_RETRY_COUNT:
+                    raise e
+                retry_count += 1
+    return func_with_rollback_async
+
+
 logging.basicConfig()
 
 if is_debug() and not os.getenv('DISABLE_DATABASE_TERMINAL_OUTPUT'):

mage_ai/orchestration/db/models/schedules.py

@@ -1,4 +1,3 @@
-import asyncio
 import collections
 import traceback
 import uuid
@@ -1400,10 +1399,8 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
 
         from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
-        asyncio.run(
-            UsageStatisticLogger().pipeline_runs_impression(
-                lambda: self.query.filter(self.status == self.PipelineRunStatus.COMPLETED).count(),
-            )
+        UsageStatisticLogger().pipeline_runs_impression_sync(
+            lambda: self.query.filter(self.status == self.PipelineRunStatus.COMPLETED).count(),
         )
 
     @safe_db_query

mage_ai/orchestration/db/models/secrets.py

@@ -3,7 +3,8 @@ from typing import Optional
 from sqlalchemy import Column, String, Text, UniqueConstraint, or_
 
 from mage_ai.orchestration.db import safe_db_query
-from mage_ai.orchestration.db.models.base import BaseModel
+from mage_ai.orchestration.db.models.base import BaseModel, classproperty
+from mage_ai.settings.repo import get_repo_path
 
 
 class Secret(BaseModel):
@@ -15,6 +16,15 @@ class Secret(BaseModel):
     key_uuid = Column(String(255), nullable=True)
     __table_args__ = (UniqueConstraint('name', 'key_uuid', name='name_key_uuid_uc'),)
 
+    @classproperty
+    def repo_query(cls):
+        return cls.query.filter(
+            or_(
+                Secret.repo_name == get_repo_path(),
+                Secret.repo_name.is_(None),
+            )
+        )
+
     @classmethod
     @safe_db_query
     def get_secret(cls, name: str, key_uuid: str) -> Optional['Secret']:

mage_ai/orchestration/metrics/pipeline_run.py

@@ -220,7 +220,7 @@ def __calculate_metrics(pipeline_run: PipelineRun, pipeline) -> Dict:
 
     pipeline_metrics_by_stream = {}
     pipeline_logs_by_stream = {}
-    pipeline_logs = pipeline_run.logs['content'].split('\n')
+    pipeline_logs = pipeline_run.logs[0]['content'].split('\n')
     for pipeline_log in pipeline_logs:
         tags = parse_line(pipeline_log)
         stream = tags.get('stream')

mage_ai/orchestration/pipeline_scheduler_original.py

@@ -1,4 +1,3 @@
-import asyncio
 import collections
 import os
 import traceback
@@ -260,7 +259,7 @@ class PipelineScheduler:
                         pipeline=self.pipeline,
                         pipeline_run=self.pipeline_run,
                     )
-                    asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
+                    UsageStatisticLogger().pipeline_run_ended_sync(self.pipeline_run)
 
                 self.logger_manager.output_logs_to_destination()
 
@@ -339,7 +338,7 @@ class PipelineScheduler:
         error_msg: str,
         status=PipelineRun.PipelineRunStatus.FAILED,
     ) -> None:
-        asyncio.run(UsageStatisticLogger().pipeline_run_ended(self.pipeline_run))
+        UsageStatisticLogger().pipeline_run_ended_sync(self.pipeline_run)
 
         if status == PipelineRun.PipelineRunStatus.FAILED:
             # Only send notification when pipeline run status is FAILED
@@ -1372,7 +1371,7 @@ def stop_pipeline_run(
     # Update pipeline run status to cancelled
     pipeline_run.update(status=status)
 
-    asyncio.run(UsageStatisticLogger().pipeline_run_ended(pipeline_run))
+    UsageStatisticLogger().pipeline_run_ended_sync(pipeline_run)
 
     # Cancel all the block runs
     cancel_block_runs_and_jobs(pipeline_run, pipeline)
@@ -1584,7 +1583,7 @@ def schedule_all():
         pipeline_runs_excluded_by_limit = []
         for pipeline_schedule in active_schedules:
             lock_key = f'pipeline_schedule_{pipeline_schedule.id}'
-            if not lock.try_acquire_lock(lock_key):
+            if not lock.try_acquire_lock(lock_key, timeout=30):
                 continue
 
             trigger_pipeline_run_limit = pipeline_schedule.get_settings().pipeline_run_limit