diff --git a/python/delta_sharing/delta_sharing.py b/python/delta_sharing/delta_sharing.py index 4c061bf2d..911bb2ebb 100644 --- a/python/delta_sharing/delta_sharing.py +++ b/python/delta_sharing/delta_sharing.py @@ -241,6 +241,19 @@ def load_table_changes_as_pandas( )) +def query_cloud_sharing_token(url: str): + profile_json, share, schema, table = _parse_url(url) + profile = DeltaSharingProfile.read_from_file(profile_json) + return DeltaSharingReader( + table=Table(name=table, share=share, schema=schema), + rest_client=DataSharingRestClient(profile), + jsonPredicateHints=None, + limit=None, + version=None, + timestamp=None + ).query_cloud_sharing_token() + + class SharingClient: """ A Delta Sharing client to query shares/schemas/tables from a Delta Sharing Server. diff --git a/python/delta_sharing/reader.py b/python/delta_sharing/reader.py index 5e995320d..9786620df 100644 --- a/python/delta_sharing/reader.py +++ b/python/delta_sharing/reader.py @@ -89,6 +89,9 @@ def limit(self, limit: Optional[int]) -> "DeltaSharingReader": timestamp=self._timestamp ) + def query_cloud_sharing_token(self) -> str: + return self._rest_client.query_cloud_sharing_token(self._table) + def __to_pandas_kernel(self): """ This function calls delta-kernel-rust python wrapper to load a df for a table diff --git a/python/delta_sharing/rest_client.py b/python/delta_sharing/rest_client.py index 002af9400..2dd537357 100644 --- a/python/delta_sharing/rest_client.py +++ b/python/delta_sharing/rest_client.py @@ -473,6 +473,14 @@ def list_table_changes(self, table: Table, cdfOptions: CdfOptions) -> ListTableC actions=actions, ) + @retry_with_exponential_backoff + def query_cloud_sharing_token(self, table: Table) -> str: + query_str = f"/shares/{table.share}/schemas/{table.schema}/tables/{table.name}/getCloudToken" + with self._get_internal(query_str, return_headers=True) as values: + headers = values[0] + cloud_token = headers.get("Cloud-Token-Header") + return cloud_token + def close(self): self._session.close() diff --git a/python/delta_sharing/tests/test_delta_sharing.py b/python/delta_sharing/tests/test_delta_sharing.py index d52d96074..243e093c2 100644 --- a/python/delta_sharing/tests/test_delta_sharing.py +++ b/python/delta_sharing/tests/test_delta_sharing.py @@ -30,6 +30,7 @@ load_table_changes_as_spark, load_table_changes_as_pandas, _parse_url, + query_cloud_sharing_token, ) from delta_sharing.protocol import Format, Metadata, Protocol, Schema, Share, Table from delta_sharing.rest_client import ( @@ -711,6 +712,31 @@ def test_load_as_pandas_success_dv_and_cm( pd.testing.assert_frame_equal(pdf, expected) +@pytest.mark.skipif(not ENABLE_INTEGRATION, reason=SKIP_MESSAGE) +@pytest.mark.parametrize( + "fragments,limit,version,expected", + [ + pytest.param( + "share8.default.dv_and_cm_table", + 0, + None, + "EXPECTED_CLOUD_TOKEN", + id="cloud token sharing", + ) + ], +) +def test_cloud_token_sharing( + profile_path: str, + fragments: str, + limit: Optional[int], + version: Optional[int], + expected: str +): + cloud_token = query_cloud_sharing_token(f"{profile_path}#{fragments}") + print("cloud token returned is: " + str(cloud_token)) + assert cloud_token == expected + + # We will test predicates with the table share8.default.cdf_table_with_partition # This table is partitioned by birthday column of type date. # There are two partitions: 2020-02-02, and 2020-01-01. diff --git a/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala b/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala index 9a69148a5..a7b71ec36 100644 --- a/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala +++ b/server/src/main/scala/io/delta/sharing/server/DeltaSharingService.scala @@ -261,6 +261,16 @@ class DeltaSharingService(serverConfig: ServerConfig) { ResponseHeaders.builder(200).set(DELTA_TABLE_VERSION_HEADER, version.toString) } + @Get("/shares/{share}/schemas/{schema}/tables/{table}/getCloudToken") + def getCloudToken( + req: HttpRequest, + @Param("share") share: String, + @Param("schema") schema: String, + @Param("table") table: String): HttpResponse = processRequest { + val headers = ResponseHeaders.builder(200).set("Cloud-Token-Header", "Cloud Token Stub").build() + HttpResponse.of(headers) + } + // TODO: deprecate HEAD request in favor of the GET request @Head("/shares/{share}/schemas/{schema}/tables/{table}") @Get("/shares/{share}/schemas/{schema}/tables/{table}/version")