From fa6e13d5006caadb36899b4e2a24ca82b7f11b17 Mon Sep 17 00:00:00 2001
From: abdelmegahed <131036743+abdelmegahed@users.noreply.github.com>
Date: Wed, 17 May 2023 15:57:28 -0400
Subject: [PATCH] fix: handle case when expirationMs is None (#1553)

* hotfix: handle case when expirationMs is None

* Add test for unsetting table exp

* Update tests/unit/test_table.py

* Update exp_resource for the unsetting_exp test

---------

Co-authored-by: Tim Swast <swast@google.com>
---
 google/cloud/bigquery/table.py |  6 +++++-
 tests/unit/test_table.py       | 19 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/google/cloud/bigquery/table.py b/google/cloud/bigquery/table.py
index a34e5dc25..bf4a90317 100644
--- a/google/cloud/bigquery/table.py
+++ b/google/cloud/bigquery/table.py
@@ -687,7 +687,11 @@ def partition_expiration(self, value):
 
         if self.time_partitioning is None:
             self._properties[api_field] = {"type": TimePartitioningType.DAY}
-        self._properties[api_field]["expirationMs"] = str(value)
+
+        if value is None:
+            self._properties[api_field]["expirationMs"] = None
+        else:
+            self._properties[api_field]["expirationMs"] = str(value)
 
     @property
     def clustering_fields(self):
diff --git a/tests/unit/test_table.py b/tests/unit/test_table.py
index 53db635fa..a221bc89e 100644
--- a/tests/unit/test_table.py
+++ b/tests/unit/test_table.py
@@ -1190,6 +1190,25 @@ def test_to_api_repr_w_custom_field(self):
         }
         self.assertEqual(resource, exp_resource)
 
+    def test_to_api_repr_w_unsetting_expiration(self):
+        from google.cloud.bigquery.table import TimePartitioningType
+
+        dataset = DatasetReference(self.PROJECT, self.DS_ID)
+        table_ref = dataset.table(self.TABLE_NAME)
+        table = self._make_one(table_ref)
+        table.partition_expiration = None
+        resource = table.to_api_repr()
+
+        exp_resource = {
+            "tableReference": table_ref.to_api_repr(),
+            "labels": {},
+            "timePartitioning": {
+                "expirationMs": None,
+                "type": TimePartitioningType.DAY,
+            },
+        }
+        self.assertEqual(resource, exp_resource)
+
     def test__build_resource_w_custom_field(self):
         dataset = DatasetReference(self.PROJECT, self.DS_ID)
         table_ref = dataset.table(self.TABLE_NAME)