[td] Add models for custom design framework (mage-ai#4192)

Author: tommydangerous

mage_ai/api/policies/CustomDesignPolicy.py

@@ -0,0 +1,33 @@
+from mage_ai.api.oauth_scope import OauthScope
+from mage_ai.api.operations.constants import OperationType
+from mage_ai.api.policies.BasePolicy import BasePolicy
+from mage_ai.api.presenters.CustomDesignPresenter import CustomDesignPresenter
+
+
+class CustomDesignPolicy(BasePolicy):
+    pass
+
+
+CustomDesignPolicy.allow_actions(
+    [
+        OperationType.LIST,
+    ],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+    override_permission_condition=lambda _policy: True,
+)
+
+
+CustomDesignPolicy.allow_read(
+    CustomDesignPresenter.default_attributes + [],
+    scopes=[
+        OauthScope.CLIENT_PRIVATE,
+    ],
+    on_action=[
+        OperationType.LIST,
+    ],
+    condition=lambda policy: policy.has_at_least_viewer_role(),
+    override_permission_condition=lambda _policy: True,
+)

mage_ai/api/presenters/CustomDesignPresenter.py

@@ -0,0 +1,14 @@
+from typing import Dict
+
+from mage_ai.api.presenters.BasePresenter import BasePresenter
+
+
+class CustomDesignPresenter(BasePresenter):
+    default_attributes = [
+        'components',
+        'pages',
+        'project',
+    ]
+
+    async def prepare_present(self, **kwargs) -> Dict:
+        return self.resource.model.to_dict()

mage_ai/api/resources/CustomDesignResource.py

@@ -0,0 +1,17 @@
+from typing import Dict
+
+from mage_ai.api.resources.AsyncBaseResource import AsyncBaseResource
+from mage_ai.presenters.design.models import CustomDesign
+
+
+class CustomDesignResource(AsyncBaseResource):
+    @classmethod
+    async def collection(self, query: Dict, _meta: Dict, user, **kwargs):
+        custom_design = CustomDesign.load_from_file(all_configurations=True)
+
+        arr = [custom_design]
+
+        if custom_design.custom_designs:
+            arr.extend(list(custom_design.custom_designs.values()))
+
+        return self.build_result_set(arr, user, **kwargs)

mage_ai/authentication/permissions/constants.py

@@ -21,6 +21,7 @@ class EntityName(str, Enum):
     ComputeConnection = 'ComputeConnection'
     ComputeService = 'ComputeService'
     ConfigurationOption = 'ConfigurationOption'
+    CustomDesign = 'CustomDesign'
     CustomTemplate = 'CustomTemplate'
     DataProvider = 'DataProvider'
     Database = 'Database'

mage_ai/data_preparation/models/project/constants.py

@@ -5,6 +5,7 @@ class FeatureUUID(str, Enum):
     ADD_NEW_BLOCK_V2 = 'add_new_block_v2'
     COMPUTE_MANAGEMENT = 'compute_management'
     DATA_INTEGRATION_IN_BATCH_PIPELINE = 'data_integration_in_batch_pipeline'
+    DBT_V2 = 'dbt_v2'
     GLOBAL_HOOKS = 'global_hooks'
     INTERACTIONS = 'interactions'
     LOCAL_TIMEZONE = 'display_local_timezone'

mage_ai/data_preparation/models/project/models.py

@@ -10,3 +10,13 @@ class ProjectPipelines(BaseDataClass):
 
     def __post_init__(self):
         self.serialize_attribute_class('settings', PipelineSettings)
+
+
+@dataclass
+class ProjectDataClass(BaseDataClass):
+    full_path: str = None
+    full_path_relative: str = None
+    path: str = None
+    root_project_name: str = None
+    root_project_full_path: str = None
+    uuid: str = None

mage_ai/frontend/interfaces/ProjectType.ts

@@ -4,6 +4,7 @@ export enum FeatureUUIDEnum {
   ADD_NEW_BLOCK_V2 = 'add_new_block_v2',
   COMPUTE_MANAGEMENT = 'compute_management',
   DATA_INTEGRATION_IN_BATCH_PIPELINE = 'data_integration_in_batch_pipeline',
+  DBT_V2 = 'dbt_v2',
   GLOBAL_HOOKS = 'global_hooks',
   INTERACTIONS = 'interactions',
   NOTEBOOK_BLOCK_OUTPUT_SPLIT_VIEW = 'notebook_block_output_split_view',

mage_ai/presenters/design/__init__.py

@@ -0,0 +1 @@
+CUSTOM_DESIGN_FILENAME = 'design.yaml'

mage_ai/presenters/design/constants.py

@@ -0,0 +1,125 @@
+import os
+from dataclasses import dataclass
+from typing import Dict
+
+import yaml
+
+from mage_ai.data_preparation.models.project.models import ProjectDataClass
+from mage_ai.presenters.design.constants import CUSTOM_DESIGN_FILENAME
+from mage_ai.settings.platform import (
+    build_repo_path_for_all_projects,
+    project_platform_activated,
+)
+from mage_ai.settings.repo import get_repo_path
+from mage_ai.shared.io import safe_write
+from mage_ai.shared.models import BaseDataClass
+
+
+@dataclass
+class DesignComponentConfigurations(BaseDataClass):
+    logo: Dict = None
+
+
+@dataclass
+class DesignComponents(BaseDataClass):
+    header: DesignComponentConfigurations = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('header', DesignComponentConfigurations)
+
+
+@dataclass
+class DesignPageConfigurations(BaseDataClass):
+    edit: Dict = None
+
+
+@dataclass
+class DesignPages(BaseDataClass):
+    pipelines: DesignPageConfigurations = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('pipelines', DesignPageConfigurations)
+
+
+@dataclass
+class CustomDesign(BaseDataClass):
+    components: DesignComponents = None
+    custom_designs: Dict = None
+    pages: DesignPages = None
+    project: ProjectDataClass = None
+
+    def __post_init__(self):
+        self.serialize_attribute_class('components', DesignComponents)
+        self.serialize_attribute_class('pages', DesignPages)
+        self.serialize_attribute_class('project', ProjectDataClass)
+
+    @classmethod
+    def file_path(self, repo_path: str = None) -> str:
+        return os.path.join(repo_path or get_repo_path(), CUSTOM_DESIGN_FILENAME)
+
+    @classmethod
+    def load_from_file(
+        self,
+        all_configurations: bool = True,
+        file_path: str = None,
+        repo_path: str = None,
+        project: Dict = None,
+    ) -> 'CustomDesign':
+        model = self.__load_from_file(file_path=file_path, project=project, repo_path=repo_path)
+
+        if all_configurations and project_platform_activated():
+            model.custom_designs = {}
+
+            for project_name, project in build_repo_path_for_all_projects().items():
+                full_path = project['full_path']
+
+                model.custom_designs[project_name] = self.__load_from_file(
+                    file_path=os.path.join(full_path, CUSTOM_DESIGN_FILENAME),
+                    project=project,
+                )
+
+        return model
+
+    @classmethod
+    def __load_from_file(
+        self,
+        file_path: str = None,
+        project: Dict = None,
+        repo_path: str = None,
+    ) -> 'CustomDesign':
+        yaml_config = {}
+
+        file_path_to_use = file_path or self.file_path(repo_path=repo_path)
+        if os.path.exists(file_path_to_use):
+            with open(file_path_to_use, 'r') as fp:
+                content = fp.read()
+                if content:
+                    yaml_config = yaml.safe_load(content) or {}
+
+        return self.load(project=project, **yaml_config)
+
+    def save(self, file_path: str = None, repo_path: str = None) -> None:
+        if not file_path:
+            file_path = self.file_path(repo_path=repo_path)
+
+        content_original = None
+        if os.path.exists(file_path):
+            with open(file_path) as f:
+                content_original = f.read()
+
+        with open(file_path, 'w'):
+            try:
+                data = self.to_dict()
+                content = yaml.safe_dump(data)
+                safe_write(file_path, content)
+            except Exception as err:
+                if content_original:
+                    safe_write(file_path, content_original)
+                raise err
+
+    def to_dict(self, **kwargs) -> Dict:
+        return super().to_dict(
+            convert_enum=True,
+            ignore_attributes=['custom_designs'],
+            ignore_empty=True,
+        )