From 5ed9ccee204b7cf8e96cb0e050f6830c05f3b4fd Mon Sep 17 00:00:00 2001
From: Chalmer Lowe <chalmerlowe@google.com>
Date: Thu, 11 Apr 2024 15:33:30 -0400
Subject: [PATCH] feat: Add compression option ZSTD. (#1890)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Add ZSTD to compression types

* feat: adds tests re Compression types

* revise datatype from Enum to object

* adds license text and docstring

* change object back to enum datatype

* updates compression object comparison

* updates Compression class

* jsonify and sort the input and output for testing

* Update tests/unit/job/test_extract.py

* moved json import statement

* removed enums test and file

---------

Co-authored-by: Ethan Steinberg <ethan.steinberg@gmail.com>
Co-authored-by: Tim Sweña (Swast) <swast@google.com>
---
 google/cloud/bigquery/enums.py |  5 ++++-
 tests/unit/job/test_extract.py | 12 +++++++++---
 2 files changed, 13 insertions(+), 4 deletions(-)

diff --git a/google/cloud/bigquery/enums.py b/google/cloud/bigquery/enums.py
index d75037ad1..1abe28381 100644
--- a/google/cloud/bigquery/enums.py
+++ b/google/cloud/bigquery/enums.py
@@ -22,7 +22,7 @@ class AutoRowIDs(enum.Enum):
     GENERATE_UUID = enum.auto()
 
 
-class Compression(object):
+class Compression(str, enum.Enum):
     """The compression type to use for exported files. The default value is
     :attr:`NONE`.
 
@@ -39,6 +39,9 @@ class Compression(object):
     SNAPPY = "SNAPPY"
     """Specifies SNAPPY format."""
 
+    ZSTD = "ZSTD"
+    """Specifies ZSTD format."""
+
     NONE = "NONE"
     """Specifies no compression."""
 
diff --git a/tests/unit/job/test_extract.py b/tests/unit/job/test_extract.py
index 76ee72f28..ee0d67d68 100644
--- a/tests/unit/job/test_extract.py
+++ b/tests/unit/job/test_extract.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 from unittest import mock
 
 from ..helpers import make_connection
@@ -45,9 +46,8 @@ def test_to_api_repr(self):
         config.print_header = False
         config._properties["extract"]["someNewField"] = "some-value"
         config.use_avro_logical_types = True
-        resource = config.to_api_repr()
-        self.assertEqual(
-            resource,
+        resource = json.dumps(config.to_api_repr(), sort_keys=True)
+        expected = json.dumps(
             {
                 "extract": {
                     "compression": "SNAPPY",
@@ -58,6 +58,12 @@ def test_to_api_repr(self):
                     "useAvroLogicalTypes": True,
                 }
             },
+            sort_keys=True,
+        )
+
+        self.assertEqual(
+            resource,
+            expected,
         )
 
     def test_from_api_repr(self):