fix: expand grants_config parsing to support more complex expressions

Author: newtonapple

sqlmesh/core/model/meta.py

@@ -523,30 +523,62 @@ def custom_materialization_properties(self) -> CustomMaterializationProperties:
     def grants(self) -> t.Optional[GrantsConfig]:
         """A dictionary of grants mapping permission names to lists of grantees."""
 
-        if not self.grants_:
+        if self.grants_ is None:
             return None
 
-        def parse_exp_to_str(e: exp.Expression) -> str:
-            if isinstance(e, exp.Literal) and e.is_string:
-                return e.this.strip()
-            if isinstance(e, exp.Identifier):
-                return e.name
-            return e.sql(dialect=self.dialect).strip()
+        if not self.grants_.expressions:
+            return {}
+
+        def expr_to_string(expr: exp.Expression, context: str) -> str:
+            if isinstance(expr, (d.MacroFunc, d.MacroVar)):
+                raise ConfigError(
+                    f"Unresolved macro in {context}: {expr.sql(dialect=self.dialect)}"
+                )
+
+            if isinstance(expr, exp.Null):
+                raise ConfigError(f"NULL value in {context}")
+
+            if isinstance(expr, exp.Literal):
+                return str(expr.this).strip()
+            if isinstance(expr, exp.Identifier):
+                return expr.name
+            if isinstance(expr, exp.Column):
+                return expr.name
+            return expr.sql(dialect=self.dialect).strip()
+
+        def normalize_to_string_list(value_expr: exp.Expression) -> t.List[str]:
+            result = []
+
+            def process_expression(expr: exp.Expression) -> None:
+                if isinstance(expr, exp.Array):
+                    for elem in expr.expressions:
+                        process_expression(elem)
+
+                elif isinstance(expr, (exp.Tuple, exp.Paren)):
+                    expressions = (
+                        [expr.unnest()] if isinstance(expr, exp.Paren) else expr.expressions
+                    )
+                    for elem in expressions:
+                        process_expression(elem)
+                else:
+                    result.append(expr_to_string(expr, "grant value"))
+
+            process_expression(value_expr)
+            return result
 
         grants_dict = {}
         for eq_expr in self.grants_.expressions:
-            permission_name = parse_exp_to_str(eq_expr.this)  # left hand side
-            grantees_expr = eq_expr.expression  # right hand side
-            if isinstance(grantees_expr, exp.Array):
-                grantee_list = []
-                for grantee_expr in grantees_expr.expressions:
-                    grantee = parse_exp_to_str(grantee_expr)
-                    if grantee:  # skip empty strings
-                        grantee_list.append(grantee)
-
-                grants_dict[permission_name.strip()] = grantee_list
-
-        return grants_dict
+            try:
+                permission_name = expr_to_string(eq_expr.left, "permission name")
+                grantee_list = normalize_to_string_list(eq_expr.expression)
+                grants_dict[permission_name] = grantee_list
+            except ConfigError as e:
+                permission_name = (
+                    eq_expr.left.name if hasattr(eq_expr.left, "name") else str(eq_expr.left)
+                )
+                raise ConfigError(f"Invalid grants configuration for '{permission_name}': {e}")
+
+        return grants_dict if grants_dict else None
 
     @property
     def all_references(self) -> t.List[Reference]:

tests/core/test_model.py

@@ -61,7 +61,7 @@
     model,
 )
 from sqlmesh.core.model.common import parse_expression
-from sqlmesh.core.model.kind import ModelKindName, _model_kind_validator
+from sqlmesh.core.model.kind import _ModelKind, ModelKindName, _model_kind_validator
 from sqlmesh.core.model.seed import CsvSettings
 from sqlmesh.core.node import IntervalUnit, _Node
 from sqlmesh.core.signal import signal
@@ -11732,90 +11732,235 @@ def test_use_original_sql():
     assert model.post_statements_[0].sql == "CREATE TABLE post (b INT)"
 
 
-def test_grants_validation_symbolic_model_error():
-    with pytest.raises(ValidationError, match=r".*grants cannot be set for EXTERNAL.*"):
-        create_sql_model(
-            "db.table",
-            parse_one("SELECT 1 AS id"),
-            kind="EXTERNAL",
-            grants={"select": ["user1", "user2"], "insert": ["admin_user"]},
-        )
+@pytest.mark.parametrize(
+    "kind",
+    [
+        "FULL",
+        "VIEW",
+        SeedKind(path="test.csv"),
+        IncrementalByTimeRangeKind(time_column="ds"),
+        IncrementalByUniqueKeyKind(unique_key="id"),
+    ],
+)
+def test_grants_valid_model_kinds(kind: t.Union[str, _ModelKind]):
+    model = create_sql_model(
+        "db.table",
+        parse_one("SELECT 1 AS id"),
+        kind=kind,
+        grants={"select": ["user1", "user2"], "insert": ["admin_user"]},
+    )
+    assert model.grants == {"select": ["user1", "user2"], "insert": ["admin_user"]}
 
 
-def test_grants_validation_embedded_model_error():
-    with pytest.raises(ValidationError, match=r".*grants cannot be set for EMBEDDED.*"):
+@pytest.mark.parametrize(
+    "kind",
+    [
+        "EXTERNAL",
+        "EMBEDDED",
+    ],
+)
+def test_grants_invalid_model_kind_errors(kind: str):
+    with pytest.raises(ValidationError, match=rf".*grants cannot be set for {kind}.*"):
         create_sql_model(
             "db.table",
             parse_one("SELECT 1 AS id"),
-            kind="EMBEDDED",
+            kind=kind,
             grants={"select": ["user1"], "insert": ["admin_user"]},
         )
 
 
-def test_grants_validation_valid_seed_model():
+def test_grants_validation_no_grants():
+    model = create_sql_model("db.table", parse_one("SELECT 1 AS id"), kind="FULL")
+    assert model.grants is None
+
+
+def test_grants_validation_empty_grantees():
     model = create_sql_model(
-        "db.table",
-        parse_one("SELECT 1 AS id"),
-        kind=SeedKind(path="test.csv"),
-        grants={"select": ["user1"], "insert": ["admin_user"]},
+        "db.table", parse_one("SELECT 1 AS id"), kind="FULL", grants={"select": []}
     )
-    assert model.grants == {"select": ["user1"], "insert": ["admin_user"]}
+    assert model.grants == {"select": []}
 
 
-def test_grants_validation_valid_materialized_model():
+def test_grants_single_value_conversions():
+    expressions = d.parse(f"""
+        MODEL (
+            name test.nested_arrays,
+            kind FULL,
+            grants (
+                'select' = "user1", update = user2
+            )
+        );
+        SELECT 1 as id
+    """)
+    model = load_sql_based_model(expressions)
+    assert model.grants == {"select": ["user1"], "update": ["user2"]}
+
     model = create_sql_model(
         "db.table",
         parse_one("SELECT 1 AS id"),
         kind="FULL",
-        grants={"select": ["user1", "user2"], "insert": ["admin_user"]},
+        grants={"select": "user1", "insert": 123},
     )
-    assert model.grants == {"select": ["user1", "user2"], "insert": ["admin_user"]}
+    assert model.grants == {"select": ["user1"], "insert": ["123"]}
 
 
-def test_grants_validation_valid_view_model():
-    model = create_sql_model(
-        "db.table", parse_one("SELECT 1 AS id"), kind="VIEW", grants={"select": ["user1", "user2"]}
+@pytest.mark.parametrize(
+    "grantees",
+    [
+        "('user1', ('user2', 'user3'), 'user4')",
+        "('user1', ['user2', 'user3'], user4)",
+        "['user1', ['user2', user3], 'user4']",
+        "[user1, ('user2', \"user3\"), 'user4']",
+    ],
+)
+def test_grants_array_flattening(grantees: str):
+    expressions = d.parse(f"""
+        MODEL (
+            name test.nested_arrays,
+            kind FULL,
+            grants (
+                'select' = {grantees}
+            )
+        );
+        SELECT 1 as id
+    """)
+    model = load_sql_based_model(expressions)
+    assert model.grants == {"select": ["user1", "user2", "user3", "user4"]}
+
+
+def test_grants_macro_var_resolved():
+    expressions = d.parse("""
+        MODEL (
+            name test.macro_grants,
+            kind FULL,
+            grants (
+                'select' = @VAR('readers'),
+                'insert' = @VAR('writers')
+            )
+        );
+        SELECT 1 as id
+    """)
+    model = load_sql_based_model(
+        expressions, variables={"readers": ["user1", "user2"], "writers": "admin"}
     )
-    assert model.grants == {"select": ["user1", "user2"]}
+    assert model.grants == {
+        "select": ["user1", "user2"],
+        "insert": ["admin"],
+    }
 
 
-def test_grants_validation_valid_incremental_model():
-    model = create_sql_model(
-        "db.table",
-        parse_one("SELECT 1 AS id, CURRENT_TIMESTAMP AS ts"),
-        kind=IncrementalByTimeRangeKind(time_column="ts"),
-        grants={"select": ["user1"], "update": ["admin_user"]},
+def test_grants_macro_var_in_array_flattening():
+    expressions = d.parse("""
+        MODEL (
+            name test.macro_in_array,
+            kind FULL,
+            grants (
+                'select' = ['user1', @VAR('admins'), 'user3']
+            )
+        );
+        SELECT 1 as id
+    """)
+
+    model = load_sql_based_model(expressions, variables={"admins": ["admin1", "admin2"]})
+    assert model.grants == {"select": ["user1", "admin1", "admin2", "user3"]}
+
+    model2 = load_sql_based_model(expressions, variables={"admins": "super_admin"})
+    assert model2.grants == {"select": ["user1", "super_admin", "user3"]}
+
+
+def test_grants_dynamic_permission_names():
+    expressions = d.parse("""
+        MODEL (
+            name test.dynamic_keys,
+            kind FULL,
+            grants (
+                @VAR('read_perm') = ['user1', 'user2'],
+                @VAR('write_perm') = ['admin']
+            )
+        );
+        SELECT 1 as id
+    """)
+    model = load_sql_based_model(
+        expressions, variables={"read_perm": "select", "write_perm": "insert"}
     )
-    assert model.grants == {"select": ["user1"], "update": ["admin_user"]}
+    assert model.grants == {"select": ["user1", "user2"], "insert": ["admin"]}
 
 
-def test_grants_validation_no_grants():
-    model = create_sql_model("db.table", parse_one("SELECT 1 AS id"), kind="FULL")
-    assert model.grants is None
+def test_grants_unresolved_macro_errors():
+    expressions1 = d.parse("""
+        MODEL (name test.bad1, kind FULL, grants ('select' = @VAR('undefined')));
+        SELECT 1 as id
+    """)
+    with pytest.raises(ConfigError, match=r"Invalid grants configuration for 'select': NULL value"):
+        load_sql_based_model(expressions1)
 
+    expressions2 = d.parse("""
+        MODEL (name test.bad2, kind FULL, grants (@VAR('undefined') = ['user']));
+        SELECT 1 as id
+    """)
+    with pytest.raises(ConfigError, match=r"Invalid grants configuration.*NULL value"):
+        load_sql_based_model(expressions2)
 
-def test_grants_validation_empty_grantees():
-    model = create_sql_model(
+    expressions3 = d.parse("""
+        MODEL (name test.bad3, kind FULL, grants ('select' = ['user', @VAR('undefined')]));
+        SELECT 1 as id
+    """)
+    with pytest.raises(ConfigError, match=r"Invalid grants configuration for 'select': NULL value"):
+        load_sql_based_model(expressions3)
+
+
+def test_grants_mixed_types_conversion():
+    expressions = d.parse("""
+        MODEL (
+            name test.mixed_types,
+            kind FULL,
+            grants (
+                'select' = ['user1', 123, admin_role, 'user2']
+            )
+        );
+        SELECT 1 as id
+    """)
+    model = load_sql_based_model(expressions)
+    assert model.grants == {"select": ["user1", "123", "admin_role", "user2"]}
+
+
+def test_grants_empty_values():
+    model1 = create_sql_model(
         "db.table", parse_one("SELECT 1 AS id"), kind="FULL", grants={"select": []}
     )
-    assert model.grants == {"select": []}
+    assert model1.grants == {"select": []}
 
+    model2 = create_sql_model("db.table", parse_one("SELECT 1 AS id"), kind="FULL")
+    assert model2.grants is None
 
-def test_grants_table_type_view():
-    model = create_sql_model("test_view", parse_one("SELECT 1 as id"), kind="VIEW")
-    assert model.grants_table_type == DataObjectType.VIEW
 
+def test_grants_backward_compatibility():
     model = create_sql_model(
-        "test_mv", parse_one("SELECT 1 as id"), kind=ViewKind(materialized=True)
+        "db.table",
+        parse_one("SELECT 1 AS id"),
+        kind="FULL",
+        grants={
+            "select": ["user1", "user2"],
+            "insert": ["admin"],
+            "roles/bigquery.dataViewer": ["user:data_eng@company.com"],
+        },
     )
-    assert model.grants_table_type == DataObjectType.MATERIALIZED_VIEW
-
-
-def test_grants_table_type_table():
-    model = create_sql_model("test_table", parse_one("SELECT 1 as id"), kind="FULL")
-    assert model.grants_table_type == DataObjectType.TABLE
+    assert model.grants == {
+        "select": ["user1", "user2"],
+        "insert": ["admin"],
+        "roles/bigquery.dataViewer": ["user:data_eng@company.com"],
+    }
 
 
-def test_grants_table_type_managed():
-    model = create_sql_model("test_managed", parse_one("SELECT 1 as id"), kind="MANAGED")
-    assert model.grants_table_type == DataObjectType.MANAGED_TABLE
+@pytest.mark.parametrize(
+    "kind, expected",
+    [
+        ("VIEW", DataObjectType.VIEW),
+        ("FULL", DataObjectType.TABLE),
+        ("MANAGED", DataObjectType.MANAGED_TABLE),
+        (ViewKind(materialized=True), DataObjectType.MATERIALIZED_VIEW),
+    ],
+)
+def test_grants_table_type(kind: t.Union[str, _ModelKind], expected: DataObjectType):
+    model = create_sql_model("test_table", parse_one("SELECT 1 as id"), kind=kind)
+    assert model.grants_table_type == expected