diff --git a/pyiceberg/catalog/rest/__init__.py b/pyiceberg/catalog/rest/__init__.py index 95db70c2a7..ddd2f0577b 100644 --- a/pyiceberg/catalog/rest/__init__.py +++ b/pyiceberg/catalog/rest/__init__.py @@ -1169,6 +1169,7 @@ def _response_to_table(self, identifier_tuple: tuple[str, ...], table_response: ), catalog=self, config=table_response.config, + labels=table_response.labels, ) def _response_to_staged_table(self, identifier_tuple: tuple[str, ...], table_response: TableResponse) -> StagedTable: diff --git a/pyiceberg/table/__init__.py b/pyiceberg/table/__init__.py index df7aaf824f..bc4cff1d5a 100644 --- a/pyiceberg/table/__init__.py +++ b/pyiceberg/table/__init__.py @@ -59,6 +59,7 @@ from pyiceberg.schema import Schema from pyiceberg.table.delete_file_index import DeleteFileIndex from pyiceberg.table.inspect import InspectTable +from pyiceberg.table.labels import Labels from pyiceberg.table.locations import LocationProvider, load_location_provider from pyiceberg.table.maintenance import MaintenanceTable from pyiceberg.table.metadata import INITIAL_SEQUENCE_NUMBER, TableMetadata @@ -1345,6 +1346,8 @@ class Table: io: FileIO catalog: Catalog config: dict[str, str] + # Catalog-provided labels, set when the catalog loads the table. They are not table state, so refresh() and commits keep them. + labels: Labels def __init__( self, @@ -1355,6 +1358,7 @@ def __init__( catalog: Catalog, config: dict[str, str] = EMPTY_DICT, kms_client: KeyManagementClient | None = None, + labels: Labels | None = None, ) -> None: self._identifier = identifier self.metadata = metadata @@ -1363,6 +1367,7 @@ def __init__( self.catalog = catalog self.config = config self._kms_client = kms_client + self.labels = labels if labels is not None else Labels() self._encryption_manager: _CachedEncryptionManager | None = None @property diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py index 632dc7b5fb..37508e9a64 100644 --- a/tests/catalog/test_rest.py +++ b/tests/catalog/test_rest.py @@ -19,6 +19,7 @@ import base64 import os +import pickle from collections.abc import Callable, Iterator from contextlib import contextmanager from typing import Any, cast @@ -3734,3 +3735,43 @@ def test_load_responses_parse_labels( assert response_type.model_validate(payload).labels == Labels() for invalid in invalid_values: assert response_type.model_validate({**payload, "labels": invalid}).labels == Labels() + + +def test_table_labels_from_rest_responses( + rest_mock: Mocker, example_table_metadata_with_snapshot_v1_rest_json: dict[str, Any], table_schema_simple: Schema +) -> None: + response = { + **example_table_metadata_with_snapshot_v1_rest_json, + "labels": {"object-labels": {"owner": "team-a"}, "fields": [{"field-id": 1, "labels": {"classification": "pii"}}]}, + } + rest_mock.get(f"{TEST_URI}v1/namespaces/fokko/tables/table", json=response, status_code=200, request_headers=TEST_HEADERS) + rest_mock.post(f"{TEST_URI}v1/namespaces/fokko/tables", json=response, status_code=200, request_headers=TEST_HEADERS) + rest_mock.post(f"{TEST_URI}v1/namespaces/fokko/register", json=response, status_code=200, request_headers=TEST_HEADERS) + catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN) + expected = Labels(object_labels={"owner": "team-a"}, fields=[FieldLabel(field_id=1, labels={"classification": "pii"})]) + + table = catalog.load_table(("fokko", "table")) + created = catalog.create_table(("fokko", "table"), schema=table_schema_simple) + registered = catalog.register_table(("fokko", "table"), metadata_location="s3://warehouse/database/table/metadata.json") + for actual in (table, created, registered): + assert actual.labels == expected + + # A pickled table keeps its labels. Java does not serialize them, because BaseTable.labels is transient. + assert pickle.loads(pickle.dumps(table)).labels == expected + + # Labels are not table state, so equality ignores them and refresh() keeps the labels from the first load. + assert table == Table( + identifier=("fokko", "table"), + metadata_location=example_table_metadata_with_snapshot_v1_rest_json["metadata-location"], + metadata=TableMetadataV1(**example_table_metadata_with_snapshot_v1_rest_json["metadata"]), + io=load_file_io(), + catalog=catalog, + ) + rest_mock.get( + f"{TEST_URI}v1/namespaces/fokko/tables/table", + json={**example_table_metadata_with_snapshot_v1_rest_json, "labels": {"object-labels": {"owner": "team-b"}}}, + status_code=200, + request_headers=TEST_HEADERS, + ) + table.refresh() + assert table.labels == expected diff --git a/tests/table/test_init.py b/tests/table/test_init.py index db6dc00e39..c6b87e7335 100644 --- a/tests/table/test_init.py +++ b/tests/table/test_init.py @@ -46,6 +46,7 @@ Table, TableIdentifier, ) +from pyiceberg.table.labels import Labels from pyiceberg.table.metadata import EncryptedKey, TableMetadataUtil, TableMetadataV1, TableMetadataV2, _generate_snapshot_id from pyiceberg.table.refs import MAIN_BRANCH, SnapshotRef, SnapshotRefType from pyiceberg.table.snapshots import ( @@ -193,6 +194,10 @@ def test_location(table_v2: Table) -> None: assert table_v2.location() == "s3://bucket/test/location" +def test_labels_default_to_empty(table_v2: Table) -> None: + assert table_v2.labels == Labels() + + def test_current_snapshot(table_v2: Table) -> None: assert table_v2.current_snapshot() == Snapshot( snapshot_id=3055729675574597004,