mage-ai 0.9.32__py3-none-any.whl → 0.9.34__py3-none-any.whl

mage_ai/api/resources/BlockResource.py

@@ -632,8 +632,9 @@ class BlockResource(GenericResource):
 
             block.update_content(content)
 
-        cache = await BlockCache.initialize_cache()
-        cache.add_pipeline(block, pipeline)
+        if pipeline:
+            cache = await BlockCache.initialize_cache()
+            cache.add_pipeline(block, pipeline)
 
         cache_block_action_object = await BlockActionObjectCache.initialize_cache()
         cache_block_action_object.update_block(block)

mage_ai/api/resources/DatabaseResource.py

@@ -67,9 +67,9 @@ class DatabaseResource(BaseResource):
             for _col, _val in merge_dict(query, where).items():
                 filters.append(self.model_class)
             return self.model_class.query.filter(
-                **merge_dict(query, where)).all()
+                **merge_dict(query, where))
         else:
-            return self.model_class.query.filter(**query).all()
+            return self.model_class.query.filter(**query)
 
     @classmethod
     @safe_db_query

mage_ai/api/resources/PermissionResource.py

@@ -1,12 +1,37 @@
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
-from mage_ai.orchestration.db import safe_db_query
-from mage_ai.orchestration.db.models.oauth import Permission, Role
-from mage_ai.shared.hash import merge_dict
+from mage_ai.orchestration.db import db_connection, safe_db_query
+from mage_ai.orchestration.db.models.oauth import Permission, Role, RolePermission
+from mage_ai.shared.hash import ignore_keys, index_by, merge_dict
 
 
 class PermissionResource(DatabaseResource):
     model_class = Permission
 
+    @classmethod
+    @safe_db_query
+    def collection(self, query_arg, meta, user, **kwargs):
+        only_entity_options = query_arg.get('only_entity_options', [False])
+        if only_entity_options:
+            only_entity_options = only_entity_options[0]
+
+        if only_entity_options:
+            return self.build_result_set(
+                [Permission()],
+                user,
+                **kwargs,
+            )
+
+        return super().collection(
+            query_arg,
+            meta,
+            user,
+            **kwargs,
+        ).order_by(
+            Permission.entity_name.asc(),
+            Permission.entity_type.asc(),
+            Permission.entity_id.asc(),
+        )
+
     @classmethod
     @safe_db_query
     def create(self, payload, user, **kwargs):
@@ -17,5 +42,42 @@ class PermissionResource(DatabaseResource):
             user_id=user.id if user else None,
         )), user, **kwargs)
 
+    @safe_db_query
+    def update(self, payload, **kwargs):
+        role_ids = [int(i) for i in payload.get('role_ids') or []]
+        roles_mapping = index_by(lambda x: x.id, self.roles or [])
+
+        role_ids_create = []
+        role_ids_delete = []
+
+        for role_id in role_ids:
+            if role_id not in roles_mapping:
+                role_ids_create.append(role_id)
+
+        for role_id in roles_mapping.keys():
+            if role_id not in role_ids:
+                role_ids_delete.append(role_id)
+
+        if role_ids_create:
+            db_connection.session.bulk_save_objects(
+                [RolePermission(
+                    permission_id=self.model.id,
+                    role_id=role_id,
+                    user_id=self.current_user.id,
+                ) for role_id in role_ids_create],
+                return_defaults=True,
+            )
+
+        if role_ids_delete:
+            delete_statement = RolePermission.__table__.delete().where(
+                RolePermission.permission_id == self.id,
+                RolePermission.role_id.in_(role_ids_delete),
+            )
+            db_connection.session.execute(delete_statement)
+
+        return super().update(ignore_keys(payload, [
+            'role_ids',
+        ]), **kwargs)
+
 
 PermissionResource.register_parent_model('role_id', Role)

mage_ai/api/resources/Resource.py

@@ -1,6 +1,13 @@
 class Resource():
-    def __init__(self, model, current_user, **kwargs):
+    def __init__(
+        self,
+        model,
+        current_user,
+        result_set_from_external=None,
+        **kwargs,
+    ):
         self.current_user = current_user
         self.model = model
         self.model_options = kwargs
         self.result_set_attr = None
+        self.result_set_from_external = result_set_from_external

mage_ai/api/resources/RolePermissionResource.py

@@ -1,9 +1,32 @@
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
+from mage_ai.orchestration.db import safe_db_query
 from mage_ai.orchestration.db.models.oauth import Role, RolePermission
+from mage_ai.shared.hash import ignore_keys, merge_dict
 
 
 class RolePermissionResource(DatabaseResource):
     model_class = RolePermission
 
+    @classmethod
+    @safe_db_query
+    def create(self, payload, user, **kwargs):
+        payload_updated = merge_dict(payload, dict(
+            user_id=user.id if user else None,
+        ))
+
+        resources = []
+        if 'permission_ids' in payload_updated:
+            permission_ids = payload_updated.get('permission_ids') or []
+
+            for permission_id in permission_ids:
+                resources.append(super().create(merge_dict(
+                    ignore_keys(payload_updated, ['permission_ids']),
+                    dict(permission_id=permission_id),
+                ), user, **kwargs))
+
+            return resources[0] if len(resources) >= 1 else self(None, user, **kwargs)
+
+        return super().create(payload_updated, user, **kwargs)
+
 
 RolePermissionResource.register_parent_model('role_id', Role)

mage_ai/api/resources/RoleResource.py

@@ -1,9 +1,9 @@
 from mage_ai.api.resources.DatabaseResource import DatabaseResource
 from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.constants import Entity
-from mage_ai.orchestration.db import safe_db_query
-from mage_ai.orchestration.db.models.oauth import Role
-from mage_ai.shared.hash import merge_dict
+from mage_ai.orchestration.db import db_connection, safe_db_query
+from mage_ai.orchestration.db.models.oauth import Role, RolePermission, UserRole
+from mage_ai.shared.hash import extract, ignore_keys, index_by, merge_dict
 
 
 class RoleResource(DatabaseResource):
@@ -41,6 +41,7 @@ class RoleResource(DatabaseResource):
             roles = []
             for permission in permissions:
                 roles.append(permission.role)
+            roles = list(filter(lambda x: x, roles))
         else:
             roles = Role.query.all()
 
@@ -63,6 +64,76 @@ class RoleResource(DatabaseResource):
     @classmethod
     @safe_db_query
     def create(self, payload, user, **kwargs):
-        return super().create(merge_dict(payload, dict(
+        return super().create(merge_dict(extract(payload, [
+            'name',
+        ]), dict(
             user_id=user.id if user else None,
         )), user, **kwargs)
+
+    @safe_db_query
+    def update(self, payload, **kwargs):
+        permission_ids = [int(i) for i in payload.get('permission_ids') or []]
+        role_permissions_mapping = index_by(lambda x: x.id, self.role_permissions or [])
+
+        permission_ids_create = []
+        permission_ids_delete = []
+
+        for permission_id in permission_ids:
+            if permission_id not in role_permissions_mapping:
+                permission_ids_create.append(permission_id)
+
+        for permission_id in role_permissions_mapping.keys():
+            if permission_id not in permission_ids:
+                permission_ids_delete.append(permission_id)
+
+        if permission_ids_create:
+            db_connection.session.bulk_save_objects(
+                [RolePermission(
+                    permission_id=permission_id,
+                    role_id=self.model.id,
+                    user_id=self.current_user.id,
+                ) for permission_id in permission_ids_create],
+                return_defaults=True,
+            )
+
+        if permission_ids_delete:
+            delete_statement = RolePermission.__table__.delete().where(
+                RolePermission.permission_id.in_(permission_ids_delete),
+                RolePermission.role_id == self.id,
+            )
+            db_connection.session.execute(delete_statement)
+
+        user_ids = [int(i) for i in payload.get('user_ids') or []]
+        user_role_mapping = index_by(lambda x: x.id, self.users or [])
+
+        user_ids_create = []
+        user_ids_delete = []
+
+        for permission_id in user_ids:
+            if permission_id not in user_role_mapping:
+                user_ids_create.append(permission_id)
+
+        for permission_id in user_role_mapping.keys():
+            if permission_id not in user_ids:
+                user_ids_delete.append(permission_id)
+
+        if user_ids_create:
+            db_connection.session.bulk_save_objects(
+                [UserRole(
+                    role_id=self.model.id,
+                    user_id=user_id,
+                ) for user_id in user_ids_create],
+                return_defaults=True,
+            )
+
+        if user_ids_delete:
+            delete_statement = UserRole.__table__.delete().where(
+                UserRole.role_id == self.id,
+                UserRole.user_id.in_(user_ids_delete),
+            )
+            db_connection.session.execute(delete_statement)
+
+        return super().update(ignore_keys(payload, [
+            'permission_ids',
+            'user_ids',
+        ]), **kwargs)

mage_ai/api/resources/StatusResource.py

@@ -9,13 +9,16 @@ from mage_ai.data_preparation.repo_manager import (
 )
 from mage_ai.data_preparation.shared.constants import MANAGE_ENV_VAR
 from mage_ai.orchestration.db import safe_db_query
+from mage_ai.orchestration.db.models.schedules import PipelineRun, PipelineSchedule
 from mage_ai.server.api.clusters import ClusterType
 from mage_ai.server.scheduler_manager import scheduler_manager
 from mage_ai.settings import (
     REQUIRE_USER_AUTHENTICATION,
+    REQUIRE_USER_PERMISSIONS,
     is_disable_pipeline_edit_access,
 )
 from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.hash import merge_dict
 
 
 class StatusResource(GenericResource):
@@ -23,6 +26,7 @@ class StatusResource(GenericResource):
     Resource to fetch block output for the notebook. Created to support legacy
     endpoint /api/pipelines/<pipeline_uuid>/blocks/<block_uuid>/outputs
     """
+
     @classmethod
     @safe_db_query
     def collection(self, query, meta, user, **kwargs):
@@ -31,6 +35,7 @@ class StatusResource(GenericResource):
             GCP_PROJECT_ID,
             KUBE_NAMESPACE,
         )
+
         instance_type = None
         project_type = get_project_type()
         repo_config = get_repo_config()
@@ -45,6 +50,7 @@ class StatusResource(GenericResource):
                 from mage_ai.cluster_manager.kubernetes.workload_manager import (
                     WorkloadManager,
                 )
+
                 if WorkloadManager.load_config() or os.getenv(KUBE_NAMESPACE):
                     instance_type = ClusterType.K8S
             except ModuleNotFoundError:
@@ -58,7 +64,35 @@ class StatusResource(GenericResource):
             'disable_pipeline_edit_access': is_disable_pipeline_edit_access(),
             'max_print_output_lines': MAX_PRINT_OUTPUT_LINES,
             'require_user_authentication': REQUIRE_USER_AUTHENTICATION,
+            'require_user_permissions': REQUIRE_USER_PERMISSIONS,
             'project_type': project_type,
             'project_uuid': get_project_uuid(),
         }
+
+        display_format = meta.get('_format')
+        if 'with_activity_details' == display_format:
+            from mage_ai.server.server import latest_user_activity
+
+            project_schedules = PipelineSchedule.repo_query.all()
+            project_schedule_ids = [schedule.id for schedule in project_schedules]
+            project_pipeline_runs = PipelineRun.query.filter(
+                PipelineRun.pipeline_schedule_id.in_(project_schedule_ids)
+            )
+            sorted_pipeline_runs = project_pipeline_runs.order_by(
+                PipelineRun.updated_at.desc()
+            )
+            if sorted_pipeline_runs.count() > 0:
+                last_scheduler_activity = sorted_pipeline_runs[0].updated_at
+            active_pipeline_run_count = project_pipeline_runs.filter(
+                PipelineRun.status == PipelineRun.PipelineRunStatus.RUNNING
+            ).count()
+
+            activity_details = {
+                'active_pipeline_run_count': active_pipeline_run_count,
+                'last_scheduler_activity': last_scheduler_activity,
+                'last_user_request': latest_user_activity.latest_activity,
+            }
+
+            status = merge_dict(status, activity_details)
+
         return self.build_result_set([status], user, **kwargs)

mage_ai/api/resources/UserResource.py

@@ -9,9 +9,9 @@ from mage_ai.authentication.passwords import (
 )
 from mage_ai.data_preparation.repo_manager import get_project_uuid
 from mage_ai.orchestration.constants import Entity
-from mage_ai.orchestration.db import safe_db_query
-from mage_ai.orchestration.db.models.oauth import Permission, Role, User
-from mage_ai.shared.hash import extract, ignore_keys
+from mage_ai.orchestration.db import db_connection, safe_db_query
+from mage_ai.orchestration.db.models.oauth import Permission, Role, User, UserRole
+from mage_ai.shared.hash import extract, ignore_keys, index_by
 from mage_ai.usage_statistics.logger import UsageStatisticLogger
 
 
@@ -176,12 +176,43 @@ class UserResource(DatabaseResource):
             payload['password_hash'] = create_bcrypt_hash(password, password_salt)
             payload['password_salt'] = password_salt
 
+        role_ids = [int(i) for i in payload.get('role_ids') or []]
+        role_mapping = index_by(lambda x: x.id, self.roles_new or [])
+
+        role_ids_create = []
+        role_ids_delete = []
+
+        for role_id in role_ids:
+            if role_id not in role_mapping:
+                role_ids_create.append(role_id)
+
+        for role_id in role_mapping.keys():
+            if role_id not in role_ids:
+                role_ids_delete.append(role_id)
+
+        if role_ids_create:
+            db_connection.session.bulk_save_objects(
+                [UserRole(
+                    role_id=role_id,
+                    user_id=self.id,
+                ) for role_id in role_ids_create],
+                return_defaults=True,
+            )
+
+        if role_ids_delete:
+            delete_statement = UserRole.__table__.delete().where(
+                UserRole.role_id.in_(role_ids_delete),
+                UserRole.user_id == self.id,
+            )
+            db_connection.session.execute(delete_statement)
+
         return super().update(ignore_keys(payload, [
+            'owner',
             'password',
             'password_confirmation',
             'password_current',
-            'owner',
             'project_access',
+            'role_ids',
             'roles_display',
         ]), **kwargs)
 

mage_ai/api/result_set.py

@@ -1,13 +1,20 @@
 from collections import UserList
+from typing import List
+
 from mage_ai.api.api_context import ApiContext
 from mage_ai.api.resources.Resource import Resource
 
 
 class ResultSet(UserList):
-    def __init__(self, results):
+    def __init__(self, results: List[Resource] = None):
         self.context = ApiContext()
         self.data = []
         self.metadata = {}
+
+        if results is not None:
+            self.add_results(results)
+
+    def add_results(self, results: List[Resource]) -> None:
         for r in results:
             if issubclass(r.__class__, Resource):
                 key = r.__class__.__name__

mage_ai/authentication/permissions/constants.py

@@ -4,6 +4,8 @@ from mage_ai.data_preparation.models.constants import BlockType, PipelineType
 
 
 class EntityName(str, Enum):
+    ALL = 'ALL'
+    ALL_EXCEPT_RESERVED = 'ALL_EXCEPT_RESERVED'
     AutocompleteItem = 'AutocompleteItem'
     Backfill = 'Backfill'
     Block = 'Block'
@@ -67,6 +69,14 @@ class EntityName(str, Enum):
     Workspace = 'Workspace'
 
 
+RESERVED_ENTITY_NAMES = [
+    EntityName.Oauth,
+    EntityName.OauthAccessToken,
+    EntityName.OauthApplication,
+    EntityName.Workspace,
+]
+
+
 class BaseEntityType(str, Enum):
     pass
 
@@ -93,3 +103,71 @@ class PipelineEntityType(BaseEntityType):
     PYTHON = PipelineType.PYTHON.value
     PYSPARK = PipelineType.PYSPARK.value
     STREAMING = PipelineType.STREAMING.value
+
+
+class PermissionAccess(int, Enum):
+    OWNER = 1
+    ADMIN = 2
+    # Editor: list, detail, create, update, delete
+    EDITOR = 4
+    # Viewer: list, detail
+    VIEWER = 8
+    LIST = 16
+    DETAIL = 32
+    CREATE = 64
+    UPDATE = 128
+    DELETE = 512
+    # Operation all will allow user to perform list, detail, create, update, delete
+    # but doesn’t include granting access for query, read, and write attributes.
+    # Query, read, and write will still need to be manually granted.
+    OPERATION_ALL = 1024
+    QUERY = 2048
+    QUERY_ALL = 4096
+    READ = 8192
+    READ_ALL = 16384
+    WRITE = 32768
+    WRITE_ALL = 65536
+    # All will allow operation all, query all, read all, and write all.
+    ALL = 131072
+    DISABLE_LIST = 262144
+    DISABLE_DETAIL = 524288
+    DISABLE_CREATE = 1048576
+    DISABLE_UPDATE = 2097152
+    DISABLE_DELETE = 4194304
+    # Disable all operations: list, detail, create, update, delete.
+    DISABLE_OPERATION_ALL = 8388608
+    DISABLE_QUERY = 16777216
+    DISABLE_QUERY_ALL = 33554432
+    DISABLE_READ = 67108864
+    DISABLE_READ_ALL = 134217728
+    DISABLE_WRITE = 268435456
+    DISABLE_WRITE_ALL = 536870912
+
+
+ACCESS_FOR_VIEWER = [
+    PermissionAccess.DETAIL,
+    PermissionAccess.LIST,
+    PermissionAccess.READ,
+    PermissionAccess.VIEWER,
+]
+ACCESS_FOR_EDITOR = ACCESS_FOR_VIEWER + [
+    PermissionAccess.CREATE,
+    PermissionAccess.DELETE,
+    PermissionAccess.EDITOR,
+    PermissionAccess.QUERY,
+    PermissionAccess.UPDATE,
+    PermissionAccess.WRITE,
+]
+ACCESS_FOR_ADMIN = ACCESS_FOR_VIEWER + ACCESS_FOR_EDITOR + [
+    PermissionAccess.ADMIN,
+]
+ACCESS_FOR_OWNER = ACCESS_FOR_ADMIN + [
+    PermissionAccess.OWNER,
+]
+
+PERMISSION_ACCESS_WITH_MULTIPLE_ACCESS = {
+    f'{PermissionAccess.ADMIN}': ACCESS_FOR_ADMIN,
+    f'{PermissionAccess.EDITOR}': ACCESS_FOR_EDITOR,
+    f'{PermissionAccess.OWNER}': ACCESS_FOR_OWNER,
+    f'{PermissionAccess.VIEWER}': ACCESS_FOR_VIEWER,
+}

mage_ai/data_integrations/destinations/constants.py

@@ -11,6 +11,7 @@ DESTINATIONS = [
         uuid='mssql',
     ),
     dict(name='MySQL'),
+    dict(name='Opensearch'),
     dict(name='OracleDB'),
     dict(name='PostgreSQL'),
     dict(name='Redshift'),

mage_ai/data_integrations/utils/config.py

@@ -57,6 +57,15 @@ def __get_settings(absolute_file_path, variables: Dict = None, pipeline=None) ->
     return settings
 
 
+def get_batch_fetch_limit(source_config: Dict):
+    from mage_integrations.sources.constants import (
+        BATCH_FETCH_LIMIT,
+        BATCH_FETCH_LIMIT_KEY,
+    )
+
+    return source_config.get(BATCH_FETCH_LIMIT_KEY, BATCH_FETCH_LIMIT)
+
+
 def get_catalog(block, variables: Dict = None, pipeline=None) -> Dict:
     return get_settings(block, variables or {}, pipeline=pipeline)['catalog']
 

mage_ai/data_integrations/utils/scheduler.py

@@ -5,6 +5,7 @@ import shutil
 from typing import Dict, List, Union
 
 from mage_ai.data_integrations.sources.constants import SQL_SOURCES
+from mage_ai.data_integrations.utils.config import build_config, get_batch_fetch_limit
 from mage_ai.data_preparation.logging.logger import DictLogger
 from mage_ai.data_preparation.models.block.data_integration.constants import (
     MAX_QUERY_STRING_SIZE,
@@ -73,7 +74,7 @@ def initialize_state_and_runs(
     try:
         update_stream_states(pipeline_run, logger, variables)
 
-        block_runs = create_block_runs(pipeline_run, logger)
+        block_runs = create_block_runs(pipeline_run, logger, variables=variables)
 
         calculate_metrics(pipeline_run, logger=logger, logging_tags=tags)
 
@@ -90,8 +91,13 @@ def initialize_state_and_runs(
         raise err
 
 
-def create_block_runs(pipeline_run: PipelineRun, logger: DictLogger) -> List[BlockRun]:
-    from mage_integrations.sources.constants import BATCH_FETCH_LIMIT
+def create_block_runs(
+    pipeline_run: PipelineRun,
+    logger: DictLogger,
+    variables: Dict = None,
+) -> List[BlockRun]:
+    if variables is None:
+        variables = dict()
 
     integration_pipeline = IntegrationPipeline.get(pipeline_run.pipeline_uuid)
 
@@ -116,6 +122,10 @@ def create_block_runs(pipeline_run: PipelineRun, logger: DictLogger) -> List[Blo
         raise Exception("No Destination Block was found. \
                          Mage expects 1 Source and 1 Destination Block")
 
+    # Get batch fetch limit
+    source_config, _ = build_config(data_loader_block.file_path, variables)
+    batch_fetch_limit = get_batch_fetch_limit(source_config)
+
     transformer_blocks = [b for b in executable_blocks if b.uuid not in [
         data_loader_block.uuid,
         data_exporter_block.uuid,
@@ -139,7 +149,7 @@ def create_block_runs(pipeline_run: PipelineRun, logger: DictLogger) -> List[Blo
         record_counts = None
         if is_sql_source:
             record_counts = record_counts_by_stream[tap_stream_id]['count']
-        number_of_batches = math.ceil((record_counts or 1) / BATCH_FETCH_LIMIT)
+        number_of_batches = math.ceil((record_counts or 1) / batch_fetch_limit)
         tags2 = merge_dict(tags, dict(
             number_of_batches=number_of_batches,
             record_counts=record_counts,
@@ -372,12 +382,11 @@ def __build_block_run_metadata_for_source(
     logging_tags: Dict = None,
     selected_streams: List[str] = None,
 ) -> List[Dict]:
-    from mage_integrations.sources.constants import BATCH_FETCH_LIMIT
-
     block_run_metadata = []
 
     catalog = data_integration_settings.get('catalog')
     config = data_integration_settings.get('config')
+    batch_fetch_limit = get_batch_fetch_limit(config)
 
     streams = selected_streams or \
         [s.get('tap_stream_id') for s in get_selected_streams(catalog)]
@@ -405,7 +414,7 @@ def __build_block_run_metadata_for_source(
         record_counts = None
         if is_sql_source:
             record_counts = record_counts_by_stream[tap_stream_id]['count']
-        number_of_batches = math.ceil((record_counts or 1) / BATCH_FETCH_LIMIT)
+        number_of_batches = math.ceil((record_counts or 1) / batch_fetch_limit)
         tags2 = merge_dict(tags, dict(
             number_of_batches=number_of_batches,
             record_counts=record_counts,