diff --git a/.github/workflows/build_and_test.yml b/.github/workflows/build_and_test.yml index 1d3124e0aee7..d62231cb8fb5 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 8bdcadbf01b0..b9fdbc8ca19b 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 532bc26e3e92..e995890145e3 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 5aefb1049e48..1faf88b7f173 100644 --- a/mage_ai/ai/constants.py +++ b/mage_ai/ai/constants.py @@ -1,7 +1,7 @@ -from enum import Enum +from mage_ai.shared.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 367e0fd02651..7d1c3fc22eff 100644 --- a/mage_ai/api/constants.py +++ b/mage_ai/api/constants.py @@ -1,20 +1,19 @@ -from enum import Enum - from mage_ai.api.operations.constants import OperationType from mage_ai.orchestration.db.models.oauth import Permission +from mage_ai.shared.enum import StrEnum -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 93f3575d6fe4..d35718c4ac46 100644 --- a/mage_ai/api/oauth_scope.py +++ b/mage_ai/api/oauth_scope.py @@ -1,4 +1,4 @@ -from enum import Enum +from mage_ai.shared.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 1d7b547412ce..9c8da09adc10 100644 --- a/mage_ai/api/operations/constants.py +++ b/mage_ai/api/operations/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from mage_ai.shared.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 bf2b1cd4b3ad..d0d6d5480446 100644 --- a/mage_ai/authentication/oauth/constants.py +++ b/mage_ai/authentication/oauth/constants.py @@ -1,8 +1,8 @@ -from enum import Enum from typing import Optional from mage_ai.settings import get_settings_value from mage_ai.settings.keys import GHE_HOSTNAME +from mage_ai.shared.enum import StrEnum ACTIVE_DIRECTORY_CLIENT_ID = '51aec820-9d49-40a9-b046-17c1f28f620d' @@ -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 f439027ccfb1..50d157ac3e3c 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 mage_ai.shared.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 cf6d064690da..5c19608ecd4f 100644 --- a/mage_ai/authentication/permissions/constants.py +++ b/mage_ai/authentication/permissions/constants.py @@ -1,10 +1,9 @@ -from enum import Enum - from mage_ai.api.operations.constants import OperationType from mage_ai.data_preparation.models.constants import BlockType, PipelineType +from mage_ai.shared.enum import IntEnum, StrEnum -class EntityName(str, Enum): +class EntityName(StrEnum): ALL = 'ALL' ALL_EXCEPT_RESERVED = 'ALL_EXCEPT_RESERVED' AutocompleteItem = 'AutocompleteItem' @@ -103,7 +102,7 @@ class EntityName(str, Enum): ] -class BaseEntityType(str, Enum): +class BaseEntityType(StrEnum): pass @@ -131,7 +130,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 +170,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 153e67897305..4902300a6f87 100644 --- a/mage_ai/cache/constants.py +++ b/mage_ai/cache/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from mage_ai.shared.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 84f2480a482b..75ed9bb3d5cb 100644 --- a/mage_ai/cache/dbt/constants.py +++ b/mage_ai/cache/dbt/constants.py @@ -1,10 +1,9 @@ -from enum import Enum - from mage_ai.data_preparation.models.constants import ( BLOCK_TYPE_DIRECTORY_NAME, PIPELINES_FOLDER, BlockType, ) +from mage_ai.shared.enum import StrEnum from mage_ai.shared.hash import merge_dict IGNORE_DIRECTORY_NAMES = merge_dict( @@ -19,7 +18,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 96e8719ff7e5..c4012d1703e5 100644 --- a/mage_ai/cluster_manager/constants.py +++ b/mage_ai/cluster_manager/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from mage_ai.shared.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 3a86cc35101d..92a1d488a57d 100644 --- a/mage_ai/command_center/constants.py +++ b/mage_ai/command_center/constants.py @@ -1,14 +1,14 @@ -from enum import Enum +from mage_ai.shared.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 37131b3aa0b9..10249a176c8c 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 mage_ai.shared.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 6f8110139a2e..e90a6784545f 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 +import numpy as np +import pandas as pd + +from mage_ai.data_cleaner.column_types.constants import ColumnType +from mage_ai.shared.enum import StrEnum 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 9183eb93fed8..18442867240a 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 typing import Dict from mage_ai.shared.config import BaseConfig +from mage_ai.shared.enum import StrEnum 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 6c5e702c440b..8451c3d05ea1 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 mage_ai.shared.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 16ef8d7b30cd..c44859a57a02 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 mage_ai.shared.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 7fde7cead9aa..b5313f197141 100644 --- a/mage_ai/data_preparation/models/block/dynamic/utils.py +++ b/mage_ai/data_preparation/models/block/dynamic/utils.py @@ -1,7 +1,6 @@ import json import os from dataclasses import dataclass, field -from enum import Enum from typing import Any, Dict, List, Tuple, Union import pandas as pd @@ -16,11 +15,12 @@ from mage_ai.server.kernel_output_parser import DataType from mage_ai.shared.array import find from mage_ai.shared.custom_logger import DX_PRINTER +from mage_ai.shared.enum import StrEnum from mage_ai.shared.hash import ignore_keys_with_blank_values 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 df1aa15b3845..38a11530f018 100644 --- a/mage_ai/data_preparation/models/constants.py +++ b/mage_ai/data_preparation/models/constants.py @@ -1,5 +1,6 @@ import os -from enum import Enum + +from mage_ai.shared.enum import StrEnum DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json' DATAFRAME_ANALYSIS_KEYS = frozenset( @@ -27,12 +28,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 +41,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 +66,7 @@ class BlockType(str, Enum): TRANSFORMER = 'transformer' -class BlockColor(str, Enum): +class BlockColor(StrEnum): BLUE = 'blue' GREY = 'grey' PINK = 'pink' @@ -74,12 +75,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 +95,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 +103,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 ed84520142f2..0fa1db6d1227 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 mage_ai.shared.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 745159274690..3f657b1e5151 100644 --- a/mage_ai/data_preparation/models/global_hooks/constants.py +++ b/mage_ai/data_preparation/models/global_hooks/constants.py @@ -1,6 +1,5 @@ -from enum import Enum - from mage_ai.authentication.permissions.constants import EntityName +from mage_ai.shared.enum import StrEnum GLOBAL_HOOKS_FILENAME = 'global_hooks.yaml' @@ -30,7 +29,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 +39,7 @@ class HookOutputKey(str, Enum): RESOURCES = 'resources' -class HookInputKey(str, Enum): +class HookInputKey(StrEnum): HOOK = 'hook' PROJECT = 'project' RESOURCE_ID = 'resource_id' @@ -67,12 +66,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 +88,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 +101,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 8dde57ca6909..95ec8744ef63 100644 --- a/mage_ai/data_preparation/models/global_hooks/models.py +++ b/mage_ai/data_preparation/models/global_hooks/models.py @@ -2,7 +2,6 @@ import os from dataclasses import dataclass, field, make_dataclass from datetime import datetime -from enum import Enum from typing import Dict, List, Tuple, Union import yaml @@ -30,6 +29,7 @@ ) from mage_ai.settings.repo import get_repo_path from mage_ai.shared.array import find, find_index, flatten +from mage_ai.shared.enum import StrEnum from mage_ai.shared.environments import is_debug, is_test from mage_ai.shared.hash import ( dig, @@ -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 a9b925125de7..9ae38751d31e 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 mage_ai.shared.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 84581e50dbde..1e2219e65f52 100644 --- a/mage_ai/data_preparation/models/triggers/__init__.py +++ b/mage_ai/data_preparation/models/triggers/__init__.py @@ -1,4 +1,3 @@ -import enum import os import traceback from dataclasses import dataclass, field @@ -12,18 +11,19 @@ from mage_ai.settings.repo import get_repo_path from mage_ai.shared.config import BaseConfig from mage_ai.shared.constants import VALID_ENVS +from mage_ai.shared.enum import StrEnum from mage_ai.shared.hash import index_by from mage_ai.shared.io import safe_write 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 f9e58ab35ef3..8127956106c8 100644 --- a/mage_ai/data_preparation/models/variable.py +++ b/mage_ai/data_preparation/models/variable.py @@ -1,7 +1,6 @@ import os import traceback from contextlib import contextmanager -from enum import Enum from typing import Any, Dict, List import numpy as np @@ -30,6 +29,7 @@ ) from mage_ai.data_preparation.storage.base_storage import BaseStorage from mage_ai.data_preparation.storage.local_storage import LocalStorage +from mage_ai.shared.enum import StrEnum from mage_ai.shared.parsers import sample_output from mage_ai.shared.utils import clean_name @@ -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 597611422251..a29b3cffa0ea 100644 --- a/mage_ai/data_preparation/models/widget/constants.py +++ b/mage_ai/data_preparation/models/widget/constants.py @@ -1,7 +1,7 @@ -from enum import Enum - from dateutil.relativedelta import relativedelta +from mage_ai.shared.enum import StrEnum + VARIABLE_NAME_BUCKETS = 'buckets' VARIABLE_NAME_GROUP_BY = 'group_by' VARIABLE_NAME_INDEX = 'index' @@ -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 03593e4420b3..5b1d77038c69 100644 --- a/mage_ai/data_preparation/repo_manager.py +++ b/mage_ai/data_preparation/repo_manager.py @@ -3,7 +3,6 @@ import os import traceback import uuid -from enum import Enum from typing import Dict, Optional from warnings import warn @@ -21,6 +20,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.shared.enum import StrEnum from mage_ai.shared.environments import is_debug yml = ruamel.yaml.YAML() @@ -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 e8cf6dddb4e5..b54364ed736a 100644 --- a/mage_ai/data_preparation/sync/__init__.py +++ b/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' diff --git a/mage_ai/io/base.py b/mage_ai/io/base.py index d91cf488508f..dab998a93bc5 100644 --- a/mage_ai/io/base.py +++ b/mage_ai/io/base.py @@ -1,19 +1,19 @@ import os from abc import ABC, abstractmethod -from enum import Enum from typing import IO, Any, Callable, Dict, Union import pandas as pd 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.utils import clean_name 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 50b14fb22cb5..88439574ddbd 100644 --- a/mage_ai/io/config.py +++ b/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 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 @@ 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 2b840bf4e261..1d3530640a3b 100644 --- a/mage_ai/io/export_utils.py +++ b/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 """ diff --git a/mage_ai/io/io_config.py b/mage_ai/io/io_config.py index 899c6cf63626..8ee59ba51812 100644 --- a/mage_ai/io/io_config.py +++ b/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' diff --git a/mage_ai/orchestration/constants.py b/mage_ai/orchestration/constants.py index 76b1b9c07e2d..2fa8137f75f6 100644 --- a/mage_ai/orchestration/constants.py +++ b/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' @@ -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 0e9a2017ff49..ada08623edee 100644 --- a/mage_ai/orchestration/db/constants.py +++ b/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' diff --git a/mage_ai/orchestration/db/models/oauth.py b/mage_ai/orchestration/db/models/oauth.py index bde7d06237f3..b85c9f2e37c4 100644 --- a/mage_ai/orchestration/db/models/oauth.py +++ b/mage_ai/orchestration/db/models/oauth.py @@ -1,4 +1,3 @@ -import enum import re from datetime import datetime from typing import Callable, Dict, List, Union @@ -32,6 +31,7 @@ from mage_ai.orchestration.db.models.base import BaseModel from mage_ai.settings.repo import get_repo_path 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 @@ -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 8ea0cc2027e2..f148b9b5eb46 100644 --- a/mage_ai/orchestration/db/models/schedules.py +++ b/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 @@ -72,6 +71,7 @@ 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 @@ -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 37df845e9aee..b15ea049a591 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 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' diff --git a/mage_ai/orchestration/monitor/monitor_stats.py b/mage_ai/orchestration/monitor/monitor_stats.py index 7aec5a0e6aac..fc8db931a0fe 100644 --- a/mage_ai/orchestration/monitor/monitor_stats.py +++ b/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.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' diff --git a/mage_ai/orchestration/notification/config.py b/mage_ai/orchestration/notification/config.py index 7a2d08e7c4e9..0267e2abcb14 100644 --- a/mage_ai/orchestration/notification/config.py +++ b/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.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' diff --git a/mage_ai/orchestration/queue/config.py b/mage_ai/orchestration/queue/config.py index 24d4438eb238..9a878d3e991f 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 mage_ai.shared.config import BaseConfig +from mage_ai.shared.enum import StrEnum -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 d01e57fff5d3..355fc81f9cde 100644 --- a/mage_ai/orchestration/queue/process_queue.py +++ b/mage_ai/orchestration/queue/process_queue.py @@ -2,7 +2,6 @@ import os import signal import time -from enum import Enum from multiprocessing import Manager from typing import Callable, Dict @@ -24,19 +23,20 @@ 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' diff --git a/mage_ai/presenters/charts/data_sources/constants.py b/mage_ai/presenters/charts/data_sources/constants.py index 816d27f9570f..b49d3663e5f2 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 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' diff --git a/mage_ai/presenters/interactions/constants.py b/mage_ai/presenters/interactions/constants.py index 740053a575d8..78385addd7f8 100644 --- a/mage_ai/presenters/interactions/constants.py +++ b/mage_ai/presenters/interactions/constants.py @@ -1,9 +1,9 @@ -from enum import Enum +from mage_ai.shared.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 1a34e7a8361e..4d730c28c864 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 mage_ai.shared.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 3cbe19804350..03cbde4a099d 100644 --- a/mage_ai/server/kernel_output_parser.py +++ b/mage_ai/server/kernel_output_parser.py @@ -1,9 +1,8 @@ -from enum import Enum - from mage_ai.data_preparation.models.constants import MAX_PRINT_OUTPUT_LINES +from mage_ai.shared.enum import StrEnum -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 04fde89ff8f4..42f7f0fd761b 100644 --- a/mage_ai/server/kernels.py +++ b/mage_ai/server/kernels.py @@ -1,11 +1,11 @@ -from enum import Enum from jupyter_client import KernelManager from jupyter_client.session import Session from mage_ai.data_preparation.models.constants import PipelineType +from mage_ai.shared.enum import StrEnum -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 f742bda5dc6b..988c5fb4ab71 100644 --- a/mage_ai/server/scheduler_manager.py +++ b/mage_ai/server/scheduler_manager.py @@ -2,7 +2,6 @@ import time import traceback from contextlib import nullcontext -from enum import Enum import newrelic.agent import sentry_sdk @@ -18,6 +17,7 @@ SERVER_LOGGING_FORMAT, SERVER_VERBOSITY, ) +from mage_ai.shared.enum import StrEnum from mage_ai.shared.logger import set_logging_format SCHEDULER_AUTO_RESTART_INTERVAL = 20_000 # in milliseconds @@ -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 a3478155afa0..eb7d80793ff3 100644 --- a/mage_ai/server/utils/frontend_renderer.py +++ b/mage_ai/server/utils/frontend_renderer.py @@ -1,14 +1,12 @@ +import os + 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 - +from mage_ai.shared.enum import StrEnum 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' @@ -36,7 +34,8 @@ def infer_notebook_type(): return None -def display_inline_iframe(host='localhost', port=SERVER_PORT, notebook_type=None, config={}): +def display_inline_iframe(host='localhost', port=SERVER_PORT, notebook_type=None, config=None): + config = dict(config) if config is not None else {} path_to_server = f'http://{host}:{port}' def __print_url(): @@ -72,7 +71,8 @@ def __print_url(): display(IFrame(path_to_server, width='95%', height=1000)) -def update_frontend_urls(host=None, port=None, notebook_type=None, config={}): +def update_frontend_urls(host=None, port=None, notebook_type=None, config=None): + config = dict(config) if config is not None else {} if notebook_type == NotebookType.DATABRICKS: required_args = [ 'cluster_id', diff --git a/mage_ai/server/websockets/constants.py b/mage_ai/server/websockets/constants.py index 6a955180572b..9ad3b3d02c6e 100644 --- a/mage_ai/server/websockets/constants.py +++ b/mage_ai/server/websockets/constants.py @@ -1,13 +1,13 @@ -from enum import Enum +from mage_ai.shared.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 7dd52d7648a8..cade13067015 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 mage_ai.shared.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 07d20d692150..a300ca0afbe1 100644 --- a/mage_ai/services/compute/aws/steps.py +++ b/mage_ai/services/compute/aws/steps.py @@ -1,6 +1,5 @@ import os import socket -from enum import Enum from typing import List import requests @@ -26,6 +25,7 @@ SetupStepStatus, ) from mage_ai.services.ssh.aws.emr.constants import SSH_DEFAULTS +from mage_ai.shared.enum import StrEnum from mage_ai.shared.hash import extract, merge_dict ERROR_MESSAGE_ACCESS_KEY_ID = ErrorMessage.load( @@ -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 f5d48421ffab..affeed2b65d5 100644 --- a/mage_ai/services/compute/constants.py +++ b/mage_ai/services/compute/constants.py @@ -1,20 +1,20 @@ -from enum import Enum +from mage_ai.shared.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 f3b7a659643b..0ce28a2ec11f 100644 --- a/mage_ai/services/compute/models.py +++ b/mage_ai/services/compute/models.py @@ -1,6 +1,5 @@ from abc import abstractmethod from dataclasses import dataclass, field -from enum import Enum from typing import Dict, List from mage_ai.data_preparation.models.project import Project @@ -10,6 +9,7 @@ ComputeManagementApplicationTab, ) from mage_ai.services.spark.constants import ComputeServiceUUID +from mage_ai.shared.enum import StrEnum from mage_ai.shared.hash import merge_dict from mage_ai.shared.models import BaseDataClass @@ -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 17a04dc7a25c..1874ae5e9dc4 100644 --- a/mage_ai/services/spark/constants.py +++ b/mage_ai/services/spark/constants.py @@ -1,12 +1,12 @@ -from enum import Enum +from mage_ai.shared.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 9db7e867ca6c..1f0c8b4bc07c 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 typing import Dict, List from mage_ai.services.spark.models.applications import Application from mage_ai.services.spark.models.base import BaseSparkModel +from mage_ai.shared.enum import StrEnum -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 5ba8ac72ccb8..2d07cc66d178 100644 --- a/mage_ai/services/spark/models/sqls.py +++ b/mage_ai/services/spark/models/sqls.py @@ -1,14 +1,14 @@ from dataclasses import dataclass, field -from enum import Enum from typing import List from mage_ai.services.spark.models.applications import Application from mage_ai.services.spark.models.base import BaseSparkModel from mage_ai.services.spark.models.jobs import Job from mage_ai.services.spark.models.stages import StageAttempt +from mage_ai.shared.enum import StrEnum -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 2572d71503af..d69bbae3b12a 100644 --- a/mage_ai/services/spark/models/stages.py +++ b/mage_ai/services/spark/models/stages.py @@ -1,25 +1,25 @@ from dataclasses import dataclass, field -from enum import Enum from typing import Dict, List from mage_ai.services.spark.models.applications import Application from mage_ai.services.spark.models.base import BaseSparkModel from mage_ai.services.spark.models.metrics import Metrics +from mage_ai.shared.enum import StrEnum -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 30101741f6fd..da575e41b846 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 typing import List from mage_ai.services.spark.models.base import BaseSparkModel +from mage_ai.shared.enum import StrEnum -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 72b96ef8c3df..cef754ea6930 100644 --- a/mage_ai/settings/backends.py +++ b/mage_ai/settings/backends.py @@ -1,13 +1,14 @@ import base64 import logging import os -from enum import Enum from typing import Optional +from mage_ai.shared.enum import StrEnum + 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 0e0add88762c..c0f227ee0f98 100644 --- a/mage_ai/settings/models/configuration_option.py +++ b/mage_ai/settings/models/configuration_option.py @@ -1,7 +1,6 @@ import asyncio import os from dataclasses import dataclass -from enum import Enum from pathlib import Path from typing import Dict, Union @@ -15,6 +14,7 @@ from mage_ai.data_preparation.shared.utils import get_template_vars from mage_ai.settings.platform import project_platform_activated from mage_ai.settings.repo import get_repo_path +from mage_ai.shared.enum import StrEnum from mage_ai.shared.files import get_full_file_paths_containing_multi_items from mage_ai.shared.hash import merge_dict from mage_ai.shared.models import BaseDataClass @@ -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 48f5a0b283db..21d6d69ed431 100644 --- a/mage_ai/shared/constants.py +++ b/mage_ai/shared/constants.py @@ -1,4 +1,4 @@ -from enum import Enum +from mage_ai.shared.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 729b35be483c..4dbcdd893172 100644 --- a/mage_ai/shared/custom_logger.py +++ b/mage_ai/shared/custom_logger.py @@ -1,15 +1,15 @@ import inspect import logging -from enum import Enum import simplejson from mage_ai.data_preparation.models.constants import BlockType +from mage_ai.shared.enum import StrEnum from mage_ai.shared.environments import is_deus_ex_machina 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/enum.py b/mage_ai/shared/enum.py new file mode 100644 index 000000000000..3dee0e9194b1 --- /dev/null +++ b/mage_ai/shared/enum.py @@ -0,0 +1,12 @@ +try: + from enum import IntEnum, StrEnum +except ImportError: + from enum import Enum, IntEnum + + class StrEnum(str, Enum): + pass + +__all__ = [ + "IntEnum", + "StrEnum", +] diff --git a/mage_ai/shared/logger.py b/mage_ai/shared/logger.py index 708be0de9eba..773c59bfd04e 100644 --- a/mage_ai/shared/logger.py +++ b/mage_ai/shared/logger.py @@ -2,10 +2,10 @@ import logging import time from contextlib import contextmanager, redirect_stdout -from enum import Enum from typing import Callable, List from mage_ai.settings import SERVER_LOGGING_TEMPLATE +from mage_ai.shared.enum import StrEnum from mage_ai.shared.hash import merge_dict logger = logging.getLogger(__name__) @@ -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 e5088013d691..b4bb38a79c67 100644 --- a/mage_ai/streaming/constants.py +++ b/mage_ai/streaming/constants.py @@ -1,10 +1,10 @@ -from enum import Enum +from mage_ai.shared.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/generic_io.py b/mage_ai/streaming/sinks/generic_io.py index 8ee631db994e..72258da2ac6f 100644 --- a/mage_ai/streaming/sinks/generic_io.py +++ b/mage_ai/streaming/sinks/generic_io.py @@ -45,9 +45,13 @@ class GenericIOConfig(BaseConfig): class GenericIOSink(BaseSink): config_class = GenericIOConfig + def __init__(self, config: Dict, config_loader_class=None, **kwargs): + self.config_loader_class = config_loader_class or ConfigFileLoader + super().__init__(config, **kwargs) + def init_client(self): config_path = os.path.join(get_repo_path(), 'io_config.yaml') - config_file_loader = ConfigFileLoader(config_path, self.config.profile) + config_file_loader = self.config_loader_class(config_path, self.config.profile) self.io_client = self.__io_class().with_config(config_file_loader) if hasattr(self.io_client, 'open') and callable(self.io_client.open): diff --git a/mage_ai/streaming/sinks/kafka.py b/mage_ai/streaming/sinks/kafka.py index 99c372d56a10..21fb500a6016 100644 --- a/mage_ai/streaming/sinks/kafka.py +++ b/mage_ai/streaming/sinks/kafka.py @@ -1,17 +1,17 @@ import json import time from dataclasses import dataclass -from enum import Enum from typing import Dict, List from kafka import KafkaProducer from mage_ai.shared.config import BaseConfig +from mage_ai.shared.enum import StrEnum from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_MS 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 846f62f5007f..893f1e840233 100644 --- a/mage_ai/streaming/sources/amazon_sqs.py +++ b/mage_ai/streaming/sources/amazon_sqs.py @@ -2,12 +2,12 @@ import threading import time from dataclasses import dataclass -from enum import Enum from typing import Callable, Dict import boto3 from mage_ai.shared.config import BaseConfig +from mage_ai.shared.enum import StrEnum from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE from mage_ai.streaming.sources.base import BaseSource from mage_ai.streaming.sources.shared import SerDeConfig, SerializationMethod @@ -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 a27d97406a18..f9bb69dce1b1 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 typing import Callable, Dict +from mage_ai.shared.enum import StrEnum 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 1198a02371b9..05ef1a91267c 100644 --- a/mage_ai/streaming/sources/kafka.py +++ b/mage_ai/streaming/sources/kafka.py @@ -2,18 +2,18 @@ import json import time from dataclasses import dataclass, field -from enum import Enum from typing import Callable, Dict, List from kafka import KafkaConsumer, TopicPartition from mage_ai.shared.config import BaseConfig +from mage_ai.shared.enum import StrEnum from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_MS from mage_ai.streaming.sources.base import BaseSource 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 90445e1c6e07..32ca4240dc55 100644 --- a/mage_ai/streaming/sources/nats_js.py +++ b/mage_ai/streaming/sources/nats_js.py @@ -1,5 +1,4 @@ import asyncio -import enum import json import ssl import threading @@ -10,6 +9,7 @@ from nats.errors import NoServersError, TimeoutError from mage_ai.shared.config import BaseConfig +from mage_ai.shared.enum import StrEnum from mage_ai.streaming.constants import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_MS from mage_ai.streaming.sources.base import BaseSource, SourceConsumeMethod @@ -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 1ace16e2ca21..9373b33a49a1 100644 --- a/mage_ai/streaming/sources/shared.py +++ b/mage_ai/streaming/sources/shared.py @@ -1,8 +1,9 @@ from dataclasses import dataclass -from enum import Enum +from mage_ai.shared.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 0767ff0cf49b..c1836ae183e3 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,4 @@ import secrets -from enum import Enum from typing import Any, Callable, Dict, List, Union from unittest.mock import patch @@ -19,13 +18,14 @@ ) from mage_ai.orchestration.db.models.oauth import Permission, Role, User, UserRole from mage_ai.shared.array import find +from mage_ai.shared.enum import StrEnum from mage_ai.shared.hash import index_by, merge_dict from mage_ai.tests.api.mixins import BootstrapMixin from mage_ai.tests.api.operations.test_base import BaseApiTestCase 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/tests/data_preparation/models/test_block.py b/mage_ai/tests/data_preparation/models/test_block.py index f2dfa4cde04e..8e15e3e1b935 100644 --- a/mage_ai/tests/data_preparation/models/test_block.py +++ b/mage_ai/tests/data_preparation/models/test_block.py @@ -1,8 +1,8 @@ +import asyncio import os from unittest.mock import patch import pandas as pd -from async_timeout import asyncio from faker import Faker from pandas.testing import assert_frame_equal diff --git a/mage_ai/tests/streaming/sinks/test_generic_io.py b/mage_ai/tests/streaming/sinks/test_generic_io.py index 160840545570..895fc40a6504 100644 --- a/mage_ai/tests/streaming/sinks/test_generic_io.py +++ b/mage_ai/tests/streaming/sinks/test_generic_io.py @@ -4,7 +4,6 @@ import pandas as pd from pandas.testing import assert_frame_equal -from mage_ai.streaming.sinks.generic_io import GenericIOSink from mage_ai.tests.base_test import TestCase TEST_DATABASES = [ @@ -52,12 +51,12 @@ class GenericIOTests(TestCase): - @patch('mage_ai.streaming.sinks.generic_io.ConfigFileLoader') @patch('mage_ai.streaming.sinks.generic_io.importlib.import_module') - def test_init_client(self, mock_import_module, mock_config_loader): + def test_init_client(self, mock_import_module): first_db = True for database in TEST_DATABASES: - mock_objects = self.__mock_objects(mock_import_module, mock_config_loader, database) + mock_objects = self.__mock_objects(mock_import_module, database) + mock_config_loader = mock_objects['config_loader'] # Assertions if first_db: @@ -74,11 +73,10 @@ def test_init_client(self, mock_import_module, mock_config_loader): mock_objects['io_client'].close.assert_called_once() first_db = False - @patch('mage_ai.streaming.sinks.generic_io.ConfigFileLoader') @patch('mage_ai.streaming.sinks.generic_io.importlib.import_module') - def test_write(self, mock_import_module, mock_config_loader): + def test_write(self, mock_import_module): database = TEST_DATABASES[0] - mock_objects = self.__mock_objects(mock_import_module, mock_config_loader, database) + mock_objects = self.__mock_objects(mock_import_module, database) mock_objects['sink'].batch_write = MagicMock() @@ -90,11 +88,10 @@ def test_write(self, mock_import_module, mock_config_loader): # Assertions mock_objects['sink'].batch_write.assert_called_once_with([message]) - @patch('mage_ai.streaming.sinks.generic_io.ConfigFileLoader') @patch('mage_ai.streaming.sinks.generic_io.importlib.import_module') - def test_batch_write(self, mock_import_module, mock_config_loader): + def test_batch_write(self, mock_import_module): database = TEST_DATABASES[0] - mock_objects = self.__mock_objects(mock_import_module, mock_config_loader, database) + mock_objects = self.__mock_objects(mock_import_module, database) mock_objects['io_client'].export = MagicMock() # Test data @@ -130,23 +127,34 @@ def test_batch_write(self, mock_import_module, mock_config_loader): ]), ) - def __mock_objects(self, mock_import_module, mock_config_loader, database): + def __mock_objects(self, mock_import_module, database): + from mage_ai.streaming.sinks.generic_io import GenericIOSink + mock_io_module = MagicMock() mock_io_class = MagicMock() mock_io_client = MagicMock() mock_import_module.return_value = mock_io_module mock_io_class.with_config.return_value = mock_io_client + + mock_config_loader = MagicMock() mock_config_loader_instance = mock_config_loader.return_value + setattr(mock_io_module, database['class_name'], mock_io_class) - generic_io_sink = GenericIOSink(dict( - connector_type=database['connector_type'], - profile='test_profile', + + generic_io_sink = GenericIOSink( config=dict( - table_name='test_table', - ) - )) + connector_type=database['connector_type'], + profile='test_profile', + config=dict( + table_name='test_table', + ) + ), + config_loader_class=mock_config_loader, + ) + return dict( config_loader_instance=mock_config_loader_instance, + config_loader=mock_config_loader, io_module=mock_io_module, io_class=mock_io_class, io_client=mock_io_client, diff --git a/mage_ai/usage_statistics/constants.py b/mage_ai/usage_statistics/constants.py index b393fd31ffff..2128cb9c2c43 100644 --- a/mage_ai/usage_statistics/constants.py +++ b/mage_ai/usage_statistics/constants.py @@ -1,9 +1,9 @@ -import enum +from mage_ai.shared.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 28cfb1d3564f..abeed9ca069f 100644 --- a/mage_integrations/mage_integrations/connections/mysql/__init__.py +++ b/mage_integrations/mage_integrations/connections/mysql/__init__.py @@ -1,4 +1,3 @@ -import enum import io import os from typing import Dict @@ -8,9 +7,10 @@ from sshtunnel import SSHTunnelForwarder from mage_integrations.connections.sql.base import Connection +from mage_integrations.utils.enum import StrEnum -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 eb913bc16d66..6012149c6dc9 100644 --- a/mage_integrations/mage_integrations/sources/couchbase/__init__.py +++ b/mage_integrations/mage_integrations/sources/couchbase/__init__.py @@ -1,7 +1,8 @@ -from enum import Enum -from mage_integrations.connections.couchbase import ( - Couchbase as CouchbaseConnection -) +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 +15,12 @@ wrap_column_in_quotes, ) from mage_integrations.sources.sql.base import Source -from singer import catalog -from typing import Any, Dict, List +from mage_integrations.utils.enum import StrEnum 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 45c7004c8915..b929b5b38592 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 mage_integrations.utils.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/mage_integrations/utils/enum.py b/mage_integrations/mage_integrations/utils/enum.py new file mode 100644 index 000000000000..3dee0e9194b1 --- /dev/null +++ b/mage_integrations/mage_integrations/utils/enum.py @@ -0,0 +1,12 @@ +try: + from enum import IntEnum, StrEnum +except ImportError: + from enum import Enum, IntEnum + + class StrEnum(str, Enum): + pass + +__all__ = [ + "IntEnum", + "StrEnum", +] diff --git a/mage_integrations/requirements.txt b/mage_integrations/requirements.txt index 4fc8b8034f83..0f7383bd88ae 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/poetry.lock b/poetry.lock index 8f2758aac5d6..2853677d149c 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1,10 +1,9 @@ -# This file is automatically @generated by Poetry 1.4.2 and should not be changed by hand. +# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand. [[package]] name = "appnope" version = "0.1.3" description = "Disable App Nap on macOS >= 10.9" -category = "dev" optional = false python-versions = "*" files = [ @@ -16,7 +15,6 @@ files = [ name = "argcomplete" version = "3.0.8" description = "Bash tab completion for argparse" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -31,7 +29,6 @@ test = ["coverage", "mypy", "pexpect", "ruff", "wheel"] name = "asttokens" version = "2.2.1" description = "Annotate AST trees with source code positions" -category = "dev" optional = false python-versions = "*" files = [ @@ -49,7 +46,6 @@ test = ["astroid", "pytest"] name = "attrs" version = "23.1.0" description = "Classes Without Boilerplate" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -68,7 +64,6 @@ tests-no-zope = ["cloudpickle", "hypothesis", "mypy (>=1.1.1)", "pympler", "pyte name = "backcall" version = "0.2.0" description = "Specifications for callback functions passed in to an API" -category = "dev" optional = false python-versions = "*" files = [ @@ -80,7 +75,6 @@ files = [ name = "black" version = "23.3.0" description = "The uncompromising code formatter." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -130,7 +124,6 @@ uvloop = ["uvloop (>=0.15.2)"] name = "cffi" version = "1.15.1" description = "Foreign Function Interface for Python calling C code." -category = "dev" optional = false python-versions = "*" files = [ @@ -207,7 +200,6 @@ pycparser = "*" name = "cfgv" version = "3.3.1" description = "Validate configuration and produce human readable error messages." -category = "dev" optional = false python-versions = ">=3.6.1" files = [ @@ -219,7 +211,6 @@ files = [ name = "charset-normalizer" version = "3.1.0" description = "The Real First Universal Charset Detector. Open, modern and actively maintained alternative to Chardet." -category = "dev" optional = false python-versions = ">=3.7.0" files = [ @@ -304,7 +295,6 @@ files = [ name = "click" version = "8.1.3" description = "Composable command line interface toolkit" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -319,7 +309,6 @@ colorama = {version = "*", markers = "platform_system == \"Windows\""} name = "colorama" version = "0.4.6" description = "Cross-platform colored terminal text." -category = "dev" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*,>=2.7" files = [ @@ -331,7 +320,6 @@ files = [ name = "comm" version = "0.1.3" description = "Jupyter Python Comm implementation, for usage in ipykernel, xeus-python etc." -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -351,7 +339,6 @@ typing = ["mypy (>=0.990)"] name = "commitizen" version = "3.2.1" description = "Python commitizen client tool" -category = "dev" optional = false python-versions = ">=3.7,<4.0" files = [ @@ -376,7 +363,6 @@ tomlkit = ">=0.5.3,<1.0.0" name = "coverage" version = "7.2.5" description = "Code coverage measurement for Python" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -443,7 +429,6 @@ toml = ["tomli"] name = "debugpy" version = "1.6.7" description = "An implementation of the Debug Adapter Protocol for Python" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -471,7 +456,6 @@ files = [ name = "decli" version = "0.6.0" description = "Minimal, easy-to-use, declarative cli tool" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -483,7 +467,6 @@ files = [ name = "decorator" version = "5.1.1" description = "Decorators for Humans" -category = "dev" optional = false python-versions = ">=3.5" files = [ @@ -495,7 +478,6 @@ files = [ name = "distlib" version = "0.3.6" description = "Distribution utilities" -category = "dev" optional = false python-versions = "*" files = [ @@ -507,7 +489,6 @@ files = [ name = "exceptiongroup" version = "1.1.1" description = "Backport of PEP 654 (exception groups)" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -522,7 +503,6 @@ test = ["pytest (>=6)"] name = "executing" version = "1.2.0" description = "Get the currently executing AST node of a frame, and other information" -category = "dev" optional = false python-versions = "*" files = [ @@ -537,7 +517,6 @@ tests = ["asttokens", "littleutils", "pytest", "rich"] name = "filelock" version = "3.12.0" description = "A platform independent file lock." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -553,7 +532,6 @@ testing = ["covdefaults (>=2.3)", "coverage (>=7.2.3)", "diff-cover (>=7.5)", "p name = "flake8" version = "5.0.4" description = "the modular source code checker: pep8 pyflakes and co" -category = "dev" optional = false python-versions = ">=3.6.1" files = [ @@ -570,7 +548,6 @@ pyflakes = ">=2.5.0,<2.6.0" name = "flake8-bugbear" version = "23.3.12" description = "A plugin for flake8 finding likely bugs and design problems in your program. Contains warnings that don't belong in pyflakes and pycodestyle." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -589,7 +566,6 @@ dev = ["coverage", "hypothesis", "hypothesmith (>=0.2)", "pre-commit", "pytest", name = "identify" version = "2.5.24" description = "File identification library for Python" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -604,7 +580,6 @@ license = ["ukkonen"] name = "importlib-metadata" version = "6.6.0" description = "Read metadata from Python packages" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -624,7 +599,6 @@ testing = ["flake8 (<5)", "flufl.flake8", "importlib-resources (>=1.3)", "packag name = "iniconfig" version = "2.0.0" description = "brain-dead simple config-ini parsing" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -636,7 +610,6 @@ files = [ name = "interrogate" version = "1.5.0" description = "Interrogate a codebase for docstring coverage." -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -662,7 +635,6 @@ tests = ["pytest", "pytest-cov", "pytest-mock"] name = "ipykernel" version = "6.23.0" description = "IPython Kernel for Jupyter" -category = "dev" optional = false python-versions = ">=3.8" files = [ @@ -676,7 +648,7 @@ comm = ">=0.1.1" debugpy = ">=1.6.5" ipython = ">=7.23.1" jupyter-client = ">=6.1.12" -jupyter-core = ">=4.12,<5.0.0 || >=5.1.0" +jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0" matplotlib-inline = ">=0.1" nest-asyncio = "*" packaging = "*" @@ -696,7 +668,6 @@ test = ["flaky", "ipyparallel", "pre-commit", "pytest (>=7.0)", "pytest-asyncio" name = "ipython" version = "8.12.2" description = "IPython: Productive Interactive Computing" -category = "dev" optional = false python-versions = ">=3.8" files = [ @@ -736,7 +707,6 @@ test-extra = ["curio", "matplotlib (!=3.2.0)", "nbformat", "numpy (>=1.21)", "pa name = "isort" version = "5.12.0" description = "A Python utility / library to sort Python imports." -category = "dev" optional = false python-versions = ">=3.8.0" files = [ @@ -754,7 +724,6 @@ requirements-deprecated-finder = ["pip-api", "pipreqs"] name = "jedi" version = "0.18.2" description = "An autocompletion tool for Python that can be used for text editors." -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -774,7 +743,6 @@ testing = ["Django (<3.1)", "attrs", "colorama", "docopt", "pytest (<7.0.0)"] name = "jinja2" version = "3.1.2" description = "A very fast and expressive template engine." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -792,7 +760,6 @@ i18n = ["Babel (>=2.7)"] name = "jupyter-client" version = "8.2.0" description = "Jupyter protocol implementation and client libraries" -category = "dev" optional = false python-versions = ">=3.8" files = [ @@ -802,7 +769,7 @@ files = [ [package.dependencies] importlib-metadata = {version = ">=4.8.3", markers = "python_version < \"3.10\""} -jupyter-core = ">=4.12,<5.0.0 || >=5.1.0" +jupyter-core = ">=4.12,<5.0.dev0 || >=5.1.dev0" python-dateutil = ">=2.8.2" pyzmq = ">=23.0" tornado = ">=6.2" @@ -816,7 +783,6 @@ test = ["coverage", "ipykernel (>=6.14)", "mypy", "paramiko", "pre-commit", "pyt name = "jupyter-core" version = "5.3.0" description = "Jupyter core package. A base package on which Jupyter projects rely." -category = "dev" optional = false python-versions = ">=3.8" files = [ @@ -837,7 +803,6 @@ test = ["ipykernel", "pre-commit", "pytest", "pytest-cov", "pytest-timeout"] name = "markupsafe" version = "2.1.2" description = "Safely add untrusted strings to HTML/XML markup." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -897,7 +862,6 @@ files = [ name = "matplotlib-inline" version = "0.1.6" description = "Inline Matplotlib backend for Jupyter" -category = "dev" optional = false python-versions = ">=3.5" files = [ @@ -912,7 +876,6 @@ traitlets = "*" name = "mccabe" version = "0.7.0" description = "McCabe checker, plugin for flake8" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -924,7 +887,6 @@ files = [ name = "mypy" version = "1.2.0" description = "Optional static typing for Python" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -971,7 +933,6 @@ reports = ["lxml"] name = "mypy-extensions" version = "1.0.0" description = "Type system extensions for programs checked with the mypy type checker." -category = "dev" optional = false python-versions = ">=3.5" files = [ @@ -983,7 +944,6 @@ files = [ name = "nest-asyncio" version = "1.5.6" description = "Patch asyncio to allow nested event loops" -category = "dev" optional = false python-versions = ">=3.5" files = [ @@ -995,7 +955,6 @@ files = [ name = "nodeenv" version = "1.7.0" description = "Node.js virtual environment builder" -category = "dev" optional = false python-versions = ">=2.7,!=3.0.*,!=3.1.*,!=3.2.*,!=3.3.*,!=3.4.*,!=3.5.*,!=3.6.*" files = [ @@ -1010,7 +969,6 @@ setuptools = "*" name = "packaging" version = "23.1" description = "Core utilities for Python packages" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1022,7 +980,6 @@ files = [ name = "pandas-stubs" version = "2.0.1.230501" description = "Type annotations for pandas" -category = "dev" optional = false python-versions = ">=3.8" files = [ @@ -1037,7 +994,6 @@ types-pytz = ">=2022.1.1" name = "parso" version = "0.8.3" description = "A Python Parser" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -1053,7 +1009,6 @@ testing = ["docopt", "pytest (<6.0.0)"] name = "pathspec" version = "0.11.1" description = "Utility library for gitignore style pattern matching of file paths." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1065,7 +1020,6 @@ files = [ name = "pexpect" version = "4.8.0" description = "Pexpect allows easy control of interactive console applications." -category = "dev" optional = false python-versions = "*" files = [ @@ -1080,7 +1034,6 @@ ptyprocess = ">=0.5" name = "pickleshare" version = "0.7.5" description = "Tiny 'shelve'-like database with concurrency support" -category = "dev" optional = false python-versions = "*" files = [ @@ -1092,7 +1045,6 @@ files = [ name = "platformdirs" version = "3.5.0" description = "A small Python package for determining appropriate platform-specific dirs, e.g. a \"user data dir\"." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1108,7 +1060,6 @@ test = ["appdirs (==1.4.4)", "covdefaults (>=2.3)", "pytest (>=7.3.1)", "pytest- name = "pluggy" version = "1.0.0" description = "plugin and hook calling mechanisms for python" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -1124,7 +1075,6 @@ testing = ["pytest", "pytest-benchmark"] name = "pre-commit" version = "3.3.1" description = "A framework for managing and maintaining multi-language pre-commit hooks." -category = "dev" optional = false python-versions = ">=3.8" files = [ @@ -1143,7 +1093,6 @@ virtualenv = ">=20.10.0" name = "prompt-toolkit" version = "3.0.38" description = "Library for building powerful interactive command lines in Python" -category = "dev" optional = false python-versions = ">=3.7.0" files = [ @@ -1158,7 +1107,6 @@ wcwidth = "*" name = "psutil" version = "5.9.5" description = "Cross-platform lib for process and system monitoring in Python." -category = "dev" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" files = [ @@ -1185,7 +1133,6 @@ test = ["enum34", "ipaddress", "mock", "pywin32", "wmi"] name = "ptyprocess" version = "0.7.0" description = "Run a subprocess in a pseudo terminal" -category = "dev" optional = false python-versions = "*" files = [ @@ -1197,7 +1144,6 @@ files = [ name = "pure-eval" version = "0.2.2" description = "Safely evaluate AST nodes without side effects" -category = "dev" optional = false python-versions = "*" files = [ @@ -1212,7 +1158,6 @@ tests = ["pytest"] name = "py" version = "1.11.0" description = "library with cross-python path, ini-parsing, io, code, log facilities" -category = "dev" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*, !=3.4.*" files = [ @@ -1224,7 +1169,6 @@ files = [ name = "pycodestyle" version = "2.9.1" description = "Python style guide checker" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -1236,7 +1180,6 @@ files = [ name = "pycparser" version = "2.21" description = "C parser in Python" -category = "dev" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" files = [ @@ -1248,7 +1191,6 @@ files = [ name = "pyflakes" version = "2.5.0" description = "passive checker of Python programs" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -1260,7 +1202,6 @@ files = [ name = "pygments" version = "2.15.1" description = "Pygments is a syntax highlighting package written in Python." -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1275,7 +1216,6 @@ plugins = ["importlib-metadata"] name = "pytest" version = "7.3.1" description = "pytest: simple powerful testing with Python" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1298,7 +1238,6 @@ testing = ["argcomplete", "attrs (>=19.2.0)", "hypothesis (>=3.56)", "mock", "no name = "pytest-cov" version = "4.0.0" description = "Pytest plugin for measuring coverage." -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -1317,7 +1256,6 @@ testing = ["fields", "hunter", "process-tests", "pytest-xdist", "six", "virtuale name = "python-dateutil" version = "2.8.2" description = "Extensions to the standard Python datetime module" -category = "dev" optional = false python-versions = "!=3.0.*,!=3.1.*,!=3.2.*,>=2.7" files = [ @@ -1332,7 +1270,6 @@ six = ">=1.5" name = "pywin32" version = "306" description = "Python for Window Extensions" -category = "dev" optional = false python-versions = "*" files = [ @@ -1356,7 +1293,6 @@ files = [ name = "pyyaml" version = "6.0" description = "YAML parser and emitter for Python" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -1406,7 +1342,6 @@ files = [ name = "pyzmq" version = "25.0.2" description = "Python bindings for 0MQ" -category = "dev" optional = false python-versions = ">=3.6" files = [ @@ -1496,7 +1431,6 @@ cffi = {version = "*", markers = "implementation_name == \"pypy\""} name = "questionary" version = "1.10.0" description = "Python library to build pretty command line user prompts ⭐️" -category = "dev" optional = false python-versions = ">=3.6,<4.0" files = [ @@ -1514,7 +1448,6 @@ docs = ["Sphinx (>=3.3,<4.0)", "sphinx-autobuild (>=2020.9.1,<2021.0.0)", "sphin name = "setuptools" version = "67.7.2" description = "Easily download, build, install, upgrade, and uninstall Python packages" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1531,7 +1464,6 @@ testing-integration = ["build[virtualenv]", "filelock (>=3.4.0)", "jaraco.envs ( name = "six" version = "1.16.0" description = "Python 2 and 3 compatibility utilities" -category = "dev" optional = false python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*" files = [ @@ -1543,7 +1475,6 @@ files = [ name = "sourcery" version = "1.2.0" description = "Magically refactor Python" -category = "dev" optional = false python-versions = "*" files = [ @@ -1556,7 +1487,6 @@ files = [ name = "stack-data" version = "0.6.2" description = "Extract data from python stack frames and tracebacks for informative displays" -category = "dev" optional = false python-versions = "*" files = [ @@ -1576,7 +1506,6 @@ tests = ["cython", "littleutils", "pygments", "pytest", "typeguard"] name = "tabulate" version = "0.9.0" description = "Pretty-print tabular data" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1591,7 +1520,6 @@ widechars = ["wcwidth"] name = "termcolor" version = "2.3.0" description = "ANSI color formatting for output in terminal" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1606,7 +1534,6 @@ tests = ["pytest", "pytest-cov"] name = "toml" version = "0.10.2" description = "Python Library for Tom's Obvious, Minimal Language" -category = "dev" optional = false python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*" files = [ @@ -1618,7 +1545,6 @@ files = [ name = "tomli" version = "2.0.1" description = "A lil' TOML parser" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1630,7 +1556,6 @@ files = [ name = "tomlkit" version = "0.11.8" description = "Style preserving TOML library" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1642,7 +1567,6 @@ files = [ name = "tornado" version = "6.3.1" description = "Tornado is a Python web framework and asynchronous networking library, originally developed at FriendFeed." -category = "dev" optional = false python-versions = ">= 3.8" files = [ @@ -1663,7 +1587,6 @@ files = [ name = "traitlets" version = "5.9.0" description = "Traitlets Python configuration system" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1679,7 +1602,6 @@ test = ["argcomplete (>=2.0)", "pre-commit", "pytest", "pytest-mock"] name = "types-pytz" version = "2023.3.0.0" description = "Typing stubs for pytz" -category = "dev" optional = false python-versions = "*" files = [ @@ -1689,21 +1611,19 @@ files = [ [[package]] name = "typing-extensions" -version = "4.5.0" -description = "Backported and Experimental Type Hints for Python 3.7+" -category = "dev" +version = "4.11.0" +description = "Backported and Experimental Type Hints for Python 3.8+" optional = false -python-versions = ">=3.7" +python-versions = ">=3.8" files = [ - {file = "typing_extensions-4.5.0-py3-none-any.whl", hash = "sha256:fb33085c39dd998ac16d1431ebc293a8b3eedd00fd4a32de0ff79002c19511b4"}, - {file = "typing_extensions-4.5.0.tar.gz", hash = "sha256:5cb5f4a79139d699607b3ef622a1dedafa84e115ab0024e0d9c044a9479ca7cb"}, + {file = "typing_extensions-4.11.0-py3-none-any.whl", hash = "sha256:c1f94d72897edaf4ce775bb7558d5b79d8126906a14ea5ed1635921406c0387a"}, + {file = "typing_extensions-4.11.0.tar.gz", hash = "sha256:83f085bd5ca59c80295fc2a82ab5dac679cbe02b9f33f7d83af68e241bea51b0"}, ] [[package]] name = "virtualenv" version = "20.23.0" description = "Virtual Python Environment builder" -category = "dev" optional = false python-versions = ">=3.7" files = [ @@ -1724,7 +1644,6 @@ test = ["covdefaults (>=2.3)", "coverage (>=7.2.3)", "coverage-enable-subprocess name = "wcwidth" version = "0.2.6" description = "Measures the displayed width of unicode strings in a terminal" -category = "dev" optional = false python-versions = "*" files = [ @@ -1736,7 +1655,6 @@ files = [ name = "zipp" version = "3.15.0" description = "Backport of pathlib-compatible object wrapper for zip files" -category = "dev" optional = false python-versions = ">=3.7" files = [ diff --git a/requirements.txt b/requirements.txt index 09da5978908b..f67b2b287b97 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,7 +44,7 @@ terminado==0.17.1 thefuzz[speedup]==0.19.0 tornado==6.3.3 typer[all]==0.9.0 -typing_extensions==4.10.0 +typing_extensions==4.11.0 watchdog==4.0.0 # extras @@ -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 diff --git a/setup.py b/setup.py index 8d92fa0440f5..aed3da152132 100644 --- a/setup.py +++ b/setup.py @@ -33,7 +33,7 @@ def readme(): 'Operating System :: OS Independent', ], install_requires=requirements, - python_requires='>=3.7', + python_requires='>=3.8', entry_points={ 'console_scripts': [ 'mage=mage_ai.cli.main:app',