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

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/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

@@ -24,5 +24,8 @@ use_tls: false
 # 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,7 +41,7 @@ class DataSource(str, Enum):
     WEAVIATE = 'weaviate'
 
 
-class FileFormat(str, Enum):
+class FileFormat(StrEnum):
     CSV = 'csv'
     JSON = 'json'
     PARQUET = 'parquet'
@@ -190,7 +191,7 @@ class BaseFile(BaseIO):
 
     def _write(
         self,
-        df: DataFrame,
+        df,
         format: Union[FileFormat, str, None],
         output: Union[IO, os.PathLike, str],
         **kwargs,
@@ -218,7 +219,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 +227,7 @@ class BaseFile(BaseIO):
 
     def __get_writer(
         self,
-        df: DataFrame,
+        df,
         format: Union[FileFormat, str, None],
     ) -> Callable:
         """
@@ -239,15 +240,27 @@ 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
+            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
+            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/postgres.py

@@ -166,13 +166,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;

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/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'
 

mage_ai/orchestration/db/models/schedules.py

@@ -1,6 +1,5 @@
 import asyncio
 import collections
-import enum
 import traceback
 import uuid
 from datetime import datetime, timedelta, timezone
@@ -73,6 +72,7 @@ from mage_ai.settings.platform import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
 from mage_ai.shared.constants import ENV_PROD
 from mage_ai.shared.dates import compare
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import ignore_keys, index_by, merge_dict
 from mage_ai.shared.utils import clean_name
 
@@ -201,13 +201,32 @@ class PipelineSchedule(PipelineScheduleProjectPlatformMixin, BaseModel):
 
     @property
     def pipeline_runs_count(self) -> int:
-        if project_platform_activated():
-            return self.pipeline_runs_count_project_platform
+        return (
+            PipelineRun.select(func.count(PipelineRun.id))
+            .filter(
+                PipelineRun.pipeline_schedule_id == self.id,
+            )
+            .scalar()
+        )
 
+    @property
+    def initial_pipeline_runs(self) -> List:
+        return (
+            PipelineRun.query
+            .filter(
+                PipelineRun.pipeline_schedule_id == self.id,
+                PipelineRun.status == PipelineRun.PipelineRunStatus.INITIAL,
+            )
+            .all()
+        )
+
+    @property
+    def running_pipeline_run_count(self) -> int:
         return (
             PipelineRun.select(func.count(PipelineRun.id))
             .filter(
                 PipelineRun.pipeline_schedule_id == self.id,
+                PipelineRun.status == PipelineRun.PipelineRunStatus.RUNNING,
             )
             .scalar()
         )
@@ -736,7 +755,7 @@ class PipelineSchedule(PipelineScheduleProjectPlatformMixin, BaseModel):
 
 
 class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
-    class PipelineRunStatus(str, enum.Enum):
+    class PipelineRunStatus(StrEnum):
         INITIAL = 'initial'
         RUNNING = 'running'
         COMPLETED = 'completed'
@@ -1538,6 +1557,7 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
         variables['execution_date'] = self.execution_date
         variables['execution_partition'] = self.execution_partition
         variables['pipeline_run_id'] = self.id
+        variables['trigger_name'] = self.pipeline_schedule.name
 
         interval_end_datetime = variables.get('interval_end_datetime')
         interval_seconds = variables.get('interval_seconds')
@@ -1632,7 +1652,7 @@ class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel):
 
 
 class BlockRun(BlockRunProjectPlatformMixin, BaseModel):
-    class BlockRunStatus(str, enum.Enum):
+    class BlockRunStatus(StrEnum):
         INITIAL = 'initial'
         QUEUED = 'queued'
         RUNNING = 'running'
@@ -1788,7 +1808,7 @@ class BlockRun(BlockRunProjectPlatformMixin, BaseModel):
 
 
 class EventMatcher(BaseModel):
-    class EventType(str, enum.Enum):
+    class EventType(StrEnum):
         AWS_EVENT = 'aws_event'
 
     event_type = Column(Enum(EventType), default=EventType.AWS_EVENT)
@@ -1889,7 +1909,7 @@ class EventMatcher(BaseModel):
 
 
 class Backfill(BaseModel):
-    class IntervalType(str, enum.Enum):
+    class IntervalType(StrEnum):
         SECOND = 'second'
         MINUTE = 'minute'
         HOUR = 'hour'
@@ -1899,7 +1919,7 @@ class Backfill(BaseModel):
         YEAR = 'year'
         CUSTOM = 'custom'
 
-    class Status(str, enum.Enum):
+    class Status(StrEnum):
         INITIAL = 'initial'
         RUNNING = 'running'
         COMPLETED = 'completed'

mage_ai/orchestration/db/models/schedules_project_platform.py

@@ -27,7 +27,6 @@ from mage_ai.settings.platform.utils import (
     get_pipeline_from_platform,
     get_pipeline_from_platform_async,
 )
-from mage_ai.shared.array import find
 from mage_ai.shared.constants import ENV_PROD
 from mage_ai.shared.dates import compare
 from mage_ai.shared.hash import merge_dict
@@ -192,15 +191,18 @@ class PipelineScheduleProjectPlatformMixin:
                     compare(current_execution_date, self.start_time.replace(tzinfo=pytz.UTC)) == -1:
                 return False
 
+            from mage_ai.orchestration.db.models.schedules import PipelineRun
+
             # If there is a pipeline_run with an execution_date the same as the
             # current_execution_date, then don’t schedule
-            if not find(
-                lambda x: compare(
-                    x.execution_date.replace(tzinfo=pytz.UTC),
-                    current_execution_date,
-                ) == 0,
-                self.fetch_pipeline_runs([self.id])
-            ):
+            run_exists = PipelineRun.select(
+                PipelineRun.query.filter(
+                    PipelineRun.pipeline_schedule_id == self.id,
+                    PipelineRun.execution_date == current_execution_date,
+                ).exists()
+            ).scalar()
+
+            if not run_exists:
                 if self.landing_time_enabled():
                     if not previous_runtimes or len(previous_runtimes) == 0:
                         return True

mage_ai/orchestration/job_manager.py

@@ -1,10 +1,10 @@
-from enum import Enum
 from typing import Callable, Dict, Union
 
 from mage_ai.orchestration.queue.queue_factory import QueueFactory
+from mage_ai.shared.enum import StrEnum
 
 
-class JobType(str, Enum):
+class JobType(StrEnum):
     BLOCK_RUN = 'block_run'
     PIPELINE_RUN = 'pipeline_run'
     INTEGRATION_STREAM = 'integration_stream'

mage_ai/orchestration/monitor/monitor_stats.py

@@ -1,4 +1,3 @@
-import enum
 from datetime import datetime, timedelta
 from functools import reduce
 from typing import Callable, Dict, List, Union
@@ -15,13 +14,14 @@ from mage_ai.orchestration.db.models.schedules import (
 )
 from mage_ai.settings.platform.constants import project_platform_activated
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.hash import group_by, merge_dict
 
 NO_PIPELINE_SCHEDULE_ID = 'no_pipeline_schedule_id'
 NO_PIPELINE_SCHEDULE_NAME = 'no_pipeline_schedule_name'
 
 
-class MonitorStatsType(str, enum.Enum):
+class MonitorStatsType(StrEnum):
     PIPELINE_RUN_COUNT = 'pipeline_run_count'
     PIPELINE_RUN_TIME = 'pipeline_run_time'
     BLOCK_RUN_COUNT = 'block_run_count'

mage_ai/orchestration/notification/config.py

@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import List
 
 from mage_ai.services.discord.config import DiscordConfig
@@ -10,9 +9,10 @@ from mage_ai.services.slack.config import SlackConfig
 from mage_ai.services.teams.config import TeamsConfig
 from mage_ai.services.telegram.config import TelegramConfig
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 
 
-class AlertOn(str, Enum):
+class AlertOn(StrEnum):
     PIPELINE_RUN_FAILURE = 'trigger_failure'
     PIPELINE_RUN_SUCCESS = 'trigger_success'
     PIPELINE_RUN_PASSED_SLA = 'trigger_passed_sla'

mage_ai/orchestration/pipeline_scheduler_original.py

@@ -1610,21 +1610,13 @@ def schedule_all():
                     pipeline=pipeline,
                     repo_config=repo_config,
                 )
-                initial_pipeline_runs = [
-                    r
-                    for r in pipeline_schedule.pipeline_runs
-                    if r.status == PipelineRun.PipelineRunStatus.INITIAL
-                ]
+                initial_pipeline_runs = pipeline_schedule.initial_pipeline_runs
 
                 if not should_schedule and not initial_pipeline_runs:
                     lock.release_lock(lock_key)
                     continue
 
-                running_pipeline_runs = [
-                    r
-                    for r in pipeline_schedule.pipeline_runs
-                    if r.status == PipelineRun.PipelineRunStatus.RUNNING
-                ]
+                running_pipeline_run_count = pipeline_schedule.running_pipeline_run_count
 
                 if (
                     should_schedule
@@ -1646,7 +1638,7 @@ def schedule_all():
                         payload['metrics'] = dict(previous_runtimes=previous_runtimes)
 
                     if pipeline_schedule.get_settings().skip_if_previous_running and (
-                        initial_pipeline_runs or running_pipeline_runs
+                        initial_pipeline_runs or running_pipeline_run_count > 0
                     ):
                         # Cancel the pipeline run if previous pipeline runs haven't completed and
                         # skip_if_previous_running is enabled
@@ -1677,7 +1669,7 @@ def schedule_all():
                 # Enforce pipeline concurrency limit
                 pipeline_run_quota = None
                 if trigger_pipeline_run_limit is not None:
-                    pipeline_run_quota = trigger_pipeline_run_limit - len(running_pipeline_runs)
+                    pipeline_run_quota = trigger_pipeline_run_limit - running_pipeline_run_count
 
                 if pipeline_run_quota is None:
                     pipeline_run_quota = len(initial_pipeline_runs)

mage_ai/orchestration/queue/config.py

@@ -1,10 +1,10 @@
 from dataclasses import dataclass
-from enum import Enum
 
 from mage_ai.shared.config import BaseConfig
+from mage_ai.shared.enum import StrEnum
 
 
-class QueueType(str, Enum):
+class QueueType(StrEnum):
     CELERY = 'celery'
     PROCESS = 'process'
 

mage_ai/orchestration/queue/process_queue.py

@@ -2,7 +2,6 @@ import multiprocessing as mp
 import os
 import signal
 import time
-from enum import Enum
 from multiprocessing import Manager
 from typing import Callable, Dict
 
@@ -24,19 +23,20 @@ from mage_ai.settings import (
     SERVER_LOGGING_FORMAT,
     SERVER_VERBOSITY,
 )
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.logger import set_logging_format
 
 LIVENESS_TIMEOUT_SECONDS = 300
 
 
-class JobStatus(str, Enum):
+class JobStatus(StrEnum):
     QUEUED = 'queued'
     RUNNING = 'running'  # Not used. The value for RUNNING job is process id.
     COMPLETED = 'completed'
     CANCELLED = 'cancelled'
 
 
-class QueueStatus(str, Enum):
+class QueueStatus(StrEnum):
     ACTIVE = 'active'
     INACTIVE = 'inactive'
 

mage_ai/presenters/charts/data_sources/constants.py

@@ -1,9 +1,9 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 DEFAULT_LIMIT = 10000
 
 
-class ChartDataSourceType(str, Enum):
+class ChartDataSourceType(StrEnum):
     BLOCK = 'block'
     BLOCK_RUNS = 'block_runs'
     CHART_CODE = 'chart_code'