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

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 @@ SKIP_LIMIT_ADAPTER_NAMES = [
 ]
 
 
-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'

mage_ai/data_preparation/models/block/dynamic/factory.py

@@ -39,6 +39,26 @@ def calculate_item_count(
     check_upstream_block_uuids_only: Optional[List[str]] = None,
     update_upstream_item_count_callback: Optional[Callable] = None,
 ) -> int:
+    """
+    Calculates the item count for a given block based on its upstream blocks' output.
+
+    This method builds counters for the block and multiplies their item counts to determine
+    the total item count. If a counter has an item count of zero and is included in
+    `check_upstream_block_uuids_only`, the method recursively calculates the item count
+    for the corresponding upstream block. An optional callback can be provided to update
+    the item count based on upstream block outputs.
+
+    Args:
+        block: The block for which to calculate the item count.
+        execution_partition (Optional[str]): The execution partition identifier.
+        check_upstream_block_uuids_only (Optional[List[str]]): A list of UUIDs for upstream blocks
+            to check when their item count is zero.
+        update_upstream_item_count_callback (Optional[Callable]): A callback function that can
+            update the item count based on the upstream block's output.
+
+    Returns:
+        int: The calculated item count for the block.
+    """
     counters = build_counters(block, execution_partition)
 
     uuid_counts = [
@@ -95,7 +115,9 @@ class DynamicBlockFactory(DynamicBlockWrapperBase):
         ])
 
         # If item count is 0 because the upstream blocks haven’t output anything yet
-        item_count_total = max(self.__calculate_item_count(), upstream_blocks_count)
+        item_count_total = self.__calculate_item_count()
+        if item_count_total == 0:
+            item_count_total = upstream_blocks_count
 
         upstream_block_runs = self.__upstream_block_runs()
         upstream_block_runs_completed = len([
@@ -174,6 +196,19 @@ class DynamicBlockFactory(DynamicBlockWrapperBase):
         return self._counters
 
     def __calculate_item_count(self) -> int:
+        """
+        Calculates the item count for the current block, ensuring downstream execution even if an
+        upstream block has no output.
+
+        This method calculates the item count by checking the output of upstream blocks. If an
+        upstream block has been spawned multiple times but has no output, it ensures that the
+        downstream block is still executed by returning an item count of 1. An internal callback
+        updates the item count based on the upstream block's completed runs.
+
+        Returns:
+            int: The calculated item count for the block, ensuring downstream block execution.
+        """
+
         # If an upstream block indeed has no output (e.g. no return statement),
         # we need to ensure that the downstream block is still executed.
 
@@ -226,10 +261,12 @@ class DynamicBlockFactory(DynamicBlockWrapperBase):
                 if completed_count == spawn_count:
                     return 1
 
-            # If no upstream block has been spawned,
+            # If no upstream block has been spawned and the upstream block runs are completed,
             # we need to return 1 so that the downstream block is still executed.
             if spawn_count == 0:
-                return 1
+                upstream_block_runs = self.__upstream_block_runs()
+                if all(b.status == BlockRun.BlockRunStatus.COMPLETED for b in upstream_block_runs):
+                    return 1
 
             return output_item_count
 

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, Optional, Tuple, Union
 
 import pandas as pd
@@ -38,11 +37,12 @@ from mage_ai.data_preparation.models.utils import (
 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'
@@ -1012,42 +1012,40 @@ def check_all_dynamic_upstreams_completed(
         execution_partition=execution_partition,
         pipeline=pipeline,
     ):
-        if not upstream_block.upstream_blocks:
-            return True
-
         from mage_ai.orchestration.db.models.schedules import BlockRun
 
-        upstreams_done = all([
-            check_all_dynamic_upstreams_completed(
-                b,
-                block_runs=block_runs,
-                execution_partition=execution_partition,
-            )
-            for b in upstream_block.upstream_blocks
-        ])
+        if upstream_block.upstream_blocks:
+            upstreams_done = all([
+                check_all_dynamic_upstreams_completed(
+                    b,
+                    block_runs=block_runs,
+                    execution_partition=execution_partition,
+                )
+                for b in upstream_block.upstream_blocks
+            ])
 
-        if not upstreams_done:
-            return False
+            if not upstreams_done:
+                return False
 
-        is_dynamic_child = is_dynamic_block_child(upstream_block)
-        if is_dynamic_child:
-            combos = build_combinations_for_dynamic_child(
-                upstream_block,
-                execution_partition=execution_partition,
-            )
+            is_dynamic_child = is_dynamic_block_child(upstream_block)
+            if is_dynamic_child:
+                combos = build_combinations_for_dynamic_child(
+                    upstream_block,
+                    execution_partition=execution_partition,
+                )
 
-            selected = [
-                br
-                for br in block_runs
-                if pipeline.get_block(
-                    br.block_uuid,
-                ).uuid
-                == upstream_block.uuid
-                and br.status == BlockRun.BlockRunStatus.COMPLETED
-            ]
-
-            if len(list(selected)) < len(combos) + 1:
-                return False
+                selected = [
+                    br
+                    for br in block_runs
+                    if pipeline.get_block(
+                        br.block_uuid,
+                    ).uuid
+                    == upstream_block.uuid
+                    and br.status == BlockRun.BlockRunStatus.COMPLETED
+                ]
+
+                if len(list(selected)) < len(combos) + 1:
+                    return False
 
         return all([
             br.status == BlockRun.BlockRunStatus.COMPLETED

mage_ai/data_preparation/models/block/dynamic/variables.py

@@ -734,7 +734,7 @@ def fetch_input_variables_for_dynamic_upstream_blocks(
                     )
                     if index is not None:
                         pair = lazy_variable_controller.render(
-                            child_dynamic_block_index=dynamic_block_index,
+                            child_dynamic_block_index=index,
                         )
                         child_data, metadata = pair
 

mage_ai/data_preparation/models/block/settings/dynamic/constants.py

@@ -1,8 +1,8 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
-DEFAULT_STREAM_POLL_INTERVAL = 60
+DEFAULT_STREAM_POLL_INTERVAL = 15
 
 
-class ModeType(str, Enum):
+class ModeType(StrEnum):
     DEFAULT = 'default'
     STREAM = 'stream'

mage_ai/data_preparation/models/block/settings/dynamic/mixins.py

@@ -45,7 +45,12 @@ class DynamicMixin:
 
     @property
     def is_dynamic_parent(self) -> bool:
-        return self.__dynamic_configuration().parent is not None
+        parent = self.__dynamic_configuration().parent
+        if isinstance(parent, bool):
+            return parent
+        elif isinstance(parent, list):
+            return len(parent) > 0
+        return parent is not None
 
     @property
     def is_dynamic_child(self) -> bool:
@@ -57,6 +62,20 @@ class DynamicMixin:
 
     @property
     def is_dynamic_child_streaming(self) -> bool:
+        """
+        Indicates whether current block is dynamic child with stream mode.
+
+        This property checks if the current block is in dynamic v2 mode (`self.is_dynamic_v2`)
+        and whether any of its upstream blocks either:
+        - Should dynamically generate the current block and have dynamic stream mode enabled, or
+        - recursively check whether upstream block is dynamic child block with stream mode.
+
+        If any of these conditions are met for any upstream block, the current block is dynamic
+        child with stream mode.
+
+        Returns:
+            bool: `True` if current block is dynamic child with stream mode., `False` otherwise.
+        """
         return self.is_dynamic_v2 and any(
             (
                 upstream_block.should_dynamically_generate_block(self)
@@ -68,10 +87,24 @@ class DynamicMixin:
 
     @property
     def should_reduce_output(self) -> bool:
-        return self.__dynamic_configuration().reduce_output is not None
+        reduce_output = self.__dynamic_configuration().reduce_output
+        if isinstance(reduce_output, bool):
+            return reduce_output
+        elif isinstance(reduce_output, list):
+            return len(reduce_output) > 0
+        return reduce_output is not None
 
     @property
     def is_dynamic_stream_mode_enabled(self) -> bool:
+        """
+        Indicates whether dynamic stream mode is enabled.
+
+        This property checks if the setting "type: stream" is in the modes list.
+        If the settings are present, dynamic stream mode is considered enabled.
+
+        Returns:
+            bool: `True` if dynamic stream mode is enabled, `False` otherwise.
+        """
         return self.settings_for_mode(ModeType.STREAM) is not None
 
     def build_dynamic_uuid(self, index: int) -> str:
@@ -98,6 +131,22 @@ class DynamicMixin:
         return False
 
     def should_dynamically_generate_block(self, block) -> bool:
+        """
+        Determines whether a block should be dynamically generated based on the parent
+        configuration.
+
+        This method checks the `parent` attribute from the dynamic configuration.
+        - If `parent` is `True`, the block should be dynamically generated.
+        - If `parent` is a list, the block should be dynamically generated only if its UUID is
+            present in the `parent` list.
+        - If `parent` is neither `True` nor a list, the block should not be dynamically generated.
+
+        Args:
+            block: The block object that contains the `uuid` attribute to be checked.
+
+        Returns:
+            bool: `True` if the block should be dynamically generated, `False` otherwise.
+        """
         parent = self.__dynamic_configuration().parent
 
         if parent is True:
@@ -111,8 +160,18 @@ class DynamicMixin:
     def __dynamic_configuration(self) -> DynamicConfiguration:
         if self.configuration:
             config = self.configuration.get('dynamic', {})
+            reduce_output = self.configuration.get('reduce_output')
             if config and isinstance(config, dict):
-                return DynamicConfiguration.load(**config)
+                dynamic_configuration = DynamicConfiguration.load(**config)
             elif isinstance(config, DynamicConfiguration):
-                return config
+                dynamic_configuration = config
+            else:
+                dynamic_configuration = DynamicConfiguration()
+
+            # For backward compatible with dynamic block v1
+            if isinstance(config, bool):
+                dynamic_configuration.parent = config
+            if isinstance(reduce_output, bool):
+                dynamic_configuration.reduce_output = reduce_output
+            return dynamic_configuration
         return DynamicConfiguration()

mage_ai/data_preparation/models/block/settings/global_data_products/models.py

@@ -1,11 +1,11 @@
 from dataclasses import dataclass
-from enum import Enum
 from typing import Optional
 
+from mage_ai.shared.enum import StrEnum
 from mage_ai.shared.models import BaseDataClass
 
 
-class GlobalDataProductObjectType(str, Enum):
+class GlobalDataProductObjectType(StrEnum):
     BLOCK = 'block'
     PIPELINE = 'pipeline'
 

mage_ai/data_preparation/models/constants.py

@@ -1,9 +1,9 @@
 import os
-from enum import Enum
 
 from mage_ai.data_preparation.models.variables.constants import (
     VariableAggregateDataType,
 )
+from mage_ai.shared.enum import StrEnum
 
 DATAFRAME_ANALYSIS_KEYS = frozenset([
     VariableAggregateDataType.INSIGHTS.value,
@@ -15,7 +15,7 @@ DATA_INTEGRATION_CATALOG_FILE = 'data_integration_catalog.json'
 DATAFRAME_ANALYSIS_MAX_COLUMNS = 100
 DATAFRAME_ANALYSIS_MAX_ROWS = 100_000
 DATAFRAME_SAMPLE_COUNT = 1000
-DATAFRAME_SAMPLE_COUNT_PREVIEW = 10
+DATAFRAME_SAMPLE_COUNT_PREVIEW = int(os.getenv('DATAFRAME_SAMPLE_COUNT_PREVIEW', 10) or 10)
 DATAFRAME_SAMPLE_MAX_COLUMNS = 1000
 DYNAMIC_CHILD_BLOCK_SAMPLE_COUNT_PREVIEW = 10
 LOGS_SUBDIR = '.logs'
@@ -31,12 +31,12 @@ VARIABLE_DIR = '.variables'
 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'
@@ -44,14 +44,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'
@@ -71,7 +71,7 @@ class BlockType(str, Enum):
     PIPELINE = 'pipeline'
 
 
-class BlockColor(str, Enum):
+class BlockColor(StrEnum):
     BLUE = 'blue'
     GREY = 'grey'
     PINK = 'pink'
@@ -80,12 +80,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'
@@ -100,7 +100,7 @@ class ExecutorType(str, Enum):
         return executor_type.upper() in cls.__members__
 
 
-class PipelineType(str, Enum):
+class PipelineType(StrEnum):
     INTEGRATION = 'integration'
     DATABRICKS = 'databricks'
     PYTHON = 'python'
@@ -109,7 +109,7 @@ class PipelineType(str, Enum):
     EXECUTION_FRAMEWORK = 'execution_framework'
 
 
-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

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 @@ RESTRICTED_RESOURCE_TYPES = [
 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 @@ VALID_KEYS_FOR_INPUT_OUTPUT_DATA_ALL = \
 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'

mage_ai/data_preparation/models/global_hooks/models.py

@@ -2,7 +2,6 @@ import hashlib
 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.platform import (
 )
 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.models import BaseDataClass
 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'
 

mage_ai/data_preparation/models/pipeline.py

@@ -4,6 +4,7 @@ import json
 import os
 import shutil
 import tempfile
+import time
 import zipfile
 from datetime import datetime, timezone
 from io import BytesIO
@@ -2328,6 +2329,10 @@ class Pipeline:
         extension_uuid: str = None,
         widget: bool = False,
     ):
+        # Introduce a small delay to prevent multiple changes from generating
+        # identical timestamps for the pipeline YAML file
+        time.sleep(0.0004)
+
         blocks_current = sorted([b.uuid for b in self.blocks_by_uuid.values()])
 
         if block_uuid is not None:

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'
     AUTOMATIC_KERNEL_CLEANUP = 'automatic_kernel_cleanup'
     CODE_BLOCK_V2 = 'code_block_v2'

mage_ai/data_preparation/models/triggers/__init__.py

@@ -1,4 +1,3 @@
-import enum
 import os
 import traceback
 from dataclasses import dataclass, field
@@ -12,6 +11,7 @@ from mage_ai.data_preparation.models.constants import PIPELINES_FOLDER
 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
 from mage_ai.shared.yaml import load_yaml
@@ -19,12 +19,12 @@ from mage_ai.shared.yaml import load_yaml
 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'
@@ -37,7 +37,7 @@ SCHEDULE_TYPE_TO_LABEL = {
 }
 
 
-class ScheduleInterval(str, enum.Enum):
+class ScheduleInterval(StrEnum):
     ONCE = '@once'
     HOURLY = '@hourly'
     DAILY = '@daily'
@@ -266,6 +266,8 @@ def build_triggers(
 
             # Add flag to settings so frontend can detect triggers with invalid cron expressions
             if not trigger.has_valid_schedule_interval:
+                if not trigger.settings:
+                    trigger.settings = dict()
                 trigger.settings['invalid_schedule_interval'] = True
 
             triggers.append(trigger)

mage_ai/data_preparation/models/variables/constants.py

@@ -1,4 +1,4 @@
-from enum import Enum
+from mage_ai.shared.enum import StrEnum
 
 CONFIG_JSON_FILE = 'config.json'
 DATAFRAME_COLUMN_TYPES_FILE = 'data_column_types.json'
@@ -14,7 +14,7 @@ MEDIA_IMAGE_VISUALIZATION_FILE = 'visualization.png'
 UBJSON_MODEL_FILENAME = 'model.ubj'
 
 
-class VariableType(str, Enum):
+class VariableType(StrEnum):
     CUSTOM_OBJECT = 'custom_object'
     DATAFRAME = 'dataframe'
     DATAFRAME_ANALYSIS = 'dataframe_analysis'
@@ -31,7 +31,7 @@ class VariableType(str, Enum):
     SPARK_DATAFRAME = 'spark_dataframe'
 
 
-class VariableAggregateDataType(str, Enum):
+class VariableAggregateDataType(StrEnum):
     INSIGHTS = 'insights'
     METADATA = 'metadata'
     RESOURCE_USAGE = 'resource_usage'
@@ -41,7 +41,7 @@ class VariableAggregateDataType(str, Enum):
     TYPE = 'type'
 
 
-class VariableAggregateDataTypeFilename(str, Enum):
+class VariableAggregateDataTypeFilename(StrEnum):
     INSIGHTS = f'{VariableAggregateDataType.INSIGHTS}.json'
     METADATA = f'{VariableAggregateDataType.METADATA}.json'
     RESOURCE_USAGE = f'{VariableAggregateDataType.RESOURCE_USAGE}.json'
@@ -67,6 +67,6 @@ JSON_SAMPLE_FILE = VariableAggregateDataTypeFilename.SAMPLE_DATA.value
 METADATA_FILE = VariableAggregateDataTypeFilename.TYPE.value
 
 
-class VariableAggregateSummaryGroupType(str, Enum):
+class VariableAggregateSummaryGroupType(StrEnum):
     DYNAMIC = 'dynamic'
     PARTS = 'parts'

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_ORDER_BY = 'order_by'
@@ -14,7 +14,7 @@ VARIABLE_NAME_X = 'x'
 VARIABLE_NAME_Y = 'y'
 
 
-class AggregationFunction(str, Enum):
+class AggregationFunction(StrEnum):
     AVERAGE = 'average'
     COUNT = 'count'
     COUNT_DISTINCT = 'count_distinct'
@@ -25,7 +25,7 @@ class AggregationFunction(str, Enum):
     SUM = 'sum'
 
 
-class TimeInterval(str, Enum):
+class TimeInterval(StrEnum):
     DAY = 'day'
     HOUR = 'hour'
     MINUTE = 'minute'
@@ -36,7 +36,7 @@ class TimeInterval(str, Enum):
     YEAR = 'year'
 
 
-class ChartType(str, Enum):
+class ChartType(StrEnum):
     BAR_CHART = 'bar chart'
     CUSTOM = 'custom'
     HISTOGRAM = 'histogram'