feat: set query label session property in bq session (#4314)

Author: ncbkr

sqlmesh/core/engine_adapter/bigquery.py

@@ -183,7 +183,31 @@ def query_factory() -> Query:
     def _begin_session(self, properties: SessionProperties) -> None:
         from google.cloud.bigquery import QueryJobConfig
 
-        job = self.client.query("SELECT 1;", job_config=QueryJobConfig(create_session=True))
+        query_label_property = properties.get("query_label")
+        parsed_query_label: list[tuple[str, str]] = []
+        if isinstance(query_label_property, (exp.Array, exp.Paren, exp.Tuple)):
+            label_tuples = (
+                [query_label_property.unnest()]
+                if isinstance(query_label_property, exp.Paren)
+                else query_label_property.expressions
+            )
+
+            # query_label is a Paren, Array or Tuple of 2-tuples and validated at load time
+            parsed_query_label.extend(
+                (label_tuple.expressions[0].name, label_tuple.expressions[1].name)
+                for label_tuple in label_tuples
+            )
+
+        if parsed_query_label:
+            query_label_str = ",".join([":".join(label) for label in parsed_query_label])
+            query = f'SET @@query_label = "{query_label_str}";SELECT 1;'
+        else:
+            query = "SELECT 1;"
+
+        job = self.client.query(
+            query,
+            job_config=QueryJobConfig(create_session=True),
+        )
         session_info = job.session_info
         session_id = session_info.session_id if session_info else None
         self._session_id = session_id

sqlmesh/core/model/common.py

@@ -441,7 +441,6 @@ def _executable_to_str(k: str, v: Executable) -> str:
 properties_validator: t.Callable = field_validator(
     "physical_properties_",
     "virtual_properties_",
-    "session_properties_",
     "materialization_properties_",
     mode="before",
     check_fields=False,

sqlmesh/core/model/meta.py

@@ -17,6 +17,7 @@
     default_catalog_validator,
     depends_on_validator,
     properties_validator,
+    parse_properties,
 )
 from sqlmesh.core.model.kind import (
     CustomKind,
@@ -310,6 +311,43 @@ def _refs_validator(cls, vs: t.Any, info: ValidationInfo) -> t.List[exp.Expressi
     def ignored_rules_validator(cls, vs: t.Any) -> t.Any:
         return LinterConfig._validate_rules(vs)
 
+    @field_validator("session_properties_", mode="before")
+    def session_properties_validator(cls, v: t.Any, info: ValidationInfo) -> t.Any:
+        # use the generic properties validator to parse the session properties
+        parsed_session_properties = parse_properties(type(cls), v, info)
+        if not parsed_session_properties:
+            return parsed_session_properties
+
+        for eq in parsed_session_properties:
+            if eq.name == "query_label":
+                query_label = eq.right
+                if not (
+                    isinstance(query_label, exp.Array)
+                    or isinstance(query_label, exp.Tuple)
+                    or isinstance(query_label, exp.Paren)
+                ):
+                    raise ConfigError(
+                        "Invalid value for `session_properties.query_label`. Must be an array or tuple."
+                    )
+
+                label_tuples: t.List[exp.Expression] = (
+                    [query_label.unnest()]
+                    if isinstance(query_label, exp.Paren)
+                    else query_label.expressions
+                )
+
+                for label_tuple in label_tuples:
+                    if not (
+                        isinstance(label_tuple, exp.Tuple)
+                        and len(label_tuple.expressions) == 2
+                        and all(isinstance(label, exp.Literal) for label in label_tuple.expressions)
+                    ):
+                        raise ConfigError(
+                            "Invalid entry in `session_properties.query_label`. Must be tuples of string literals with length 2."
+                        )
+
+        return parsed_session_properties
+
     @model_validator(mode="before")
     def _pre_root_validator(cls, data: t.Any) -> t.Any:
         if not isinstance(data, dict):

tests/core/engine_adapter/test_bigquery.py

@@ -531,6 +531,7 @@ def test_begin_end_session(mocker: MockerFixture):
 
     adapter = BigQueryEngineAdapter(lambda: connection_mock, job_retries=0)
 
+    # starting a session without session properties
     with adapter.session({}):
         assert adapter._connection_pool.get_attribute("session_id") is not None
         adapter.execute("SELECT 2;")
@@ -551,6 +552,18 @@ def test_begin_end_session(mocker: MockerFixture):
     assert execute_b_call[1]["query"] == "SELECT 3;"
     assert not execute_b_call[1]["job_config"].connection_properties
 
+    # starting a new session with session property query_label and array value
+    with adapter.session({"query_label": parse_one("[('key1', 'value1'), ('key2', 'value2')]")}):
+        adapter.execute("SELECT 4;")
+    begin_new_session_call = connection_mock._client.query.call_args_list[3]
+    assert begin_new_session_call[0][0] == 'SET @@query_label = "key1:value1,key2:value2";SELECT 1;'
+
+    # starting a new session with session property query_label and Paren value
+    with adapter.session({"query_label": parse_one("(('key1', 'value1'))")}):
+        adapter.execute("SELECT 5;")
+    begin_new_session_call = connection_mock._client.query.call_args_list[5]
+    assert begin_new_session_call[0][0] == 'SET @@query_label = "key1:value1";SELECT 1;'
+
 
 def _to_sql_calls(execute_mock: t.Any, identify: bool = True) -> t.List[str]:
     output = []

tests/core/test_model.py

@@ -3757,7 +3757,10 @@ def my_model(context, **kwargs):
         """('key_a' = 'value_a', 'key_b' = 1, 'key_c' = TRUE, 'key_d' = 2.0)"""
     )
 
-    with pytest.raises(ConfigError, match=r"Invalid property 'invalid'.*"):
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid property 'invalid'. Properties must be specified as key-value pairs <key> = <value>. ",
+    ):
         load_sql_based_model(
             d.parse(
                 """
@@ -4418,6 +4421,108 @@ def test_model_session_properties(sushi_context):
         "warehouse": "test_warehouse",
     }
 
+    model = load_sql_based_model(
+        d.parse(
+            """
+        MODEL (
+            name test_schema.test_model,
+            session_properties (
+                'query_label' = [
+                    ('key1', 'value1'),
+                    ('key2', 'value2')
+                ]
+            )
+        );
+        SELECT a FROM tbl;
+        """,
+            default_dialect="bigquery",
+        )
+    )
+    assert model.session_properties == {
+        "query_label": parse_one("[('key1', 'value1'), ('key2', 'value2')]")
+    }
+
+    model = load_sql_based_model(
+        d.parse(
+            """
+        MODEL (
+            name test_schema.test_model,
+            session_properties (
+                'query_label' = (
+                    ('key1', 'value1')
+                )
+            )
+        );
+        SELECT a FROM tbl;
+        """,
+            default_dialect="bigquery",
+        )
+    )
+    assert model.session_properties == {"query_label": parse_one("(('key1', 'value1'))")}
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid value for `session_properties.query_label`. Must be an array or tuple.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                session_properties (
+                    'query_label' = 'invalid value'
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="bigquery",
+            )
+        )
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid entry in `session_properties.query_label`. Must be tuples of string literals with length 2.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                session_properties (
+                    'query_label' = (
+                        ('key1', 'value1', 'another_value')
+                    )
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="bigquery",
+            )
+        )
+
+    with pytest.raises(
+        ConfigError,
+        match=r"Invalid entry in `session_properties.query_label`. Must be tuples of string literals with length 2.",
+    ):
+        load_sql_based_model(
+            d.parse(
+                """
+            MODEL (
+                name test_schema.test_model,
+                session_properties (
+                    'query_label' = (
+                        'some value', 
+                        'another value',
+                        'yet another value',
+                    )
+                )
+            );
+            SELECT a FROM tbl;
+            """,
+                default_dialect="bigquery",
+            )
+        )
+
 
 def test_model_jinja_macro_rendering():
     expressions = d.parse(