Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions google/cloud/dataproc_spark_connect/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -592,6 +592,16 @@ def getOrCreate(self) -> "DataprocSparkSession":
session = PySparkSQLSession.builder.getOrCreate()
return session # type: ignore

if self._project_id is None:
raise DataprocSparkConnectException(
f"Error while creating Dataproc Session: project ID is not set"
)

if self._region is None:
raise DataprocSparkConnectException(
f"Error while creating Dataproc Session: location is not set"
)

# Handle custom session ID by setting it early and letting existing logic handle it
if self._custom_session_id:
self._handle_custom_session_id()
Expand Down
22 changes: 22 additions & 0 deletions tests/unit/test_session.py
Original file line number Diff line number Diff line change
Expand Up @@ -1471,9 +1471,31 @@ def test_create_session_with_valid_notebook_id(
)
self.stopSession(mock_session_controller_client_instance, session)

def test_create_session_without_project_id(self):
"""Tests that an exception is raised when project ID is not provided."""
os.environ.clear()
try:
DataprocSparkSession.builder.location("test-region").getOrCreate()
except DataprocSparkConnectException as e:
self.assertIn("project ID is not set", str(e))

def test_create_session_without_location(self):
"""Tests that an exception is raised when location is not provided."""
os.environ.clear()
try:
DataprocSparkSession.builder.projectId("test-project").getOrCreate()
except DataprocSparkConnectException as e:
self.assertIn("location is not set", str(e))


class DataprocSparkConnectClientTest(unittest.TestCase):

def setUp(self):
self.original_environment = dict(os.environ)
os.environ.clear()
os.environ["GOOGLE_CLOUD_PROJECT"] = "test-project"
os.environ["GOOGLE_CLOUD_REGION"] = "test-region"

@staticmethod
def stopSession(mock_session_controller_client_instance, session):
session_response = Session()
Expand Down