Treat an empty project or region as unset

README.md

@@ -69,6 +69,30 @@ in your code using the builder API:
    spark = ManagedSparkSession.builder.projectId('my-project').location('us-central1').sessionConfig(session_config).getOrCreate()
    ```
 
+5. To start from a [Session Template](https://cloud.google.com/dataproc-serverless/docs/concepts/session-templates),
+   pass its ID. The template is resolved against the project and region you
+   configured, so you don't have to repeat them:
+
+   ```python
+   from google.cloud.managed_spark_connect import ManagedSparkSession
+   spark = (
+       ManagedSparkSession.builder
+           .projectId('my-project')
+           .location('us-central1')
+           .sessionTemplate('my-template')
+           .getOrCreate()
+   )
+   ```
+
+   A full resource name is still accepted, and is required when the template
+   lives in a different project or region than the session:
+
+   ```python
+   spark = ManagedSparkSession.builder.sessionTemplate(
+       'projects/other-project/locations/us-east1/sessionTemplates/my-template'
+   ).projectId('my-project').location('us-central1').getOrCreate()
+   ```
+
 ### Builder Configuration
 
 The `ManagedSparkSession.builder` provides a fluent API to configure the session. Below is a list of available methods:
@@ -85,7 +109,7 @@ The `ManagedSparkSession.builder` provides a fluent API to configure the session
 | `projectId(project_id)` | Sets the Google Cloud project ID. |
 | `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)` | Sets the Session Template to use. Accepts a bare template ID or a full resource name. |
 | `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. |
 

google/cloud/managed_spark_connect/session.py

@@ -108,6 +108,46 @@ def _is_valid_session_id(session_id: str) -> bool:
     return bool(re.match(pattern, session_id))
 
 
+def _qualify_session_template(
+    template: str, project_id: Optional[str], region: Optional[str]
+) -> str:
+    """
+    Resolve a session template to a fully qualified resource name.
+
+    The Sessions API accepts only resource names that include the project and
+    location, so a bare template ID is expanded against the session's own
+    project and region. Values that already contain a path separator (the
+    ``projects/...`` resource name and the ``https://...`` URL forms) carry
+    their own project and location, and are returned unchanged.
+
+    Raises:
+        ManagedSparkConnectException: If a bare template ID was given but the
+            project or region needed to resolve it is missing.
+    """
+    if not template or "/" in template:
+        return template
+
+    missing = [
+        field
+        for field, value in (("project ID", project_id), ("location", region))
+        if not value
+    ]
+    if missing:
+        raise ManagedSparkConnectException(
+            f"Error while creating Managed Spark Session: cannot resolve the"
+            f" '{template}' session template because the"
+            f" {' and '.join(missing)}"
+            f" {'are' if len(missing) > 1 else 'is'} not set."
+            f" Either set the project and location, or pass the template's"
+            f" full resource name"
+            f" (projects/<project>/locations/<location>/sessionTemplates/{template})."
+        )
+
+    return (
+        f"projects/{project_id}/locations/{region}/sessionTemplates/{template}"
+    )
+
+
 class ManagedSparkSession(SparkSession):
     """The entry point to programming Spark with the Dataset and DataFrame API.
 
@@ -258,7 +298,18 @@ def idleTtlSeconds(self, seconds: int):
             return self
 
         def sessionTemplate(self, profile: str):
-            """Set the Session Template to use for the session."""
+            """Set the Session Template to use for the session.
+
+            Accepts either a bare template ID, which is resolved against the
+            session's project and region, or a fully qualified resource name.
+
+            Args:
+                profile: The template ID (``my-template``) or resource name
+                    (``projects/p/locations/r/sessionTemplates/my-template``)
+
+            Returns:
+                This Builder instance for method chaining
+            """
             self.session_config.session_template = profile
             return self
 
@@ -603,12 +654,14 @@ def getOrCreate(self) -> "ManagedSparkSession":
                     session = PySparkSQLSession.builder.getOrCreate()
                     return session  # type: ignore
 
-                if self._project_id is None:
+                # Falsy rather than None: an environment variable that is set
+                # but empty reads back as "", which is just as unusable.
+                if not self._project_id:
                     raise ManagedSparkConnectException(
                         f"Error while creating Managed Spark Session: project ID is not set"
                     )
 
-                if self._region is None:
+                if not self._region:
                     raise ManagedSparkConnectException(
                         f"Error while creating Managed Spark Session: location is not set"
                     )
@@ -650,6 +703,13 @@ def _get_session_config(self):
             for k, v in self._options.items():
                 session_config.runtime_config.properties[k] = v
             session_config.spark_connect_session = sessions.SparkConnectConfig()
+            # Resolved here rather than in sessionTemplate() so that the
+            # template may be set before the project and region are.
+            session_config.session_template = _qualify_session_template(
+                session_config.session_template,
+                self._project_id,
+                self._region,
+            )
             if not session_config.runtime_config.version:
                 session_config.runtime_config.version = (
                     ManagedSparkSession._DEFAULT_RUNTIME_VERSION

tests/unit/test_session.py

@@ -1486,6 +1486,24 @@ def test_create_session_without_location(self):
         except ManagedSparkConnectException as e:
             self.assertIn("location is not set", str(e))
 
+    def test_create_session_with_empty_project_id(self):
+        """Tests that a set-but-empty project ID is treated as not provided."""
+        os.environ.clear()
+        os.environ["GOOGLE_CLOUD_PROJECT"] = ""
+        os.environ["GOOGLE_CLOUD_REGION"] = "test-region"
+        with self.assertRaises(ManagedSparkConnectException) as context:
+            ManagedSparkSession.builder.getOrCreate()
+        self.assertIn("project ID is not set", str(context.exception))
+
+    def test_create_session_with_empty_location(self):
+        """Tests that a set-but-empty location is treated as not provided."""
+        os.environ.clear()
+        os.environ["GOOGLE_CLOUD_PROJECT"] = "test-project"
+        os.environ["GOOGLE_CLOUD_REGION"] = ""
+        with self.assertRaises(ManagedSparkConnectException) as context:
+            ManagedSparkSession.builder.getOrCreate()
+        self.assertIn("location is not set", str(context.exception))
+
     def test_create_session_without_application_default_credentials(self):
         """Tests that an exception is raised when application default credentials is not provided."""
         os.environ.clear()
@@ -2064,6 +2082,59 @@ def test_builder_pattern_ttl_with_timedelta(
             )
             self.stopSession(mock_session_controller_client_instance, session)
 
+    @mock.patch("google.auth.default")
+    @mock.patch("google.cloud.dataproc_v1.SessionControllerClient")
+    @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config")
+    @mock.patch(
+        "google.cloud.managed_spark_connect.ManagedSparkSession.Builder.generate_session_id"
+    )
+    @mock.patch(
+        "google.cloud.managed_spark_connect.session.is_s8s_session_active"
+    )
+    def test_create_session_with_bare_session_template_id(
+        self,
+        mock_is_s8s_session_active,
+        mock_session_id,
+        mock_client_config,
+        mock_session_controller_client,
+        mock_credentials,
+    ):
+        """A bare template ID is expanded before the session is created."""
+        session = None
+        mock_session_controller_client_instance = (
+            self._setup_session_creation_mocks(
+                mock_is_s8s_session_active,
+                mock_session_id,
+                mock_client_config,
+                mock_session_controller_client,
+                mock_credentials,
+            )
+        )
+
+        try:
+            session = (
+                ManagedSparkSession.builder.projectId("test-project")
+                .location("us-central1")
+                .sessionTemplate("test-template")
+                .getOrCreate()
+            )
+
+            create_session_request = mock_session_controller_client_instance.create_session.call_args[
+                0
+            ][
+                0
+            ]
+            self.assertEqual(
+                create_session_request.session.session_template,
+                "projects/test-project/locations/us-central1/sessionTemplates/test-template",
+            )
+
+        finally:
+            mock_session_controller_client_instance.terminate_session.return_value = (
+                mock.Mock()
+            )
+            self.stopSession(mock_session_controller_client_instance, session)
+
     @mock.patch("google.auth.default")
     @mock.patch("google.cloud.dataproc_v1.SessionControllerClient")
     @mock.patch("pyspark.sql.connect.client.SparkConnectClient.config")
@@ -2641,5 +2712,142 @@ def test_session_skip_terminated(self, mock_session_controller_client):
         mock_client.get_session.assert_called_once()
 
 
+class SessionTemplateExpansionTests(unittest.TestCase):
+    """Test cases for resolving bare session template IDs to resource names."""
+
+    _EXPANDED = (
+        "projects/test-project/locations/test-region/sessionTemplates/tmpl"
+    )
+    _QUALIFIED = (
+        "projects/other-project/locations/other-region/sessionTemplates/tmpl"
+    )
+    _URL = (
+        "https://www.googleapis.com/compute/v1/projects/other-project"
+        "/locations/other-region/sessionTemplates/tmpl"
+    )
+
+    def setUp(self):
+        self.original_environment = dict(os.environ)
+        os.environ.clear()
+
+    def tearDown(self):
+        os.environ.clear()
+        os.environ.update(self.original_environment)
+
+    @staticmethod
+    def _builder():
+        builder = ManagedSparkSession.Builder()
+        builder._project_id = "test-project"
+        builder._region = "test-region"
+        return builder
+
+    def test_bare_template_id_is_expanded(self):
+        """A bare template ID resolves against the session project and region."""
+        builder = self._builder().sessionTemplate("tmpl")
+        self.assertEqual(
+            builder._get_session_config().session_template, self._EXPANDED
+        )
+
+    def test_resource_name_is_left_unchanged(self):
+        """A fully qualified resource name is passed through untouched."""
+        builder = self._builder().sessionTemplate(self._QUALIFIED)
+        self.assertEqual(
+            builder._get_session_config().session_template, self._QUALIFIED
+        )
+
+    def test_url_is_left_unchanged(self):
+        """The googleapis.com URL form is passed through untouched."""
+        builder = self._builder().sessionTemplate(self._URL)
+        self.assertEqual(
+            builder._get_session_config().session_template, self._URL
+        )
+
+    def test_unset_template_is_left_unset(self):
+        """A session without a template does not get an empty template name."""
+        builder = self._builder()
+        self.assertEqual(builder._get_session_config().session_template, "")
+
+    def test_expansion_is_independent_of_builder_call_order(self):
+        """The template may be set before or after the project and region."""
+        template_first = ManagedSparkSession.Builder()
+        template_first.sessionTemplate("tmpl")
+        template_first.projectId("test-project").location("test-region")
+
+        template_last = ManagedSparkSession.Builder()
+        template_last.projectId("test-project").location("test-region")
+        template_last.sessionTemplate("tmpl")
+
+        self.assertEqual(
+            template_first._get_session_config().session_template,
+            self._EXPANDED,
+        )
+        self.assertEqual(
+            template_last._get_session_config().session_template,
+            self._EXPANDED,
+        )
+
+    def test_bare_template_id_in_session_config_is_expanded(self):
+        """A template set through sessionConfig() is expanded too."""
+        session_config = Session()
+        session_config.session_template = "tmpl"
+        builder = self._builder().sessionConfig(session_config)
+        self.assertEqual(
+            builder._get_session_config().session_template, self._EXPANDED
+        )
+
+    def test_bare_template_id_without_project_is_rejected(self):
+        """A bare template ID cannot be resolved without a project."""
+        builder = self._builder().sessionTemplate("tmpl")
+        builder._project_id = None
+        with self.assertRaises(ManagedSparkConnectException) as context:
+            builder._get_session_config()
+        message = str(context.exception)
+        self.assertIn("'tmpl' session template", message)
+        self.assertIn("project ID is not set", message)
+
+    def test_bare_template_id_without_region_is_rejected(self):
+        """A bare template ID cannot be resolved without a location."""
+        builder = self._builder().sessionTemplate("tmpl")
+        builder._region = None
+        with self.assertRaises(ManagedSparkConnectException) as context:
+            builder._get_session_config()
+        self.assertIn("location is not set", str(context.exception))
+
+    def test_bare_template_id_without_either_names_both(self):
+        """The error names every field that is missing."""
+        builder = self._builder().sessionTemplate("tmpl")
+        builder._project_id = None
+        builder._region = None
+        with self.assertRaises(ManagedSparkConnectException) as context:
+            builder._get_session_config()
+        self.assertIn(
+            "project ID and location are not set", str(context.exception)
+        )
+
+    def test_empty_project_is_rejected(self):
+        """A set-but-empty project is treated as missing, not interpolated."""
+        builder = self._builder().sessionTemplate("tmpl")
+        builder._project_id = ""
+        with self.assertRaises(ManagedSparkConnectException) as context:
+            builder._get_session_config()
+        self.assertIn("project ID is not set", str(context.exception))
+
+    def test_resource_name_needs_no_project_or_region(self):
+        """A full resource name carries its own project and location."""
+        builder = self._builder().sessionTemplate(self._QUALIFIED)
+        builder._project_id = None
+        builder._region = None
+        self.assertEqual(
+            builder._get_session_config().session_template, self._QUALIFIED
+        )
+
+    def test_unset_template_needs_no_project_or_region(self):
+        """A session without a template is unaffected by the guard."""
+        builder = self._builder()
+        builder._project_id = None
+        builder._region = None
+        self.assertEqual(builder._get_session_config().session_template, "")
+
+
 if __name__ == "__main__":
     unittest.main()