Skip to content

Commit

Permalink
[SPARK-43172][CONNECT] Expose host and token from spark connect client
Browse files Browse the repository at this point in the history
### What changes were proposed in this pull request?

Expose the host and bearer token as properties of the
`SparkConnectClient`.

### Why are the changes needed?

This allows for querying the connecting host and bearer token
given an instance of spark connect client.

### Does this PR introduce _any_ user-facing change?

Introduces two new properties `host` and `token` to the
`SparkConnectClient` class

### How was this patch tested?

Unit tests

Closes apache#40836 from nija-at/expose-host-token.

Authored-by: Niranjan Jayakar <[email protected]>
Signed-off-by: Hyukjin Kwon <[email protected]>
  • Loading branch information
nija-at authored and HyukjinKwon committed Apr 18, 2023
1 parent 1c057f5 commit 5cb1c63
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 0 deletions.
15 changes: 15 additions & 0 deletions python/pyspark/sql/connect/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -784,6 +784,21 @@ def close(self) -> None:
"""
self._channel.close()

@property
def host(self) -> str:
"""
The hostname where this client intends to connect.
"""
return self._builder.host

@property
def token(self) -> Optional[str]:
"""
The authentication bearer token during connection.
If authentication is not using a bearer token, None will be returned.
"""
return self._builder._token

def _execute_plan_request_with_metadata(self) -> pb2.ExecutePlanRequest:
req = pb2.ExecutePlanRequest()
req.session_id = self._session_id
Expand Down
8 changes: 8 additions & 0 deletions python/pyspark/sql/tests/connect/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,14 @@ def test_user_agent_default(self):
self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected")
self.assertEqual(mock.req.client_type, "_SPARK_CONNECT_PYTHON")

def test_properties(self):
client = SparkConnectClient("sc://foo/;token=bar")
self.assertEqual(client.token, "bar")
self.assertEqual(client.host, "foo")

client = SparkConnectClient("sc://foo/")
self.assertIsNone(client.token)


class MockService:
# Simplest mock of the SparkConnectService.
Expand Down

0 comments on commit 5cb1c63

Please sign in to comment.