From 207751739680e5892bbcfee7ef26e347bac5b44e Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Thu, 29 Oct 2026 10:37:57 -0700 Subject: [PATCH 1/2] fix: qualify short session template names in _get_session_template CreateSession accepts either a bare session template ID or a fully-qualified projects/.../locations/.../sessionTemplates/... resource name on Session.session_template, whereas GetSessionTemplateRequest.name requires a fully-qualified resource name. When a bare session template ID is configured and MANAGED_SPARK_CONNECT_* execution environment variables are present, qualify the template name to projects/{project}/locations/{region}/sessionTemplates/{template} before calling GetSessionTemplate so template field precedence is respected. --- google/cloud/managed_spark_connect/session.py | 5 +++ tests/unit/test_session.py | 36 +++++++++++++++++++ 2 files changed, 41 insertions(+) diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index e00a3dc1..b5951768 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -1348,6 +1348,11 @@ def _get_session_template(self, template_name: str): if not template_name: return None + if "/" not in template_name and self._project_id and self._region: + template_name = ( + f"projects/{self._project_id}/locations/{self._region}" + f"/sessionTemplates/{template_name}" + ) try: get_session_template_request = GetSessionTemplateRequest() get_session_template_request.name = template_name diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 7d64d779..f5ecb75b 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -39,6 +39,7 @@ CreateSessionRequest, DeleteSessionRequest, GetSessionRequest, + GetSessionTemplateRequest, Session, SessionTemplate, SparkConnectConfig, @@ -3775,6 +3776,41 @@ def test_no_template_lookup_without_exec_env_vars( mock_template_client.return_value.get_session_template.assert_not_called() self.assertEqual(session_config.session_template, self._TEMPLATE) + @mock.patch("google.cloud.dataproc_v1.SessionTemplateControllerClient") + def test_short_session_template_name_qualifies_get_session_template_request( + self, mock_template_client + ): + """A bare session template name is qualified to a full sessionTemplates URI for GetSessionTemplate.""" + template = SessionTemplate(name=self._TEMPLATE) + template.environment_config.execution_config.subnetwork_uri = ( + "template-subnet" + ) + template.environment_config.execution_config.ttl = {"seconds": 3600} + mock_template_client.return_value.get_session_template.return_value = ( + template + ) + + with mock.patch.dict(os.environ, self._EXEC_ENV): + session_config = ( + self._builder() + .sessionTemplate("test_template") + ._get_session_config() + ) + + mock_template_client.return_value.get_session_template.assert_called_once_with( + GetSessionTemplateRequest(name=self._TEMPLATE) + ) + self.assertEqual(session_config.session_template, "test_template") + exec_config = session_config.environment_config.execution_config + self.assertEqual(exec_config.service_account, "env-acc@example.com") + self.assertEqual( + exec_config.authentication_config.user_workload_authentication_type, + AuthenticationConfig.AuthenticationType.SERVICE_ACCOUNT, + ) + self.assertEqual(exec_config.subnetwork_uri, "") + self.assertFalse(exec_config.ttl) + self.assertEqual(exec_config.idle_ttl, datetime.timedelta(seconds=1800)) + if __name__ == "__main__": unittest.main() From fc932a72966a1401af2d63088e81fd120cc0c053 Mon Sep 17 00:00:00 2001 From: Dave Borowitz Date: Thu, 29 Oct 2026 10:53:50 -0700 Subject: [PATCH 2/2] feat: add runtimeProfile builder method as alias for sessionTemplate Add ManagedSparkSession.Builder.runtimeProfile(profile) to align with the Runtime Profile terminology in Cloud Console and public documentation, while keeping sessionTemplate(profile) as a supported alias. --- README.md | 5 +++-- google/cloud/managed_spark_connect/session.py | 11 +++++++++-- tests/unit/test_session.py | 9 +++++++++ 3 files changed, 21 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 97508091..5a39d858 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ extra instead, uninstall `pyspark` too if nothing else needs it. ## Setup This client requires permissions to -manage [Managed Spark Sessions and Session Templates](https://cloud.google.com/dataproc-serverless/docs/concepts/iam). +manage [Managed Spark Sessions and Runtime Profiles (formerly known as Session Templates)](https://cloud.google.com/dataproc-serverless/docs/concepts/iam). If you are running the client outside of Google Cloud, you need to provide authentication credentials. Set the `GOOGLE_APPLICATION_CREDENTIALS` environment @@ -111,9 +111,10 @@ The `ManagedSparkSession.builder` provides a fluent API to configure the session | `labels(labels)` | Adds multiple labels to the session. | | `location(location)` | Sets the Google Cloud region. | | `projectId(project_id)` | Sets the Google Cloud project ID. | +| `runtimeProfile(profile)` | Sets the Runtime Profile to use. | | `runtimeVersion(version)` | Sets the Managed Spark runtime version (e.g., "3.0"). | | `serviceAccount(account)` | Sets the service account for the session. | -| `sessionTemplate(profile)` | Sets the Session Template to use. | +| `sessionTemplate(profile)` | Alias for `runtimeProfile(profile)`. | | `subnetwork(subnet)` | Sets the subnetwork URI for the session. | | `ttl(duration)` | Sets the time-to-live (TTL) for the session using a `datetime.timedelta` object. | | `instrumentation(inst)` | Records what the library does on your behalf; see [Instrumentation (opt-in)](#instrumentation-opt-in). | diff --git a/google/cloud/managed_spark_connect/session.py b/google/cloud/managed_spark_connect/session.py index b5951768..09312647 100644 --- a/google/cloud/managed_spark_connect/session.py +++ b/google/cloud/managed_spark_connect/session.py @@ -593,11 +593,18 @@ def idleTtlSeconds(self, seconds: int): } return self - def sessionTemplate(self, profile: str): - """Set the Session Template to use for the session.""" + def runtimeProfile(self, profile: str): + """Set the Runtime Profile to use for the session.""" self.session_config.session_template = profile return self + def sessionTemplate(self, profile: str): + """Set the Runtime Profile (formerly known as Session Template) to use for the session. + + Alias for :meth:`runtimeProfile`. + """ + return self.runtimeProfile(profile) + def label(self, key: str, value: str): """Add a single label to the session.""" return self.labels({key: value}) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index f5ecb75b..b4eea4c1 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -3811,6 +3811,15 @@ def test_short_session_template_name_qualifies_get_session_template_request( self.assertFalse(exec_config.ttl) self.assertEqual(exec_config.idle_ttl, datetime.timedelta(seconds=1800)) + def test_runtime_profile_and_session_template_alias(self): + """runtimeProfile and sessionTemplate both set session_config.session_template and return the builder.""" + builder = self._builder() + self.assertIs(builder.runtimeProfile("profile-a"), builder) + self.assertEqual(builder.session_config.session_template, "profile-a") + + self.assertIs(builder.sessionTemplate("profile-b"), builder) + self.assertEqual(builder.session_config.session_template, "profile-b") + if __name__ == "__main__": unittest.main()