From 781096eb0c71fa540357e0e6e3b51104ad6469ee Mon Sep 17 00:00:00 2001
From: Anupam Saini <4581797+anupam-saini@users.noreply.github.com>
Date: Thu, 14 Mar 2024 07:06:45 -0400
Subject: [PATCH] Add `table_exists` method to Catalog (#512)

* Add `table_exist` method to Catalog

* `table_exist` -> `table_exists`

* Add `table_exists` for RestCatalog

* Add doc
---
 mkdocs/docs/api.md            | 10 ++++++++++
 pyiceberg/catalog/__init__.py |  7 +++++++
 pyiceberg/catalog/rest.py     |  8 ++++++++
 tests/catalog/test_base.py    | 11 +++++++++++
 tests/catalog/test_rest.py    | 20 ++++++++++++++++++++
 5 files changed, 56 insertions(+)

diff --git a/mkdocs/docs/api.md b/mkdocs/docs/api.md
index 724a45c52f..35b271ae9e 100644
--- a/mkdocs/docs/api.md
+++ b/mkdocs/docs/api.md
@@ -194,6 +194,16 @@ static_table = StaticTable.from_metadata(
 
 The static-table is considered read-only.
 
+## Check if a table exists
+
+To check whether the `bids` table exists:
+
+```python
+catalog.table_exists("docs_example.bids")
+```
+
+Returns `True` if the table already exists.
+
 ## Write support
 
 With PyIceberg 0.6.0 write support is added through Arrow. Let's consider an Arrow Table:
diff --git a/pyiceberg/catalog/__init__.py b/pyiceberg/catalog/__init__.py
index 4f7b5a3292..f2b46fcde7 100644
--- a/pyiceberg/catalog/__init__.py
+++ b/pyiceberg/catalog/__init__.py
@@ -646,6 +646,13 @@ def purge_table(self, identifier: Union[str, Identifier]) -> None:
         delete_files(io, prev_metadata_files, PREVIOUS_METADATA)
         delete_files(io, {table.metadata_location}, METADATA)
 
+    def table_exists(self, identifier: Union[str, Identifier]) -> bool:
+        try:
+            self.load_table(identifier)
+            return True
+        except NoSuchTableError:
+            return False
+
     @staticmethod
     def _write_metadata(metadata: TableMetadata, io: FileIO, metadata_path: str) -> None:
         ToOutputFile.table_metadata(metadata, io.new_output(metadata_path))
diff --git a/pyiceberg/catalog/rest.py b/pyiceberg/catalog/rest.py
index f67fe2c1fd..9f0d054493 100644
--- a/pyiceberg/catalog/rest.py
+++ b/pyiceberg/catalog/rest.py
@@ -717,3 +717,11 @@ def update_namespace_properties(
             updated=parsed_response.updated,
             missing=parsed_response.missing,
         )
+
+    @retry(**_RETRY_ARGS)
+    def table_exists(self, identifier: Union[str, Identifier]) -> bool:
+        identifier_tuple = self.identifier_to_tuple_without_catalog(identifier)
+        response = self._session.head(
+            self.url(Endpoints.load_table, prefixed=True, **self._split_identifier_for_path(identifier_tuple))
+        )
+        return response.status_code == 200
diff --git a/tests/catalog/test_base.py b/tests/catalog/test_base.py
index 44c36a7d2d..5f78eb3bc4 100644
--- a/tests/catalog/test_base.py
+++ b/tests/catalog/test_base.py
@@ -413,6 +413,17 @@ def test_table_raises_error_on_table_not_found(catalog: InMemoryCatalog) -> None
         catalog.load_table(TEST_TABLE_IDENTIFIER)
 
 
+def test_table_exists(catalog: InMemoryCatalog) -> None:
+    # Given
+    given_catalog_has_a_table(catalog)
+    # Then
+    assert catalog.table_exists(TEST_TABLE_IDENTIFIER)
+
+
+def test_table_exists_on_table_not_found(catalog: InMemoryCatalog) -> None:
+    assert not catalog.table_exists(TEST_TABLE_IDENTIFIER)
+
+
 def test_drop_table(catalog: InMemoryCatalog) -> None:
     # Given
     given_catalog_has_a_table(catalog)
diff --git a/tests/catalog/test_rest.py b/tests/catalog/test_rest.py
index 850e5f0180..4956fffe6c 100644
--- a/tests/catalog/test_rest.py
+++ b/tests/catalog/test_rest.py
@@ -644,6 +644,26 @@ def test_load_table_404(rest_mock: Mocker) -> None:
     assert "Table does not exist" in str(e.value)
 
 
+def test_table_exist_200(rest_mock: Mocker) -> None:
+    rest_mock.head(
+        f"{TEST_URI}v1/namespaces/fokko/tables/table",
+        status_code=200,
+        request_headers=TEST_HEADERS,
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+    assert catalog.table_exists(("fokko", "table"))
+
+
+def test_table_exist_500(rest_mock: Mocker) -> None:
+    rest_mock.head(
+        f"{TEST_URI}v1/namespaces/fokko/tables/table",
+        status_code=500,
+        request_headers=TEST_HEADERS,
+    )
+    catalog = RestCatalog("rest", uri=TEST_URI, token=TEST_TOKEN)
+    assert not catalog.table_exists(("fokko", "table"))
+
+
 def test_drop_table_404(rest_mock: Mocker) -> None:
     rest_mock.delete(
         f"{TEST_URI}v1/namespaces/fokko/tables/does_not_exists",