diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 1d3124e0aee..d62231cb8fb 100644 --- a/.github/workflows/build_and_test.yml +++ b/.github/workflows/build_and_test.yml @@ -50,7 +50,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.8", "3.9", "3.10"] + python-version: ["3.8", "3.9", "3.10", "3.11"] needs: check-code-quality steps: - uses: actions/checkout@v3 diff --git a/Dockerfile b/Dockerfile index 8bdcadbf01b..b9fdbc8ca19 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,4 +1,4 @@ -FROM python:3.10-bookworm +FROM python:3.11-bookworm LABEL description="Deploy Mage on ECS" ARG FEATURE_BRANCH USER root diff --git a/docs/contributing/charts/how-to-add.mdx b/docs/contributing/charts/how-to-add.mdx index 532bc26e3e9..e995890145e 100644 --- a/docs/contributing/charts/how-to-add.mdx +++ b/docs/contributing/charts/how-to-add.mdx @@ -8,7 +8,7 @@ sidebarTitle: "Charts" Add a new type in [`mage_ai/data_preparation/models/widget/constants.py`](https://github.com/mage-ai/mage-ai/blob/master/mage_ai/data_preparation/models/widget/constants.py): ```python -class ChartType(str, Enum): +class ChartType(StrEnum): # ... PIE_CHART = 'pie_chart' ``` diff --git a/mage_ai/ai/constants.py b/mage_ai/ai/constants.py index 5aefb1049e4..e8b2797a775 100644 --- a/mage_ai/ai/constants.py +++ b/mage_ai/ai/constants.py @@ -1,7 +1,7 @@ -from enum import Enum +from enum import StrEnum -class LLMUseCase(str, Enum): +class LLMUseCase(StrEnum): GENERATE_DOC_FOR_BLOCK = 'generate_doc_for_block' GENERATE_DOC_FOR_PIPELINE = 'generate_doc_for_pipeline' GENERATE_BLOCK_WITH_DESCRIPTION = 'generate_block_with_description' diff --git a/mage_ai/api/constants.py b/mage_ai/api/constants.py index 367e0fd0265..6fd637b729d 100644 --- a/mage_ai/api/constants.py +++ b/mage_ai/api/constants.py @@ -1,20 +1,20 @@ -from enum import Enum +from enum import StrEnum from mage_ai.api.operations.constants import OperationType from mage_ai.orchestration.db.models.oauth import Permission -class AttributeOperationType(str, Enum): +class AttributeOperationType(StrEnum): QUERY = 'query' READ = 'read' WRITE = 'write' -class AttributeType(str, Enum): +class AttributeType(StrEnum): ALL = '__*MAGE*__' -class AuthorizeStatusType(str, Enum): +class AuthorizeStatusType(StrEnum): ALL = 'all' FAILED = 'failed' SUCCEEDED = 'succeeded' diff --git a/mage_ai/api/oauth_scope.py b/mage_ai/api/oauth_scope.py index 93f3575d6fe..a8d6f9557a9 100644 --- a/mage_ai/api/oauth_scope.py +++ b/mage_ai/api/oauth_scope.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum class OauthScope(): @@ -18,7 +18,7 @@ class OauthScope(): TOKEN_SCOPES = [] -class OauthScopeType(str, Enum): +class OauthScopeType(StrEnum): CLIENT_ALL = 'all' CLIENT_INTERNAL = 'internal' CLIENT_PRIVATE = 'private' diff --git a/mage_ai/api/operations/constants.py b/mage_ai/api/operations/constants.py index 1d7b547412c..21e475e8c83 100644 --- a/mage_ai/api/operations/constants.py +++ b/mage_ai/api/operations/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum ALL = 'all' CREATE = 'create' @@ -20,7 +20,7 @@ COOKIE_PREFIX = '__COOKIE__' -class OperationType(str, Enum): +class OperationType(StrEnum): ALL = ALL CREATE = CREATE DELETE = DELETE diff --git a/mage_ai/authentication/oauth/constants.py b/mage_ai/authentication/oauth/constants.py index bf2b1cd4b3a..296aeec3e17 100644 --- a/mage_ai/authentication/oauth/constants.py +++ b/mage_ai/authentication/oauth/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum from typing import Optional from mage_ai.settings import get_settings_value @@ -10,7 +10,7 @@ GITHUB_STATE = '1337' -class ProviderName(str, Enum): +class ProviderName(StrEnum): ACTIVE_DIRECTORY = 'active_directory' AZURE_DEVOPS = 'azure_devops' BITBUCKET = 'bitbucket' diff --git a/mage_ai/authentication/operation_history/constants.py b/mage_ai/authentication/operation_history/constants.py index f439027ccfb..e07c5eb18a6 100644 --- a/mage_ai/authentication/operation_history/constants.py +++ b/mage_ai/authentication/operation_history/constants.py @@ -1,8 +1,8 @@ -from enum import Enum +from enum import StrEnum MAGE_OPERATION_HISTORY_DIRECTORY_DEFAULT = '.operation_history' MAGE_OPERATION_HISTORY_DIRECTORY_ENVIRONMENT_VARIABLE_NAME = 'MAGE_OPERATION_HISTORY_DIRECTORY' -class ResourceType(str, Enum): +class ResourceType(StrEnum): PIPELINE = 'pipeline' diff --git a/mage_ai/authentication/permissions/constants.py b/mage_ai/authentication/permissions/constants.py index cf6d064690d..c85032a2be7 100644 --- a/mage_ai/authentication/permissions/constants.py +++ b/mage_ai/authentication/permissions/constants.py @@ -1,10 +1,10 @@ -from enum import Enum +from enum import IntEnum, StrEnum from mage_ai.api.operations.constants import OperationType from mage_ai.data_preparation.models.constants import BlockType, PipelineType -class EntityName(str, Enum): +class EntityName(StrEnum): ALL = 'ALL' ALL_EXCEPT_RESERVED = 'ALL_EXCEPT_RESERVED' AutocompleteItem = 'AutocompleteItem' @@ -103,7 +103,7 @@ class EntityName(str, Enum): ] -class BaseEntityType(str, Enum): +class BaseEntityType(StrEnum): pass @@ -131,7 +131,7 @@ class PipelineEntityType(BaseEntityType): STREAMING = PipelineType.STREAMING.value -class PermissionAccess(int, Enum): +class PermissionAccess(IntEnum): OWNER = 1 ADMIN = 2 # Editor: list, detail, create, update, delete @@ -171,7 +171,7 @@ class PermissionAccess(int, Enum): DISABLE_UNLESS_CONDITIONS = 1073741824 -class PermissionCondition(str, Enum): +class PermissionCondition(StrEnum): HAS_NOTEBOOK_EDIT_ACCESS = 'HAS_NOTEBOOK_EDIT_ACCESS' HAS_PIPELINE_EDIT_ACCESS = 'HAS_PIPELINE_EDIT_ACCESS' USER_OWNS_ENTITY = 'USER_OWNS_ENTITY' diff --git a/mage_ai/cache/constants.py b/mage_ai/cache/constants.py index 153e6789730..0c5a0dc5a99 100644 --- a/mage_ai/cache/constants.py +++ b/mage_ai/cache/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum CACHE_KEY_BLOCKS_TO_PIPELINE_MAPPING = 'blocks_to_pipeline_mapping' CACHE_KEY_BLOCK_ACTION_OBJECTS_MAPPING = 'block_action_objects_mapping' @@ -11,5 +11,5 @@ MAGE_CACHE_DIRECTORY_ENVIRONMENT_VARIABLE_NAME = 'MAGE_CACHE_DIRECTORY' -class CacheItemType(str, Enum): +class CacheItemType(StrEnum): DBT = 'dbt' diff --git a/mage_ai/cache/dbt/constants.py b/mage_ai/cache/dbt/constants.py index 84f2480a482..1a44a2bf6e1 100644 --- a/mage_ai/cache/dbt/constants.py +++ b/mage_ai/cache/dbt/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum from mage_ai.data_preparation.models.constants import ( BLOCK_TYPE_DIRECTORY_NAME, @@ -19,7 +19,7 @@ PROJECT_FILENAMES = [PROJECT_FILENAME, 'dbt_project.yaml'] -class FileType(str, Enum): +class FileType(StrEnum): MODEL = 'model' PROFILES = 'profiles' PROJECT = 'project' diff --git a/mage_ai/cluster_manager/constants.py b/mage_ai/cluster_manager/constants.py index 96e8719ff7e..5cd3942694e 100644 --- a/mage_ai/cluster_manager/constants.py +++ b/mage_ai/cluster_manager/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum # ECS environment variables ECS_CLUSTER_NAME = 'ECS_CLUSTER_NAME' @@ -29,7 +29,7 @@ NODE_PORT_SERVICE_TYPE = 'NodePort' -class ClusterType(str, Enum): +class ClusterType(StrEnum): EMR = 'emr' ECS = 'ecs' CLOUD_RUN = 'cloud_run' diff --git a/mage_ai/command_center/constants.py b/mage_ai/command_center/constants.py index 3a86cc35101..dc2bdf8c29e 100644 --- a/mage_ai/command_center/constants.py +++ b/mage_ai/command_center/constants.py @@ -1,14 +1,14 @@ -from enum import Enum +from enum import StrEnum SETTINGS_FILENAME = '.command_center.yaml' -class ItemTagEnum(str, Enum): +class ItemTagEnum(StrEnum): PINNED = 'pinned' RECENT = 'recent' -class ItemType(str, Enum): +class ItemType(StrEnum): ACTION = 'action' # Cast spell CREATE = 'create' # Conjure DELETE = 'delete' @@ -22,7 +22,7 @@ class ItemType(str, Enum): UPDATE = 'update' -class ObjectType(str, Enum): +class ObjectType(StrEnum): APPLICATION = 'application' APPLICATION_EXPANSION = 'application_expansion' AUTHENTICATION = 'authentication' @@ -42,11 +42,11 @@ class ObjectType(str, Enum): VERSION_CONTROL_FILE = 'version_control_file' -class ModeType(str, Enum): +class ModeType(StrEnum): VERSION_CONTROL = 'version_control' -class FileExtension(str, Enum): +class FileExtension(StrEnum): CSV = 'csv' JSON = 'json' MD = 'md' @@ -59,7 +59,7 @@ class FileExtension(str, Enum): YML = 'yml' -class ButtonActionType(str, Enum): +class ButtonActionType(StrEnum): ADD_APPLICATION = 'add_application' CLOSE_APPLICATION = 'close_application' # Go back out of the current application. CLOSE_COMMAND_CENTER = 'close_command_center' @@ -70,7 +70,7 @@ class ButtonActionType(str, Enum): SELECT_ITEM_FROM_REQUEST = 'select_item_from_request' -class InteractionType(str, Enum): +class InteractionType(StrEnum): CLICK = 'click' CLOSE_APPLICATION = ButtonActionType.CLOSE_APPLICATION.value CLOSE_COMMAND_CENTER = ButtonActionType.CLOSE_COMMAND_CENTER.value @@ -81,7 +81,7 @@ class InteractionType(str, Enum): SELECT_ITEM = 'select_item' -class ApplicationType(str, Enum): +class ApplicationType(StrEnum): DETAIL = 'detail' DETAIL_LIST = 'detail_list' EXPANSION = 'expansion' @@ -89,22 +89,22 @@ class ApplicationType(str, Enum): LIST = 'list' -class ValidationType(str, Enum): +class ValidationType(StrEnum): CONFIRMATION = 'confirmation' CUSTOM_VALIDATION_PARSERS = 'custom_validation_parsers' -class RenderLocationType(str, Enum): +class RenderLocationType(StrEnum): ITEMS_CONTAINER_AFTER = 'items_container_after' -class ApplicationExpansionUUID(str, Enum): +class ApplicationExpansionUUID(StrEnum): ArcaneLibrary = 'ArcaneLibrary' PortalTerminal = 'PortalTerminal' VersionControlFileDiffs = 'VersionControlFileDiffs' -class ApplicationExpansionStatus(str, Enum): +class ApplicationExpansionStatus(StrEnum): ACTIVE = 'ACTIVE' CLOSED = 'CLOSED' INACTIVE = 'INACTIVE' diff --git a/mage_ai/data_cleaner/column_types/constants.py b/mage_ai/data_cleaner/column_types/constants.py index 37131b3aa0b..39bcc6f7cd4 100644 --- a/mage_ai/data_cleaner/column_types/constants.py +++ b/mage_ai/data_cleaner/column_types/constants.py @@ -1,7 +1,7 @@ -from enum import Enum +from enum import StrEnum -class ColumnType(str, Enum): +class ColumnType(StrEnum): CATEGORY = 'category' CATEGORY_HIGH_CARDINALITY = 'category_high_cardinality' DATETIME = 'datetime' diff --git a/mage_ai/data_cleaner/transformer_actions/constants.py b/mage_ai/data_cleaner/transformer_actions/constants.py index 6f8110139a2..0786dfb0ace 100644 --- a/mage_ai/data_cleaner/transformer_actions/constants.py +++ b/mage_ai/data_cleaner/transformer_actions/constants.py @@ -1,9 +1,10 @@ -from mage_ai.data_cleaner.column_types.constants import ColumnType -from enum import Enum -import pandas as pd -import numpy as np import re +from enum import StrEnum +import numpy as np +import pandas as pd + +from mage_ai.data_cleaner.column_types.constants import ColumnType CONSTANT_IMPUTATION_DEFAULTS = { ColumnType.CATEGORY: 'missing', @@ -34,7 +35,7 @@ } -class ActionType(str, Enum): +class ActionType(StrEnum): ADD = 'add' AVERAGE = 'average' CLEAN_COLUMN_NAME = 'clean_column_name' @@ -73,18 +74,18 @@ class ActionType(str, Enum): STANDARDIZE = 'standardize' -class Axis(str, Enum): +class Axis(StrEnum): COLUMN = 'column' ROW = 'row' -class VariableType(str, Enum): +class VariableType(StrEnum): FEATURE = 'feature' FEATURE_SET = 'feature_set' FEATURE_SET_VERSION = 'feature_set_version' -class Operator(str, Enum): +class Operator(StrEnum): CONTAINS = 'contains' NOT_CONTAINS = 'not contains' EQUALS = '==' @@ -95,7 +96,7 @@ class Operator(str, Enum): LESS_THAN_OR_EQUALS = '<=' -class ImputationStrategy(str, Enum): +class ImputationStrategy(StrEnum): AVERAGE = 'average' COLUMN = 'column' CONSTANT = 'constant' diff --git a/mage_ai/data_preparation/logging/__init__.py b/mage_ai/data_preparation/logging/__init__.py index 9183eb93fed..027650b9160 100644 --- a/mage_ai/data_preparation/logging/__init__.py +++ b/mage_ai/data_preparation/logging/__init__.py @@ -1,12 +1,12 @@ from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import Dict from mage_ai.shared.config import BaseConfig from mage_ai.shared.logger import LoggingLevel -class LoggerType(str, Enum): +class LoggerType(StrEnum): DEFAULT = 'file' S3 = 's3' GCS = 'gcs' diff --git a/mage_ai/data_preparation/models/block/data_integration/constants.py b/mage_ai/data_preparation/models/block/data_integration/constants.py index 6c5e702c440..54c2a6aabd6 100644 --- a/mage_ai/data_preparation/models/block/data_integration/constants.py +++ b/mage_ai/data_preparation/models/block/data_integration/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum BLOCK_CATALOG_FILENAME = 'catalog.json' REPLICATION_METHOD_INCREMENTAL = 'INCREMENTAL' @@ -81,6 +81,6 @@ VARIABLE_BOOKMARK_VALUES_KEY = '__bookmark_values__' -class IngestMode(str, Enum): +class IngestMode(StrEnum): DISK = 'disk' MEMORY = 'memory' diff --git a/mage_ai/data_preparation/models/block/dbt/constants.py b/mage_ai/data_preparation/models/block/dbt/constants.py index 16ef8d7b30c..63243d4132f 100644 --- a/mage_ai/data_preparation/models/block/dbt/constants.py +++ b/mage_ai/data_preparation/models/block/dbt/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum DBT_DIRECTORY_NAME = 'dbt' @@ -13,12 +13,12 @@ ] -class Flag(str, Enum): +class Flag(StrEnum): PROFILES_DIR = 'profiles-dir' PROJECT_DIR = 'project-dir' -class LogLevel(str, Enum): +class LogLevel(StrEnum): DEBUG = 'debug' INFO = 'info' WARN = 'warn' diff --git a/mage_ai/data_preparation/models/block/dynamic/utils.py b/mage_ai/data_preparation/models/block/dynamic/utils.py index 7fde7cead9a..6a7b69c3eb9 100644 --- a/mage_ai/data_preparation/models/block/dynamic/utils.py +++ b/mage_ai/data_preparation/models/block/dynamic/utils.py @@ -1,7 +1,7 @@ import json import os from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import Any, Dict, List, Tuple, Union import pandas as pd @@ -20,7 +20,7 @@ from mage_ai.shared.models import BaseDataClass -class DynamicBlockFlag(str, Enum): +class DynamicBlockFlag(StrEnum): CLONE_OF_ORIGINAL = 'clone_of_original' DYNAMIC = 'dynamic' DYNAMIC_CHILD = 'dynamic_child' diff --git a/mage_ai/data_preparation/models/constants.py b/mage_ai/data_preparation/models/constants.py index df1aa15b384..5d9f404a97b 100644 --- a/mage_ai/data_preparation/models/constants.py +++ b/mage_ai/data_preparation/models/constants.py @@ -1,5 +1,5 @@ import os -from enum import Enum +from enum import StrEnum DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json' DATAFRAME_ANALYSIS_KEYS = frozenset( @@ -27,12 +27,12 @@ PIPELINE_RUN_STATUS_LAST_RUN_FAILED = 'last_run_failed' -class AIMode(str, Enum): +class AIMode(StrEnum): OPEN_AI = 'open_ai' HUGGING_FACE = 'hugging_face' -class BlockLanguage(str, Enum): +class BlockLanguage(StrEnum): MARKDOWN = 'markdown' PYTHON = 'python' R = 'r' @@ -40,14 +40,14 @@ class BlockLanguage(str, Enum): YAML = 'yaml' -class BlockStatus(str, Enum): +class BlockStatus(StrEnum): EXECUTED = 'executed' FAILED = 'failed' NOT_EXECUTED = 'not_executed' UPDATED = 'updated' -class BlockType(str, Enum): +class BlockType(StrEnum): CALLBACK = 'callback' CONDITIONAL = 'conditional' CHART = 'chart' @@ -65,7 +65,7 @@ class BlockType(str, Enum): TRANSFORMER = 'transformer' -class BlockColor(str, Enum): +class BlockColor(StrEnum): BLUE = 'blue' GREY = 'grey' PINK = 'pink' @@ -74,12 +74,12 @@ class BlockColor(str, Enum): YELLOW = 'yellow' -class CallbackStatus(str, Enum): +class CallbackStatus(StrEnum): FAILURE = 'failure' SUCCESS = 'success' -class ExecutorType(str, Enum): +class ExecutorType(StrEnum): AZURE_CONTAINER_INSTANCE = 'azure_container_instance' ECS = 'ecs' GCP_CLOUD_RUN = 'gcp_cloud_run' @@ -94,7 +94,7 @@ def is_valid_type(cls, executor_type: str) -> bool: return executor_type.upper() in cls.__members__ -class PipelineType(str, Enum): +class PipelineType(StrEnum): INTEGRATION = 'integration' DATABRICKS = 'databricks' PYTHON = 'python' @@ -102,7 +102,7 @@ class PipelineType(str, Enum): STREAMING = 'streaming' -class PipelineStatus(str, Enum): +class PipelineStatus(StrEnum): ACTIVE = 'active', # At least one active trigger INACTIVE = 'inactive', # All inactive triggers NO_SCHEDULES = 'no_schedules', # No triggers diff --git a/mage_ai/data_preparation/models/global_data_product/constants.py b/mage_ai/data_preparation/models/global_data_product/constants.py index ed84520142f..92269db4143 100644 --- a/mage_ai/data_preparation/models/global_data_product/constants.py +++ b/mage_ai/data_preparation/models/global_data_product/constants.py @@ -1,6 +1,6 @@ -import enum +from enum import StrEnum -class GlobalDataProductObjectType(str, enum.Enum): +class GlobalDataProductObjectType(StrEnum): BLOCK = 'block' PIPELINE = 'pipeline' diff --git a/mage_ai/data_preparation/models/global_hooks/constants.py b/mage_ai/data_preparation/models/global_hooks/constants.py index 74515927469..3281aae74e7 100644 --- a/mage_ai/data_preparation/models/global_hooks/constants.py +++ b/mage_ai/data_preparation/models/global_hooks/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum from mage_ai.authentication.permissions.constants import EntityName @@ -30,7 +30,7 @@ RESOURCE_TYPES = [en for en in EntityName if en not in DISABLED_RESOURCE_TYPES] -class HookOutputKey(str, Enum): +class HookOutputKey(StrEnum): ERROR = 'error' META = 'meta' METADATA = 'metadata' @@ -40,7 +40,7 @@ class HookOutputKey(str, Enum): RESOURCES = 'resources' -class HookInputKey(str, Enum): +class HookInputKey(StrEnum): HOOK = 'hook' PROJECT = 'project' RESOURCE_ID = 'resource_id' @@ -67,12 +67,12 @@ class HookInputKey(str, Enum): INTERNAL_DEFAULT_PREDICATE_VALUE = '__INTERNAL_DEFAULT_PREDICATE_VALUE__' -class PredicateAndOrOperator(str, Enum): +class PredicateAndOrOperator(StrEnum): AND = 'and' OR = 'or' -class PredicateObjectType(str, Enum): +class PredicateObjectType(StrEnum): ERROR = HookOutputKey.ERROR.value HOOK = HookInputKey.HOOK.value META = HookOutputKey.META.value @@ -89,7 +89,7 @@ class PredicateObjectType(str, Enum): USER = HookInputKey.USER.value -class PredicateOperator(str, Enum): +class PredicateOperator(StrEnum): EQUALS = 'EQUALS' GREATER_THAN = 'GREATER_THAN' GREATER_THAN_OR_EQUALS = 'GREATER_THAN_OR_EQUALS' @@ -102,7 +102,7 @@ class PredicateOperator(str, Enum): PRESENT = 'PRESENT' -class PredicateValueDataType(str, Enum): +class PredicateValueDataType(StrEnum): BOOLEAN = 'BOOLEAN' DICTIONARY = 'DICTIONARY' FLOAT = 'FLOAT' diff --git a/mage_ai/data_preparation/models/global_hooks/models.py b/mage_ai/data_preparation/models/global_hooks/models.py index 8dde57ca690..dc3822df873 100644 --- a/mage_ai/data_preparation/models/global_hooks/models.py +++ b/mage_ai/data_preparation/models/global_hooks/models.py @@ -2,7 +2,7 @@ import os from dataclasses import dataclass, field, make_dataclass from datetime import datetime -from enum import Enum +from enum import StrEnum from typing import Dict, List, Tuple, Union import yaml @@ -44,7 +44,7 @@ from mage_ai.shared.multi import run_parallel_multiple_args -class HookOperation(str, Enum): +class HookOperation(StrEnum): CREATE = OperationType.CREATE.value DELETE = OperationType.DELETE.value DETAIL = OperationType.DETAIL.value @@ -54,18 +54,18 @@ class HookOperation(str, Enum): UPDATE_ANYWHERE = 'update_anywhere' -class HookCondition(str, Enum): +class HookCondition(StrEnum): FAILURE = 'failure' SUCCESS = 'success' -class HookStrategy(str, Enum): +class HookStrategy(StrEnum): BREAK = 'break' CONTINUE = 'continue' RAISE = 'raise' -class HookStage(str, Enum): +class HookStage(StrEnum): AFTER = 'after' BEFORE = 'before' diff --git a/mage_ai/data_preparation/models/project/constants.py b/mage_ai/data_preparation/models/project/constants.py index a9b925125de..f4b97f79bd5 100644 --- a/mage_ai/data_preparation/models/project/constants.py +++ b/mage_ai/data_preparation/models/project/constants.py @@ -1,7 +1,7 @@ -from enum import Enum +from enum import StrEnum -class FeatureUUID(str, Enum): +class FeatureUUID(StrEnum): ADD_NEW_BLOCK_V2 = 'add_new_block_v2' CODE_BLOCK_V2 = 'code_block_v2' COMMAND_CENTER = 'command_center' diff --git a/mage_ai/data_preparation/models/triggers/__init__.py b/mage_ai/data_preparation/models/triggers/__init__.py index 84581e50dbd..6b4a807ff26 100644 --- a/mage_ai/data_preparation/models/triggers/__init__.py +++ b/mage_ai/data_preparation/models/triggers/__init__.py @@ -1,8 +1,8 @@ -import enum import os import traceback from dataclasses import dataclass, field from datetime import datetime +from enum import StrEnum from typing import Dict, List import yaml @@ -18,12 +18,12 @@ TRIGGER_FILE_NAME = 'triggers.yaml' -class ScheduleStatus(str, enum.Enum): +class ScheduleStatus(StrEnum): ACTIVE = 'active' INACTIVE = 'inactive' -class ScheduleType(str, enum.Enum): +class ScheduleType(StrEnum): API = 'api' EVENT = 'event' TIME = 'time' @@ -36,7 +36,7 @@ class ScheduleType(str, enum.Enum): } -class ScheduleInterval(str, enum.Enum): +class ScheduleInterval(StrEnum): ONCE = '@once' HOURLY = '@hourly' DAILY = '@daily' diff --git a/mage_ai/data_preparation/models/variable.py b/mage_ai/data_preparation/models/variable.py index f9e58ab35ef..aa22f468135 100644 --- a/mage_ai/data_preparation/models/variable.py +++ b/mage_ai/data_preparation/models/variable.py @@ -1,7 +1,7 @@ import os import traceback from contextlib import contextmanager -from enum import Enum +from enum import StrEnum from typing import Any, Dict, List import numpy as np @@ -44,7 +44,7 @@ JSON_SAMPLE_FILE = 'sample_data.json' -class VariableType(str, Enum): +class VariableType(StrEnum): DATAFRAME = 'dataframe' DATAFRAME_ANALYSIS = 'dataframe_analysis' GEO_DATAFRAME = 'geo_dataframe' diff --git a/mage_ai/data_preparation/models/widget/constants.py b/mage_ai/data_preparation/models/widget/constants.py index 59761142225..86b20664527 100644 --- a/mage_ai/data_preparation/models/widget/constants.py +++ b/mage_ai/data_preparation/models/widget/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum from dateutil.relativedelta import relativedelta @@ -12,7 +12,7 @@ VARIABLE_NAME_Y = 'y' -class AggregationFunction(str, Enum): +class AggregationFunction(StrEnum): AVERAGE = 'average' COUNT = 'count' COUNT_DISTINCT = 'count_distinct' @@ -23,7 +23,7 @@ class AggregationFunction(str, Enum): SUM = 'sum' -class TimeInterval(str, Enum): +class TimeInterval(StrEnum): DAY = 'day' HOUR = 'hour' MINUTE = 'minute' @@ -34,7 +34,7 @@ class TimeInterval(str, Enum): YEAR = 'year' -class ChartType(str, Enum): +class ChartType(StrEnum): BAR_CHART = 'bar chart' CUSTOM = 'custom' HISTOGRAM = 'histogram' diff --git a/mage_ai/data_preparation/repo_manager.py b/mage_ai/data_preparation/repo_manager.py index 03593e4420b..fd35516c03b 100644 --- a/mage_ai/data_preparation/repo_manager.py +++ b/mage_ai/data_preparation/repo_manager.py @@ -3,7 +3,7 @@ import os import traceback import uuid -from enum import Enum +from enum import StrEnum from typing import Dict, Optional from warnings import warn @@ -31,7 +31,7 @@ logger = logging.getLogger(__name__) -class ProjectType(str, Enum): +class ProjectType(StrEnum): MAIN = 'main' SUB = 'sub' STANDALONE = 'standalone' diff --git a/mage_ai/data_preparation/sync/__init__.py b/mage_ai/data_preparation/sync/__init__.py index e8cf6dddb4e..687db0f3f68 100644 --- a/mage_ai/data_preparation/sync/__init__.py +++ b/mage_ai/data_preparation/sync/__init__.py @@ -1,7 +1,7 @@ import inspect import os from dataclasses import dataclass -from enum import Enum +from enum import StrEnum from mage_ai.shared.config import BaseConfig @@ -10,7 +10,7 @@ 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' diff --git a/mage_ai/io/base.py b/mage_ai/io/base.py index d91cf488508..29161e16af3 100644 --- a/mage_ai/io/base.py +++ b/mage_ai/io/base.py @@ -1,6 +1,6 @@ import os from abc import ABC, abstractmethod -from enum import Enum +from enum import StrEnum from typing import IO, Any, Callable, Dict, Union import pandas as pd @@ -13,7 +13,7 @@ QUERY_ROW_LIMIT = 10_000_000 -class DataSource(str, Enum): +class DataSource(StrEnum): ALGOLIA = 'algolia' API = 'api' BIGQUERY = 'bigquery' @@ -39,7 +39,7 @@ class DataSource(str, Enum): WEAVIATE = 'weaviate' -class FileFormat(str, Enum): +class FileFormat(StrEnum): CSV = 'csv' JSON = 'json' PARQUET = 'parquet' @@ -47,7 +47,7 @@ class FileFormat(str, Enum): XML = 'xml' -class ExportWritePolicy(str, Enum): +class ExportWritePolicy(StrEnum): APPEND = 'append' FAIL = 'fail' REPLACE = 'replace' diff --git a/mage_ai/io/config.py b/mage_ai/io/config.py index 50b14fb22cb..77575c0ed6f 100644 --- a/mage_ai/io/config.py +++ b/mage_ai/io/config.py @@ -1,6 +1,6 @@ import os from abc import ABC, abstractmethod -from enum import Enum +from enum import StrEnum from pathlib import Path from typing import Any, Dict, Union @@ -11,7 +11,7 @@ from mage_ai.settings.repo import get_repo_path -class ConfigKey(str, Enum): +class ConfigKey(StrEnum): """ List of configuration settings for use with data IO clients. """ @@ -333,7 +333,7 @@ def get(self, env_var: Union[ConfigKey, str]) -> Any: return os.getenv(env_var) -class VerboseConfigKey(str, Enum): +class VerboseConfigKey(StrEnum): """ Config key headers for the verbose configuration file format. """ diff --git a/mage_ai/io/export_utils.py b/mage_ai/io/export_utils.py index 2b840bf4e26..f8b0db451b3 100644 --- a/mage_ai/io/export_utils.py +++ b/mage_ai/io/export_utils.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum from typing import Callable, Dict, List, Mapping from pandas import DataFrame, Series @@ -19,7 +19,7 @@ class BadConversionError(Exception): pass -class PandasTypes(str, Enum): +class PandasTypes(StrEnum): """ Internal datatypes defined by the pandas Public API """ diff --git a/mage_ai/io/io_config.py b/mage_ai/io/io_config.py index 899c6cf6362..de80508e2cb 100644 --- a/mage_ai/io/io_config.py +++ b/mage_ai/io/io_config.py @@ -1,12 +1,12 @@ import os -from enum import Enum +from enum import StrEnum from pathlib import Path from typing import Any, Mapping, Optional, Union import yaml -class IOConfigKeys(str, Enum): +class IOConfigKeys(StrEnum): AWS = 'AWS' BIGQUERY = 'BigQuery' CLICKHOUSE = 'ClickHouse' diff --git a/mage_ai/orchestration/constants.py b/mage_ai/orchestration/constants.py index 76b1b9c07e2..3a7de6a9b69 100644 --- a/mage_ai/orchestration/constants.py +++ b/mage_ai/orchestration/constants.py @@ -1,4 +1,4 @@ -import enum +from enum import StrEnum DATABASE_CONNECTION_URL_ENV_VAR = 'MAGE_DATABASE_CONNECTION_URL' @@ -17,7 +17,7 @@ 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' diff --git a/mage_ai/orchestration/db/constants.py b/mage_ai/orchestration/db/constants.py index 0e9a2017ff4..436beb0d8b7 100644 --- a/mage_ai/orchestration/db/constants.py +++ b/mage_ai/orchestration/db/constants.py @@ -1,6 +1,6 @@ -from enum import Enum +from enum import StrEnum -class DatabaseType(str, Enum): +class DatabaseType(StrEnum): POSTGRESQL = 'postgresql' SQLITE = 'sqlite' diff --git a/mage_ai/orchestration/db/models/oauth.py b/mage_ai/orchestration/db/models/oauth.py index bde7d06237f..d06d922ab53 100644 --- a/mage_ai/orchestration/db/models/oauth.py +++ b/mage_ai/orchestration/db/models/oauth.py @@ -1,6 +1,6 @@ -import enum import re from datetime import datetime +from enum import IntEnum, StrEnum from typing import Callable, Dict, List, Union from sqlalchemy import ( @@ -306,7 +306,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' @@ -568,7 +568,7 @@ def validate_user_id(self, key, value): class Permission(BaseModel): - class Access(int, enum.Enum): + class Access(IntEnum): OWNER = PermissionAccess.OWNER.value ADMIN = PermissionAccess.ADMIN.value EDITOR = PermissionAccess.EDITOR.value @@ -848,11 +848,11 @@ def validate_role_id(self, key, value): 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' diff --git a/mage_ai/orchestration/db/models/schedules.py b/mage_ai/orchestration/db/models/schedules.py index 8ea0cc2027e..9ab4d35f6ec 100644 --- a/mage_ai/orchestration/db/models/schedules.py +++ b/mage_ai/orchestration/db/models/schedules.py @@ -1,9 +1,9 @@ import asyncio import collections -import enum import traceback import uuid from datetime import datetime, timedelta, timezone +from enum import StrEnum from math import ceil from statistics import stdev from typing import DefaultDict, Dict, List @@ -731,7 +731,7 @@ def runtime_average( class PipelineRun(PipelineRunProjectPlatformMixin, BaseModel): - class PipelineRunStatus(str, enum.Enum): + class PipelineRunStatus(StrEnum): INITIAL = 'initial' RUNNING = 'running' COMPLETED = 'completed' @@ -1612,7 +1612,7 @@ def get_variables(self, extra_variables: Dict = None, pipeline_uuid: str = None) class BlockRun(BlockRunProjectPlatformMixin, BaseModel): - class BlockRunStatus(str, enum.Enum): + class BlockRunStatus(StrEnum): INITIAL = 'initial' QUEUED = 'queued' RUNNING = 'running' @@ -1764,7 +1764,7 @@ def get_outputs(self, sample_count: int = None) -> List[Dict]: class EventMatcher(BaseModel): - class EventType(str, enum.Enum): + class EventType(StrEnum): AWS_EVENT = 'aws_event' event_type = Column(Enum(EventType), default=EventType.AWS_EVENT) @@ -1867,7 +1867,7 @@ def __match_dict(sub_pattern, sub_config): class Backfill(BaseModel): - class IntervalType(str, enum.Enum): + class IntervalType(StrEnum): SECOND = 'second' MINUTE = 'minute' HOUR = 'hour' @@ -1877,7 +1877,7 @@ class IntervalType(str, enum.Enum): YEAR = 'year' CUSTOM = 'custom' - class Status(str, enum.Enum): + class Status(StrEnum): INITIAL = 'initial' RUNNING = 'running' COMPLETED = 'completed' diff --git a/mage_ai/orchestration/job_manager.py b/mage_ai/orchestration/job_manager.py index efa873932bb..c09f888370d 100644 --- a/mage_ai/orchestration/job_manager.py +++ b/mage_ai/orchestration/job_manager.py @@ -1,11 +1,11 @@ import traceback -from enum import Enum +from enum import StrEnum from typing import Callable, Dict, Union from mage_ai.orchestration.queue.queue_factory import QueueFactory -class JobType(str, Enum): +class JobType(StrEnum): BLOCK_RUN = 'block_run' PIPELINE_RUN = 'pipeline_run' INTEGRATION_STREAM = 'integration_stream' diff --git a/mage_ai/orchestration/monitor/monitor_stats.py b/mage_ai/orchestration/monitor/monitor_stats.py index 7aec5a0e6aa..9f345a1854b 100644 --- a/mage_ai/orchestration/monitor/monitor_stats.py +++ b/mage_ai/orchestration/monitor/monitor_stats.py @@ -1,5 +1,5 @@ -import enum from datetime import datetime, timedelta +from enum import StrEnum from functools import reduce from typing import Callable, Dict, List, Union @@ -21,7 +21,7 @@ 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' diff --git a/mage_ai/orchestration/notification/config.py b/mage_ai/orchestration/notification/config.py index 7a2d08e7c4e..4f6cf801898 100644 --- a/mage_ai/orchestration/notification/config.py +++ b/mage_ai/orchestration/notification/config.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import List from mage_ai.services.discord.config import DiscordConfig @@ -12,7 +12,7 @@ from mage_ai.shared.config import BaseConfig -class AlertOn(str, Enum): +class AlertOn(StrEnum): PIPELINE_RUN_FAILURE = 'trigger_failure' PIPELINE_RUN_SUCCESS = 'trigger_success' PIPELINE_RUN_PASSED_SLA = 'trigger_passed_sla' diff --git a/mage_ai/orchestration/queue/config.py b/mage_ai/orchestration/queue/config.py index 24d4438eb23..41a5773fedf 100644 --- a/mage_ai/orchestration/queue/config.py +++ b/mage_ai/orchestration/queue/config.py @@ -1,10 +1,10 @@ from dataclasses import dataclass -from enum import Enum +from enum import StrEnum from mage_ai.shared.config import BaseConfig -class QueueType(str, Enum): +class QueueType(StrEnum): CELERY = 'celery' PROCESS = 'process' diff --git a/mage_ai/orchestration/queue/process_queue.py b/mage_ai/orchestration/queue/process_queue.py index d01e57fff5d..c1f43097dfe 100644 --- a/mage_ai/orchestration/queue/process_queue.py +++ b/mage_ai/orchestration/queue/process_queue.py @@ -2,7 +2,7 @@ import os import signal import time -from enum import Enum +from enum import StrEnum from multiprocessing import Manager from typing import Callable, Dict @@ -29,14 +29,14 @@ 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' diff --git a/mage_ai/presenters/charts/data_sources/constants.py b/mage_ai/presenters/charts/data_sources/constants.py index 816d27f9570..275eb039aa8 100644 --- a/mage_ai/presenters/charts/data_sources/constants.py +++ b/mage_ai/presenters/charts/data_sources/constants.py @@ -1,9 +1,9 @@ -from enum import Enum +from enum import StrEnum DEFAULT_LIMIT = 10000 -class ChartDataSourceType(str, Enum): +class ChartDataSourceType(StrEnum): BLOCK = 'block' BLOCK_RUNS = 'block_runs' CHART_CODE = 'chart_code' diff --git a/mage_ai/presenters/interactions/constants.py b/mage_ai/presenters/interactions/constants.py index 740053a575d..7bd63093e4a 100644 --- a/mage_ai/presenters/interactions/constants.py +++ b/mage_ai/presenters/interactions/constants.py @@ -1,9 +1,9 @@ -from enum import Enum +from enum import StrEnum INTERACTIONS_DIRECTORY_NAME = 'interactions' -class InteractionInputType(str, Enum): +class InteractionInputType(StrEnum): CHECKBOX = 'checkbox' CODE = 'code' DROPDOWN_MENU = 'dropdown_menu' @@ -11,11 +11,11 @@ class InteractionInputType(str, Enum): TEXT_FIELD = 'text_field' -class InteractionInputStyleInputType(str, Enum): +class InteractionInputStyleInputType(StrEnum): NUMBER = 'number' -class InteractionVariableType(str, Enum): +class InteractionVariableType(StrEnum): BOOLEAN = 'boolean' DATE = 'date' DATETIME = 'datetime' diff --git a/mage_ai/presenters/pages/models/constants.py b/mage_ai/presenters/pages/models/constants.py index 1a34e7a8361..9c79b822ee2 100644 --- a/mage_ai/presenters/pages/models/constants.py +++ b/mage_ai/presenters/pages/models/constants.py @@ -1,16 +1,16 @@ -from enum import Enum +from enum import StrEnum -class ComponentCategory(str, Enum): +class ComponentCategory(StrEnum): BUTTON = 'button' FIELD = 'field' FORM = 'form' -class PageCategory(str, Enum): +class PageCategory(StrEnum): COMMUNITY = 'community' -class ResourceType(str, Enum): +class ResourceType(StrEnum): PIPELINE = 'pipeline' PIPELINE_SCHEDULE = 'pipeline_schedule' diff --git a/mage_ai/server/kernel_output_parser.py b/mage_ai/server/kernel_output_parser.py index 3cbe1980435..f5793360230 100644 --- a/mage_ai/server/kernel_output_parser.py +++ b/mage_ai/server/kernel_output_parser.py @@ -1,9 +1,9 @@ -from enum import Enum +from enum import StrEnum from mage_ai.data_preparation.models.constants import MAX_PRINT_OUTPUT_LINES -class DataType(str, Enum): +class DataType(StrEnum): DATA_FRAME = 'data_frame' IMAGE_PNG = 'image/png' PROGRESS = 'progress' diff --git a/mage_ai/server/kernels.py b/mage_ai/server/kernels.py index 04fde89ff8f..2a1aa262796 100644 --- a/mage_ai/server/kernels.py +++ b/mage_ai/server/kernels.py @@ -1,11 +1,12 @@ -from enum import Enum +from enum import StrEnum + from jupyter_client import KernelManager from jupyter_client.session import Session from mage_ai.data_preparation.models.constants import PipelineType -class KernelName(str, Enum): +class KernelName(StrEnum): PYSPARK = 'pysparkkernel' PYTHON3 = 'python3' diff --git a/mage_ai/server/scheduler_manager.py b/mage_ai/server/scheduler_manager.py index f742bda5dc6..bed281aa4e9 100644 --- a/mage_ai/server/scheduler_manager.py +++ b/mage_ai/server/scheduler_manager.py @@ -2,7 +2,7 @@ import time import traceback from contextlib import nullcontext -from enum import Enum +from enum import StrEnum import newrelic.agent import sentry_sdk @@ -63,7 +63,7 @@ class SchedulerManager: Singleton class to manage scheduler process. """ - class SchedulerStatus(str, Enum): + class SchedulerStatus(StrEnum): RUNNING = 'running' STOPPED = 'stopped' diff --git a/mage_ai/server/utils/frontend_renderer.py b/mage_ai/server/utils/frontend_renderer.py index a3478155afa..dfa30856a30 100644 --- a/mage_ai/server/utils/frontend_renderer.py +++ b/mage_ai/server/utils/frontend_renderer.py @@ -1,15 +1,13 @@ +import os +from enum import StrEnum + from IPython import get_ipython from IPython.display import IFrame, Javascript, display -from enum import Enum -from mage_ai.server.app import ( - server_config, -) + +from mage_ai.server.app import server_config from mage_ai.server.constants import SERVER_PORT from mage_ai.server.logger import Logger -import os - - FRONTEND_DIST_PATH = os.path.abspath( os.path.join(os.path.dirname(os.path.dirname(__file__)), 'frontend_dist'), ) @@ -19,7 +17,7 @@ logger = Logger().new_server_logger(__name__) -class NotebookType(str, Enum): +class NotebookType(StrEnum): DATABRICKS = 'databricks' GOOGLE_COLAB = 'google_colab' SAGEMAKER = 'sagemaker' diff --git a/mage_ai/server/websockets/constants.py b/mage_ai/server/websockets/constants.py index 6a955180572..f6f2e8df48d 100644 --- a/mage_ai/server/websockets/constants.py +++ b/mage_ai/server/websockets/constants.py @@ -1,13 +1,13 @@ -from enum import Enum +from enum import StrEnum -class Channel(str, Enum): +class Channel(StrEnum): TERMINAL = 'TERMINAL' -class ExecutionState(str, Enum): +class ExecutionState(StrEnum): IDLE = 'idle' -class MessageType(str, Enum): +class MessageType(StrEnum): DISPLAY_DATA = 'display_data' diff --git a/mage_ai/services/compute/aws/constants.py b/mage_ai/services/compute/aws/constants.py index 7dd52d7648a..35c201526a1 100644 --- a/mage_ai/services/compute/aws/constants.py +++ b/mage_ai/services/compute/aws/constants.py @@ -1,10 +1,10 @@ -from enum import Enum +from enum import StrEnum CONNECTION_CREDENTIAL_AWS_ACCESS_KEY_ID = 'AWS_ACCESS_KEY_ID' CONNECTION_CREDENTIAL_AWS_SECRET_ACCESS_KEY = 'AWS_SECRET_ACCESS_KEY' -class ClusterStatusState(str, Enum): +class ClusterStatusState(StrEnum): BOOTSTRAPPING = 'BOOTSTRAPPING' RUNNING = 'RUNNING' STARTING = 'STARTING' diff --git a/mage_ai/services/compute/aws/steps.py b/mage_ai/services/compute/aws/steps.py index 07d20d69215..75f9936662e 100644 --- a/mage_ai/services/compute/aws/steps.py +++ b/mage_ai/services/compute/aws/steps.py @@ -1,6 +1,6 @@ import os import socket -from enum import Enum +from enum import StrEnum from typing import List import requests @@ -53,7 +53,7 @@ ) -class SetupStepUUID(str, Enum): +class SetupStepUUID(StrEnum): ACTIVATE_CLUSTER = 'activate_cluster' AWS_ACCESS_KEY_ID = CONNECTION_CREDENTIAL_AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY = CONNECTION_CREDENTIAL_AWS_SECRET_ACCESS_KEY diff --git a/mage_ai/services/compute/constants.py b/mage_ai/services/compute/constants.py index f5d48421ffa..6885f6ea4bd 100644 --- a/mage_ai/services/compute/constants.py +++ b/mage_ai/services/compute/constants.py @@ -1,20 +1,20 @@ -from enum import Enum +from enum import StrEnum -class ComputeConnectionActionUUID(str, Enum): +class ComputeConnectionActionUUID(StrEnum): CREATE = 'CREATE' DELETE = 'DELETE' DESELECT = 'DESELECT' UPDATE = 'UPDATE' -class ComputeConnectionState(str, Enum): +class ComputeConnectionState(StrEnum): ACTIVE = 'ACTIVE' INACTIVE = 'INACTIVE' PENDING = 'PENDING' -class ComputeManagementApplicationTab(str, Enum): +class ComputeManagementApplicationTab(StrEnum): CLUSTERS = 'clusters' CONNECTION = 'connection' MONITORING = 'monitoring' diff --git a/mage_ai/services/compute/models.py b/mage_ai/services/compute/models.py index f3b7a659643..c6192c2613d 100644 --- a/mage_ai/services/compute/models.py +++ b/mage_ai/services/compute/models.py @@ -1,6 +1,6 @@ from abc import abstractmethod from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import Dict, List from mage_ai.data_preparation.models.project import Project @@ -35,7 +35,7 @@ def __post_init__(self): self.error = ErrorMessage.load(**self.error) -class SetupStepStatus(str, Enum): +class SetupStepStatus(StrEnum): COMPLETED = 'completed' ERROR = 'error' INCOMPLETE = 'incomplete' @@ -137,7 +137,6 @@ def __post_init__(self): class ComputeService: - uuid = ComputeServiceUUID.STANDALONE_CLUSTER def __init__(self, project: Project, with_clusters: bool = False): self.project = project @@ -169,7 +168,7 @@ def to_dict(self, **kwargs) -> Dict: @property def uuid(self) -> ComputeServiceUUID: - return self.__class__.uuid + return ComputeServiceUUID.STANDALONE_CLUSTER @abstractmethod def active_cluster(self, **kwargs) -> Dict: diff --git a/mage_ai/services/spark/constants.py b/mage_ai/services/spark/constants.py index 17a04dc7a25..bb9df8e55d5 100644 --- a/mage_ai/services/spark/constants.py +++ b/mage_ai/services/spark/constants.py @@ -1,12 +1,12 @@ -from enum import Enum +from enum import StrEnum SPARK_DIRECTORY_NAME = '.spark' -class ComputeServiceUUID(str, Enum): +class ComputeServiceUUID(StrEnum): AWS_EMR = 'AWS_EMR' STANDALONE_CLUSTER = 'STANDALONE_CLUSTER' -class SparkMaster(str, Enum): +class SparkMaster(StrEnum): LOCAL = 'local' diff --git a/mage_ai/services/spark/models/jobs.py b/mage_ai/services/spark/models/jobs.py index 9db7e867ca6..649c7bc184d 100644 --- a/mage_ai/services/spark/models/jobs.py +++ b/mage_ai/services/spark/models/jobs.py @@ -1,12 +1,12 @@ from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import Dict, List from mage_ai.services.spark.models.applications import Application from mage_ai.services.spark.models.base import BaseSparkModel -class JobStatus(str, Enum): +class JobStatus(StrEnum): FAILED = 'FAILED' RUNNING = 'RUNNING' SUCCEEDED = 'SUCCEEDED' diff --git a/mage_ai/services/spark/models/sqls.py b/mage_ai/services/spark/models/sqls.py index 5ba8ac72ccb..156363c8e31 100644 --- a/mage_ai/services/spark/models/sqls.py +++ b/mage_ai/services/spark/models/sqls.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import List from mage_ai.services.spark.models.applications import Application @@ -8,7 +8,7 @@ from mage_ai.services.spark.models.stages import StageAttempt -class SqlStatus(str, Enum): +class SqlStatus(StrEnum): COMPLETED = 'COMPLETED' diff --git a/mage_ai/services/spark/models/stages.py b/mage_ai/services/spark/models/stages.py index 2572d71503a..eed9a7cf6d4 100644 --- a/mage_ai/services/spark/models/stages.py +++ b/mage_ai/services/spark/models/stages.py @@ -1,5 +1,5 @@ from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import Dict, List from mage_ai.services.spark.models.applications import Application @@ -7,19 +7,19 @@ from mage_ai.services.spark.models.metrics import Metrics -class Locality(str, Enum): +class Locality(StrEnum): NODE_LOCAL = 'NODE_LOCAL' PROCESS_LOCAL = 'PROCESS_LOCAL' RACK_LOCAL = 'RACK_LOCAL' -class StageStatus(str, Enum): +class StageStatus(StrEnum): COMPLETE = 'COMPLETE' PENDING = 'PENDING' SKIPPED = 'SKIPPED' -class TaskStatus(str, Enum): +class TaskStatus(StrEnum): SUCCESS = 'SUCCESS' diff --git a/mage_ai/services/spark/models/threads.py b/mage_ai/services/spark/models/threads.py index 30101741f6f..b558ca41550 100644 --- a/mage_ai/services/spark/models/threads.py +++ b/mage_ai/services/spark/models/threads.py @@ -1,11 +1,11 @@ from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import List from mage_ai.services.spark.models.base import BaseSparkModel -class ThreadState(str, Enum): +class ThreadState(StrEnum): RUNNABLE = 'RUNNABLE' TIMED_WAITING = 'TIMED_WAITING' WAITING = 'WAITING' diff --git a/mage_ai/settings/backends.py b/mage_ai/settings/backends.py index 72b96ef8c3d..c5b241597ce 100644 --- a/mage_ai/settings/backends.py +++ b/mage_ai/settings/backends.py @@ -1,13 +1,13 @@ import base64 import logging import os -from enum import Enum +from enum import StrEnum from typing import Optional logger = logging.getLogger(__name__) -class BackendType(str, Enum): +class BackendType(StrEnum): """ Enum for the different types of settings backends. """ diff --git a/mage_ai/settings/models/configuration_option.py b/mage_ai/settings/models/configuration_option.py index 0e0add88762..ba2d2562ad4 100644 --- a/mage_ai/settings/models/configuration_option.py +++ b/mage_ai/settings/models/configuration_option.py @@ -1,7 +1,7 @@ import asyncio import os from dataclasses import dataclass -from enum import Enum +from enum import StrEnum from pathlib import Path from typing import Dict, Union @@ -37,11 +37,11 @@ async def read_file(full_path: str) -> str: return config -class ConfigurationType(str, Enum): +class ConfigurationType(StrEnum): DBT = 'dbt' -class OptionType(str, Enum): +class OptionType(StrEnum): PROFILES = 'profiles' PROJECTS = 'projects' TARGETS = 'targets' diff --git a/mage_ai/shared/constants.py b/mage_ai/shared/constants.py index 48f5a0b283d..477f39d76de 100644 --- a/mage_ai/shared/constants.py +++ b/mage_ai/shared/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from enum import StrEnum ENV_DEV = 'dev' ENV_PROD = 'prod' @@ -20,7 +20,7 @@ ENV_VAR_INSTANCE_TYPE = 'INSTANCE_TYPE' -class InstanceType(str, Enum): +class InstanceType(StrEnum): SERVER_AND_SCHEDULER = 'server_and_scheduler' SCHEDULER = 'scheduler' WEB_SERVER = 'web_server' diff --git a/mage_ai/shared/custom_logger.py b/mage_ai/shared/custom_logger.py index 729b35be483..55ab9001c8e 100644 --- a/mage_ai/shared/custom_logger.py +++ b/mage_ai/shared/custom_logger.py @@ -1,6 +1,6 @@ import inspect import logging -from enum import Enum +from enum import StrEnum import simplejson @@ -9,7 +9,7 @@ from mage_ai.shared.parsers import encode_complex -class Color(str, Enum): +class Color(StrEnum): BLUE = "\x1b[1;34m" BOLD_RED = "\x1b[31;1m" GREEN = "\x1b[1;32m" diff --git a/mage_ai/shared/logger.py b/mage_ai/shared/logger.py index 708be0de9eb..1a693f0f6c4 100644 --- a/mage_ai/shared/logger.py +++ b/mage_ai/shared/logger.py @@ -2,7 +2,7 @@ import logging import time from contextlib import contextmanager, redirect_stdout -from enum import Enum +from enum import StrEnum from typing import Callable, List from mage_ai.settings import SERVER_LOGGING_TEMPLATE @@ -107,7 +107,7 @@ def set_logging_format(logging_format: str = None, level: str = None) -> None: root_logger.exception('Invalid logging level %s', level) -class LoggingLevel(str, Enum): +class LoggingLevel(StrEnum): DEBUG = 'DEBUG' INFO = 'INFO' WARNING = 'WARNING' diff --git a/mage_ai/streaming/constants.py b/mage_ai/streaming/constants.py index e5088013d69..96a1d06fb2e 100644 --- a/mage_ai/streaming/constants.py +++ b/mage_ai/streaming/constants.py @@ -1,10 +1,10 @@ -from enum import Enum +from enum import StrEnum DEFAULT_BATCH_SIZE = 100 DEFAULT_TIMEOUT_MS = 500 -class SourceType(str, Enum): +class SourceType(StrEnum): ACTIVEMQ = 'activemq' AMAZON_SQS = 'amazon_sqs' AZURE_EVENT_HUB = 'azure_event_hub' @@ -17,7 +17,7 @@ class SourceType(str, Enum): MONGODB = 'mongodb' -class SinkType(str, Enum): +class SinkType(StrEnum): ACTIVEMQ = 'activemq' AMAZON_S3 = 'amazon_s3' AZURE_DATA_LAKE = 'azure_data_lake' diff --git a/mage_ai/streaming/sinks/kafka.py b/mage_ai/streaming/sinks/kafka.py index 99c372d56a1..9af1c020d4c 100644 --- a/mage_ai/streaming/sinks/kafka.py +++ b/mage_ai/streaming/sinks/kafka.py @@ -1,7 +1,7 @@ import json import time from dataclasses import dataclass -from enum import Enum +from enum import StrEnum from typing import Dict, List from kafka import KafkaProducer @@ -11,7 +11,7 @@ from mage_ai.streaming.sinks.base import BaseSink -class SecurityProtocol(str, Enum): +class SecurityProtocol(StrEnum): SASL_SSL = 'SASL_SSL' SSL = 'SSL' diff --git a/mage_ai/streaming/sources/amazon_sqs.py b/mage_ai/streaming/sources/amazon_sqs.py index 846f62f5007..f02a3ffd6ac 100644 --- a/mage_ai/streaming/sources/amazon_sqs.py +++ b/mage_ai/streaming/sources/amazon_sqs.py @@ -2,7 +2,7 @@ import threading import time from dataclasses import dataclass -from enum import Enum +from enum import StrEnum from typing import Callable, Dict import boto3 @@ -15,7 +15,7 @@ DEFAULT_WAIT_TIME_SECONDS = 1 -class MessageDeletionMethod(str, Enum): +class MessageDeletionMethod(StrEnum): AFTER_RECEIVED = 'AFTER_RECEIVED' MANUAL = 'MANUAL' diff --git a/mage_ai/streaming/sources/base.py b/mage_ai/streaming/sources/base.py index a27d97406a1..7fd75907e56 100644 --- a/mage_ai/streaming/sources/base.py +++ b/mage_ai/streaming/sources/base.py @@ -1,12 +1,12 @@ import json from abc import ABC, abstractmethod -from enum import Enum +from enum import StrEnum from typing import Callable, Dict from mage_ai.shared.environments import is_test -class SourceConsumeMethod(str, Enum): +class SourceConsumeMethod(StrEnum): BATCH_READ = 'BATCH_READ' READ = 'READ' READ_ASYNC = 'READ_ASYNC' diff --git a/mage_ai/streaming/sources/kafka.py b/mage_ai/streaming/sources/kafka.py index 1198a02371b..c0d1812f973 100644 --- a/mage_ai/streaming/sources/kafka.py +++ b/mage_ai/streaming/sources/kafka.py @@ -2,7 +2,7 @@ import json import time from dataclasses import dataclass, field -from enum import Enum +from enum import StrEnum from typing import Callable, Dict, List from kafka import KafkaConsumer, TopicPartition @@ -13,7 +13,7 @@ from mage_ai.streaming.sources.shared import SerDeConfig, SerializationMethod -class SecurityProtocol(str, Enum): +class SecurityProtocol(StrEnum): SASL_PLAINTEXT = 'SASL_PLAINTEXT' SASL_SSL = 'SASL_SSL' SSL = 'SSL' diff --git a/mage_ai/streaming/sources/nats_js.py b/mage_ai/streaming/sources/nats_js.py index 90445e1c6e0..347bca25b42 100644 --- a/mage_ai/streaming/sources/nats_js.py +++ b/mage_ai/streaming/sources/nats_js.py @@ -1,9 +1,9 @@ import asyncio -import enum import json import ssl import threading from dataclasses import dataclass +from enum import StrEnum from typing import Callable, Dict, Optional import nats @@ -22,7 +22,7 @@ class SSLConfig: check_hostname: bool = False -class ConsumerType(str, enum.Enum): +class ConsumerType(StrEnum): PULL = "PULL" PUSH = "PUSH" diff --git a/mage_ai/streaming/sources/shared.py b/mage_ai/streaming/sources/shared.py index 1ace16e2ca2..c8a65b9cfbc 100644 --- a/mage_ai/streaming/sources/shared.py +++ b/mage_ai/streaming/sources/shared.py @@ -1,8 +1,8 @@ from dataclasses import dataclass -from enum import Enum +from enum import StrEnum -class SerializationMethod(str, Enum): +class SerializationMethod(StrEnum): AVRO = 'AVRO' JSON = 'JSON' PROTOBUF = 'PROTOBUF' diff --git a/mage_ai/tests/api/policies/permissions/test_base_policy_with_permissions.py b/mage_ai/tests/api/policies/permissions/test_base_policy_with_permissions.py index 0767ff0cf49..d5cad11ea18 100644 --- a/mage_ai/tests/api/policies/permissions/test_base_policy_with_permissions.py +++ b/mage_ai/tests/api/policies/permissions/test_base_policy_with_permissions.py @@ -1,5 +1,5 @@ import secrets -from enum import Enum +from enum import StrEnum from typing import Any, Callable, Dict, List, Union from unittest.mock import patch @@ -25,7 +25,7 @@ from mage_ai.tests.api.policies.permissions.mixins import PermissionsMixin -class TestSuite(str, Enum): +class TestSuite(StrEnum): AUTHORIZED = 'AUTHORIZED' DISABLED = 'DISABLED' INVERSE = 'INVERSE' diff --git a/mage_ai/usage_statistics/constants.py b/mage_ai/usage_statistics/constants.py index b393fd31fff..bbc40629194 100644 --- a/mage_ai/usage_statistics/constants.py +++ b/mage_ai/usage_statistics/constants.py @@ -1,9 +1,9 @@ -import enum +from enum import StrEnum API_ENDPOINT = 'https://api.mage.ai/v1/usage_statistics' -class EventNameType(str, enum.Enum): +class EventNameType(StrEnum): API_ERROR = 'api_error' APPLICATION_ERROR = 'application_error' BLOCK_RUN_ENDED = 'block_run_ended' @@ -13,14 +13,14 @@ class EventNameType(str, enum.Enum): USAGE_STATISTIC_CREATE = 'usage_statistic.create' -class EventActionType(str, enum.Enum): +class EventActionType(StrEnum): CREATE = 'create' DENY = 'deny' EXECUTE = 'execute' IMPRESSION = 'impression' -class EventObjectType(str, enum.Enum): +class EventObjectType(StrEnum): BLOCK = 'block' BLOCK_RUN = 'block_run' CHART = 'chart' diff --git a/mage_integrations/mage_integrations/connections/mysql/__init__.py b/mage_integrations/mage_integrations/connections/mysql/__init__.py index 28cfb1d3564..bc9a3b6b8a0 100644 --- a/mage_integrations/mage_integrations/connections/mysql/__init__.py +++ b/mage_integrations/mage_integrations/connections/mysql/__init__.py @@ -1,6 +1,6 @@ -import enum import io import os +from enum import StrEnum from typing import Dict import paramiko @@ -10,7 +10,7 @@ from mage_integrations.connections.sql.base import Connection -class ConnectionMethod(str, enum.Enum): +class ConnectionMethod(StrEnum): DIRECT = 'direct' SSH_TUNNEL = 'ssh_tunnel' diff --git a/mage_integrations/mage_integrations/sources/couchbase/__init__.py b/mage_integrations/mage_integrations/sources/couchbase/__init__.py index eb913bc16d6..fe7762ecf99 100644 --- a/mage_integrations/mage_integrations/sources/couchbase/__init__.py +++ b/mage_integrations/mage_integrations/sources/couchbase/__init__.py @@ -1,7 +1,9 @@ -from enum import Enum -from mage_integrations.connections.couchbase import ( - Couchbase as CouchbaseConnection -) +from enum import StrEnum +from typing import Any, Dict, List + +from singer import catalog + +from mage_integrations.connections.couchbase import Couchbase as CouchbaseConnection from mage_integrations.sources.base import main from mage_integrations.sources.catalog import Catalog from mage_integrations.sources.constants import ( @@ -14,13 +16,11 @@ wrap_column_in_quotes, ) from mage_integrations.sources.sql.base import Source -from singer import catalog -from typing import Any, Dict, List DEFAULT_COLUMN_NAME = '_document' -class SchemaStrategy(str, Enum): +class SchemaStrategy(StrEnum): INFER = 'infer' COMBINE = 'combine' @@ -42,21 +42,21 @@ def discover(self, streams: List[str] = None) -> Catalog: catalog_entries = [] for stream_id in collection_names: properties = dict() - + infer_query = f""" INFER `{stream_id}` WITH {{"sample_size": 1000, "similarity_metric": 0.4, "dictionary_threshold": 3}} """ infer_result = connection.load(infer_query)[0] - + def get_infer_result_doc_count(result): props = result.get('properties', {}) doc_count = props.get('#docs', 0) if type(doc_count) is list: doc_count = sum(doc_count) return doc_count - + strategy = self.config.get('strategy') if strategy == SchemaStrategy.INFER or \ (strategy is None and len(infer_result) == 1): @@ -79,7 +79,7 @@ def get_infer_result_doc_count(result): dict( type=[COLUMN_TYPE_NULL, COLUMN_TYPE_OBJECT], additionalProperties=True - ) + ) schema = catalog.Schema.from_dict(dict( properties=properties, @@ -96,7 +96,7 @@ def __get_type(self, dtype: str) -> str: column_type = dtype if dtype == 'missing': column_type = COLUMN_TYPE_STRING - + return column_type def _build_comparison_statement( diff --git a/mage_integrations/mage_integrations/sources/sql/constants.py b/mage_integrations/mage_integrations/sources/sql/constants.py index 45c7004c891..5e0e5a61b8f 100644 --- a/mage_integrations/mage_integrations/sources/sql/constants.py +++ b/mage_integrations/mage_integrations/sources/sql/constants.py @@ -1,7 +1,7 @@ -from enum import Enum +from enum import StrEnum -class PredicateOperator(str, Enum): +class PredicateOperator(StrEnum): EQUALS = 'EQUALS' GREATER_THAN = 'GREATER_THAN' GREATER_THAN_OR_EQUALS = 'GREATER_THAN_OR_EQUALS' diff --git a/mage_integrations/requirements.txt b/mage_integrations/requirements.txt index 4fc8b8034f8..0f7383bd88a 100644 --- a/mage_integrations/requirements.txt +++ b/mage_integrations/requirements.txt @@ -2,7 +2,7 @@ attrs==17.4.0 azure-storage-blob==12.14.1 backoff clickhouse_sqlalchemy~=0.2.4 -couchbase==4.1.1 +couchbase==4.1.6 deltalake==0.7.0 elasticsearch==8.9.0 facebook_business==17.0.2 diff --git a/requirements.txt b/requirements.txt index 09da5978908..17af376624a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -84,7 +84,8 @@ redshift-connector==2.0.915 lxml==4.9.4 snowflake-connector-python==3.5.0 sshtunnel==0.4.0 -tables==3.7.0 +tables==3.7.0; python_version == '3.8' +tables==3.9.2; python_version >= '3.9' trino~=0.326 # Azure