Implement AGGREGATION_POLICY and PROJECTION_POLICY; Add exempt_other_…

…policies for MASKING_POLICY; Add test objects for policies; Reformat some code
littleK0i · Jun 12, 2024 · 549fcc8 · 549fcc8
1 parent 8dcb184
commit 549fcc8
Show file tree

Hide file tree

Showing 57 changed files with 955 additions and 13 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,12 @@
 # Changelog
 
+## [0.29.0] - 2024-06-12
+
+- Implemented `AGGREGATION_POLICY`, `PROJECTION_POLICY` object types.
+- Added property `exempt_other_policies` for `MASKING_POLICY`.
+- Added CLI option `--apply-all-policy` to execute SQL for all types of policies.
+- Prepared test objects for all types of policies.
+
 ## [0.28.3] - 2024-06-08
 
 - Implemented graceful warning when encounter identifier which does not conform to SnowDDL standards while processing existing role grants. Previously it caused SnowDDL to stop with hard error.

diff --git a/snowddl/app/base.py b/snowddl/app/base.py
@@ -157,9 +157,24 @@ def init_arguments_parser(self):
             default=False,
             action="store_true",
         )
+        parser.add_argument(
+            "--apply-all-policy", help="Additionally apply changes to all types of POLICIES", default=False, action="store_true"
+        )
+        parser.add_argument(
+            "--apply-aggregation-policy",
+            help="Additionally apply changes to AGGREGATION POLICIES",
+            default=False,
+            action="store_true",
+        )
         parser.add_argument(
             "--apply-masking-policy", help="Additionally apply changes to MASKING POLICIES", default=False, action="store_true"
         )
+        parser.add_argument(
+            "--apply-projection-policy",
+            help="Additionally apply changes to PROJECTION POLICIES",
+            default=False,
+            action="store_true",
+        )
         parser.add_argument(
             "--apply-row-access-policy",
             help="Additionally apply changes to ROW ACCESS POLICIES",
@@ -337,9 +352,22 @@ def init_settings(self):
             if self.args.get("apply_replace_table"):
                 settings.execute_replace_table = True
 
+            if self.args.get("apply_all_policy"):
+                settings.execute_aggregation_policy = True
+                settings.execute_masking_policy = True
+                settings.execute_projection_policy = True
+                settings.execute_row_access_policy = True
+                settings.execute_network_policy = True
+
+            if self.args.get("apply_aggregation_policy"):
+                settings.execute_aggregation_policy = True
+
             if self.args.get("apply_masking_policy"):
                 settings.execute_masking_policy = True
 
+            if self.args.get("apply_projection_policy"):
+                settings.execute_projection_policy = True
+
             if self.args.get("apply_row_access_policy"):
                 settings.execute_row_access_policy = True
 
@@ -543,8 +571,10 @@ def output_engine_stats(self):
     def output_engine_warnings(self):
         for object_type, object_names in self.engine.intention_cache.invalid_name_warning.items():
             for name in object_names:
-                self.logger.warning(f"Detected {object_type.name} with name [{name}] "
-                                    f"which does not conform to SnowDDL standards, please rename or drop it manually")
+                self.logger.warning(
+                    f"Detected {object_type.name} with name [{name}] "
+                    f"which does not conform to SnowDDL standards, please rename or drop it manually"
+                )
 
     def output_suggested_ddl(self):
         if self.engine.suggested_ddl:

diff --git a/snowddl/app/singledb.py b/snowddl/app/singledb.py
@@ -152,9 +152,24 @@ def init_arguments_parser(self):
             default=False,
             action="store_true",
         )
+        parser.add_argument(
+            "--apply-all-policy", help="Additionally apply changes to all types of POLICIES", default=False, action="store_true"
+        )
+        parser.add_argument(
+            "--apply-aggregation-policy",
+            help="Additionally apply changes to AGGREGATION POLICIES",
+            default=False,
+            action="store_true",
+        )
         parser.add_argument(
             "--apply-masking-policy", help="Additionally apply changes to MASKING POLICIES", default=False, action="store_true"
         )
+        parser.add_argument(
+            "--apply-projection-policy",
+            help="Additionally apply changes to PROJECTION POLICIES",
+            default=False,
+            action="store_true",
+        )
         parser.add_argument(
             "--apply-row-access-policy",
             help="Additionally apply changes to ROW ACCESS POLICIES",

diff --git a/snowddl/blueprint/__init__.py b/snowddl/blueprint/__init__.py
@@ -4,6 +4,7 @@
     DependsOnMixin,
     RoleBlueprint,
     AccountParameterBlueprint,
+    AggregationPolicyBlueprint,
     AlertBlueprint,
     BusinessRoleBlueprint,
     DatabaseBlueprint,
@@ -26,6 +27,7 @@
     PipeBlueprint,
     PrimaryKeyBlueprint,
     ProcedureBlueprint,
+    ProjectionPolicyBlueprint,
     ResourceMonitorBlueprint,
     RowAccessPolicyBlueprint,
     SchemaBlueprint,
@@ -85,5 +87,13 @@
 
 from .object_type import ObjectType
 from .permission_model import PermissionModel, PermissionModelCreateGrant, PermissionModelFutureGrant, PermissionModelRuleset
-from .reference import ForeignKeyReference, IndexReference, MaskingPolicyReference, RowAccessPolicyReference, TagReference
+from .reference import (
+    AggregationPolicyReference,
+    ForeignKeyReference,
+    IndexReference,
+    MaskingPolicyReference,
+    ProjectionPolicyReference,
+    RowAccessPolicyReference,
+    TagReference,
+)
 from .stage import StageWithPath, StageUploadFile
diff --git a/snowddl/blueprint/blueprint.py b/snowddl/blueprint/blueprint.py
@@ -27,7 +27,15 @@
     TableConstraintIdent,
 )
 from .object_type import ObjectType
-from .reference import ForeignKeyReference, IndexReference, MaskingPolicyReference, RowAccessPolicyReference, TagReference
+from .reference import (
+    AggregationPolicyReference,
+    ForeignKeyReference,
+    IndexReference,
+    MaskingPolicyReference,
+    ProjectionPolicyReference,
+    RowAccessPolicyReference,
+    TagReference,
+)
 from .stage import StageWithPath
 from ..model import BaseModelWithConfig
 
@@ -57,6 +65,11 @@ class AccountParameterBlueprint(AbstractBlueprint):
     value: Union[bool, float, int, str]
 
 
+class AggregationPolicyBlueprint(SchemaObjectBlueprint):
+    body: str
+    references: List[AggregationPolicyReference]
+
+
 class AlertBlueprint(SchemaObjectBlueprint):
     full_name: SchemaObjectIdent
     warehouse: AccountObjectIdent
@@ -200,6 +213,7 @@ class MaskingPolicyBlueprint(SchemaObjectBlueprint):
     arguments: List[NameWithType]
     returns: DataType
     body: str
+    exempt_other_policies: bool = False
     references: List[MaskingPolicyReference]
 
 
@@ -258,6 +272,11 @@ class ProcedureBlueprint(SchemaObjectBlueprint):
     secrets: Optional[Dict[str, SchemaObjectIdent]] = None
 
 
+class ProjectionPolicyBlueprint(SchemaObjectBlueprint):
+    body: str
+    references: List[ProjectionPolicyReference]
+
+
 class ResourceMonitorBlueprint(AbstractBlueprint):
     full_name: AccountObjectIdent
     credit_quota: int

diff --git a/snowddl/blueprint/object_type.py b/snowddl/blueprint/object_type.py
@@ -15,6 +15,12 @@ class ObjectType(Enum):
         "blueprint_cls": "AccountParameterBlueprint",
     }
 
+    AGGREGATION_POLICY = {
+        "singular": "AGGREGATION POLICY",
+        "plural": "AGGREGATION POLICIES",
+        "blueprint_cls": "AggregationPolicyBlueprint",
+    }
+
     ALERT = {
         "singular": "ALERT",
         "plural": "ALERTS",
@@ -142,6 +148,12 @@ class ObjectType(Enum):
         "blueprint_cls": "ProcedureBlueprint",
     }
 
+    PROJECTION_POLICY = {
+        "singular": "PROJECTION POLICY",
+        "plural": "PROJECTION POLICIES",
+        "blueprint_cls": "ProjectionPolicyBlueprint",
+    }
+
     RESOURCE_MONITOR = {
         "singular": "RESOURCE MONITOR",
         "plural": "RESOURCE MONITORS",

diff --git a/snowddl/blueprint/reference.py b/snowddl/blueprint/reference.py
@@ -5,6 +5,12 @@
 from ..model import BaseModelWithConfig
 
 
+class AggregationPolicyReference(BaseModelWithConfig):
+    object_type: ObjectType
+    object_name: SchemaObjectIdent
+    columns: List[Ident]
+
+
 class ForeignKeyReference(BaseModelWithConfig):
     columns: List[Ident]
     ref_table_name: SchemaObjectIdent
@@ -22,6 +28,12 @@ class MaskingPolicyReference(BaseModelWithConfig):
     columns: List[Ident]
 
 
+class ProjectionPolicyReference(BaseModelWithConfig):
+    object_type: ObjectType
+    object_name: SchemaObjectIdent
+    column: Ident
+
+
 class RowAccessPolicyReference(BaseModelWithConfig):
     object_type: ObjectType
     object_name: SchemaObjectIdent

diff --git a/snowddl/config.py b/snowddl/config.py
@@ -79,7 +79,11 @@ def get_blueprints_by_type_and_pattern(self, cls: Type[T_Blueprint], pattern: st
             else:
                 include_full_names.update(full_name for full_name in all_blueprints if regexp.match(full_name))
 
-        return {full_name: bp for full_name, bp in all_blueprints.items() if full_name in include_full_names and full_name not in exclude_full_names}
+        return {
+            full_name: bp
+            for full_name, bp in all_blueprints.items()
+            if full_name in include_full_names and full_name not in exclude_full_names
+        }
 
     def get_placeholder(self, name: str) -> Union[bool, float, int, str]:
         if name not in self.placeholders:

diff --git a/snowddl/parser/__init__.py b/snowddl/parser/__init__.py
@@ -1,5 +1,6 @@
 from ._parsed_file import ParsedFile
 from .account_params import AccountParameterParser
+from .aggregation_policy import AggregationPolicyParser
 from .alert import AlertParser
 from .business_role import BusinessRoleParser
 from .database import DatabaseParser
@@ -21,6 +22,7 @@
 from .pipe import PipeParser
 from .placeholder import PlaceholderParser
 from .procedure import ProcedureParser
+from .projection_policy import ProjectionPolicyParser
 from .resource_monitor import ResourceMonitorParser
 from .row_access_policy import RowAccessPolicyParser
 from .schema import SchemaParser
@@ -63,7 +65,9 @@
     ViewParser,
     PipeParser,
     TaskParser,
+    AggregationPolicyParser,
     MaskingPolicyParser,
+    ProjectionPolicyParser,
     RowAccessPolicyParser,
     OutboundShareParser,
     TechnicalRoleParser,
@@ -94,7 +98,9 @@
     ViewParser,
     PipeParser,
     TaskParser,
+    AggregationPolicyParser,
     MaskingPolicyParser,
+    ProjectionPolicyParser,
     RowAccessPolicyParser,
     AlertParser,
 ]
diff --git a/snowddl/parser/aggregation_policy.py b/snowddl/parser/aggregation_policy.py
@@ -0,0 +1,76 @@
+from snowddl.blueprint import (
+    AggregationPolicyBlueprint,
+    AggregationPolicyReference,
+    Ident,
+    ObjectType,
+    SchemaObjectIdent,
+    build_schema_object_ident,
+)
+from snowddl.parser.abc_parser import AbstractParser, ParsedFile
+
+
+# fmt: off
+aggregation_policy_json_schema = {
+    "type": "object",
+    "properties": {
+        "body": {
+            "type": "string"
+        },
+        "references": {
+            "type": "array",
+            "items": {
+                "type": "object",
+                "properties": {
+                    "object_type": {
+                        "type": "string"
+                    },
+                    "object_name": {
+                        "type": "string"
+                    },
+                    "columns": {
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "minItems": 1
+                    }
+                },
+                "required": ["object_type", "object_name"],
+                "additionalProperties": False
+            },
+            "minItems": 1
+        },
+        "comment": {
+            "type": "string"
+        }
+    },
+    "required": ["body"],
+    "additionalProperties": False
+}
+# fmt: on
+
+
+class AggregationPolicyParser(AbstractParser):
+    def load_blueprints(self):
+        self.parse_schema_object_files("aggregation_policy", aggregation_policy_json_schema, self.process_aggregation_policy)
+
+    def process_aggregation_policy(self, f: ParsedFile):
+        references = []
+
+        for a in f.params.get("references", []):
+            ref = AggregationPolicyReference(
+                object_type=ObjectType[a["object_type"].upper()],
+                object_name=build_schema_object_ident(self.env_prefix, a["object_name"], f.database, f.schema),
+                columns=[Ident(c) for c in a.get("columns", [])],
+            )
+
+            references.append(ref)
+
+        bp = AggregationPolicyBlueprint(
+            full_name=SchemaObjectIdent(self.env_prefix, f.database, f.schema, f.name),
+            body=f.params["body"],
+            references=references,
+            comment=f.params.get("comment"),
+        )
+
+        self.config.add_blueprint(bp)
diff --git a/snowddl/parser/database.py b/snowddl/parser/database.py
@@ -64,8 +64,10 @@ def load_blueprints(self):
             database_name = database_path.name.upper()
             database_params = self.parse_single_file(database_path / "params.yaml", database_json_schema)
 
+            # fmt: off
             databases_permission_model_name = database_params.get("permission_model", self.config.DEFAULT_PERMISSION_MODEL).upper()
             database_permission_model = self.config.get_permission_model(databases_permission_model_name)
+            # fmt: on
 
             if not database_permission_model.ruleset.create_database_owner_role:
                 for k in database_params:

diff --git a/snowddl/parser/masking_policy.py b/snowddl/parser/masking_policy.py
@@ -27,6 +27,9 @@
         "body": {
             "type": "string"
         },
+        "exempt_other_policies": {
+            "type": "boolean"
+        },
         "references": {
             "type": "array",
             "items": {
@@ -83,6 +86,7 @@ def process_masking_policy(self, f: ParsedFile):
             body=f.params["body"],
             arguments=arguments,
             returns=DataType(f.params["returns"]),
+            exempt_other_policies=f.params.get("exempt_other_policies", False),
             references=references,
             comment=f.params.get("comment"),
         )