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 <[email protected]>
googleapis · May 17, 2023 · fa6e13d · fa6e13d
1 parent 075aa66
commit fa6e13d
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 1 deletion.
diff --git 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
@@ -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)