Skip to content

Commit 5cb1c63

Browse files
nija-atHyukjinKwon
authored andcommitted
[SPARK-43172][CONNECT] Expose host and token from spark connect client
### 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]>
1 parent 1c057f5 commit 5cb1c63

File tree

2 files changed

+23
-0
lines changed

2 files changed

+23
-0
lines changed

python/pyspark/sql/connect/client.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -784,6 +784,21 @@ def close(self) -> None:
784784
"""
785785
self._channel.close()
786786

787+
@property
788+
def host(self) -> str:
789+
"""
790+
The hostname where this client intends to connect.
791+
"""
792+
return self._builder.host
793+
794+
@property
795+
def token(self) -> Optional[str]:
796+
"""
797+
The authentication bearer token during connection.
798+
If authentication is not using a bearer token, None will be returned.
799+
"""
800+
return self._builder._token
801+
787802
def _execute_plan_request_with_metadata(self) -> pb2.ExecutePlanRequest:
788803
req = pb2.ExecutePlanRequest()
789804
req.session_id = self._session_id

python/pyspark/sql/tests/connect/test_client.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -50,6 +50,14 @@ def test_user_agent_default(self):
5050
self.assertIsNotNone(mock.req, "ExecutePlan API was not called when expected")
5151
self.assertEqual(mock.req.client_type, "_SPARK_CONNECT_PYTHON")
5252

53+
def test_properties(self):
54+
client = SparkConnectClient("sc://foo/;token=bar")
55+
self.assertEqual(client.token, "bar")
56+
self.assertEqual(client.host, "foo")
57+
58+
client = SparkConnectClient("sc://foo/")
59+
self.assertIsNone(client.token)
60+
5361

5462
class MockService:
5563
# Simplest mock of the SparkConnectService.

0 commit comments

Comments
 (0)