Repository: mozilla/bugbug

Issue:
Remove fallback on downloading xz files
It's added in #524 because the current task have xz as artifacts. Once we make a successful deployment, we can switch to zst fully and remove the xz fallback.


Discussion:
@x249wang do you want to do this cleanup?
@marco-c Sure! Does this involve removing the code for downloading xz files in case the zst files are unavailable? When should this change occur? 
> Does this involve removing the code for downloading xz files in case the zst files are unavailable?

Exactly, and the related code.

> When should this change occur?

It can happen already, as I've deployed the zst change and we have zst datasets and models now (e.g. https://tools.taskcluster.net/groups/M4YroRtZRbyprTYhnmn2Kw/tasks/KJlvJjJvS-G5Ha1xHe6KlQ/runs/0/artifacts)

Resolving patch:
diff --git a/bugbug/db.py b/bugbug/db.py
--- a/bugbug/db.py
+++ b/bugbug/db.py
@@ -6,10 +6,8 @@
 import gzip
 import io
 import json
-import lzma
 import os
 import pickle
-import shutil
 from contextlib import contextmanager
 from urllib.parse import urljoin
 
@@ -49,11 +47,6 @@ def extract_file(path):
             dctx = zstandard.ZstdDecompressor()
             with open(f"{path}.zst", "rb") as input_f:
                 dctx.copy_stream(input_f, output_f)
-
-        elif compression_type == ".xz":
-            with lzma.open(f"{path}.xz") as input_f:
-                shutil.copyfileobj(input_f, output_f)
-
         else:
             assert False, f"Unexpected compression type: {compression_type}"
 
@@ -66,20 +59,10 @@ def download_support_file(path, file_name):
         print(f"Downloading {url} to {path}")
         utils.download_check_etag(url, path)
 
-        if path.endswith(".zst") or path.endswith(".xz"):
+        if path.endswith(".zst"):
             extract_file(path)
     except requests.exceptions.HTTPError:
-        try:
-            url = f"{os.path.splitext(url)[0]}.xz"
-            path = f"{os.path.splitext(path)[0]}.xz"
-
-            print(f"Downloading {url} to {path}")
-            utils.download_check_etag(url, path)
-
-            extract_file(path)
-
-        except requests.exceptions.HTTPError:
-            print(f"{file_name} is not yet available to download for {path}")
+        print(f"{file_name} is not yet available to download for {path}")
 
 
 def download_version(path):
@@ -92,34 +75,19 @@ def download(path, force=False, support_files_too=False):
         return
 
     zst_path = f"{path}.zst"
-    xz_path = f"{path}.xz"
 
     # Only download if the file is not there yet.
-    if (not os.path.exists(zst_path) and not os.path.exists(xz_path)) or force:
+    if not os.path.exists(zst_path) or force:
         url = DATABASES[path]["url"]
         try:
-            path_compressed = zst_path
-            print(f"Downloading {url} to {path_compressed}")
-            utils.download_check_etag(url, path_compressed)
+            print(f"Downloading {url} to {zst_path}")
+            utils.download_check_etag(url, zst_path)
 
         except requests.exceptions.HTTPError:
-            try:
-                url_xz = f"{os.path.splitext(url)[0]}.xz"
-                path_compressed = xz_path
-                print(f"Downloading {url_xz} to {path_compressed} instead")
-                utils.download_check_etag(url_xz, path_compressed)
-
-            except requests.exceptions.HTTPError:
-                print(f"{url} is not yet available to download")
-                raise
+            print(f"{url} is not yet available to download")
+            raise
 
-    else:
-        if os.path.exists(zst_path) or not os.path.exists(xz_path):
-            path_compressed = zst_path
-        else:
-            path_compressed = xz_path
-
-    extract_file(path_compressed)
+    extract_file(zst_path)
 
     if support_files_too:
         for support_file in DATABASES[path]["support_files"]:
@@ -130,10 +98,6 @@ def last_modified(path):
     url = DATABASES[path]["url"]
     last_modified = utils.get_last_modified(url)
 
-    if last_modified is None:
-        base_url = os.path.splitext(url)[0]
-        last_modified = utils.get_last_modified(f"{base_url}.xz")
-
     if last_modified is None:
         raise Exception("Last-Modified is not available")
 


Test patch:
diff --git a/tests/test_db.py b/tests/test_db.py
--- a/tests/test_db.py
+++ b/tests/test_db.py
@@ -4,7 +4,6 @@
 # You can obtain one at http://mozilla.org/MPL/2.0/.
 
 import json
-import lzma
 import os
 import pickle
 from datetime import datetime
@@ -120,15 +119,6 @@ def create_zst_file(db_path):
     return create_zst_file
 
 
-@pytest.fixture
-def mock_xz():
-    def create_xz_file(db_path):
-        with lzma.open(db_path, "wb") as output_f:
-            output_f.write(b'{"Hello": "World"}')
-
-    return create_xz_file
-
-
 def test_extract_db_zst(tmp_path, mock_zst):
     db_path = tmp_path / f"prova.zst"
 
@@ -142,19 +132,6 @@ def test_extract_db_zst(tmp_path, mock_zst):
     assert file_decomp == {"Hello": "World"}
 
 
-def test_extract_db_xz(tmp_path, mock_xz):
-    db_path = tmp_path / f"prova.xz"
-
-    mock_xz(db_path)
-
-    db.extract_file(db_path)
-
-    with open(f"{os.path.splitext(db_path)[0]}", "rb") as f:
-        file_decomp = json.load(f)
-
-    assert file_decomp == {"Hello": "World"}
-
-
 def test_extract_db_bad_format(tmp_path):
     db_path = tmp_path / "prova.pickle"
 
@@ -198,84 +175,22 @@ def test_download_zst(tmp_path, mock_zst):
     assert os.path.exists(db_path.with_suffix(db_path.suffix + ".zst.etag"))
 
 
-@responses.activate
-def test_download_xz(tmp_path, mock_xz):
-    url_zst = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
-    url_xz = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.xz"
-
-    db_path = tmp_path / "prova.json"
-    db.register(db_path, url_zst, 1)
-
-    responses.add(
-        responses.HEAD,
-        url_zst,
-        status=404,
-        headers={"ETag": "123", "Accept-Encoding": "zstd"},
-    )
-
-    responses.add(
-        responses.GET,
-        url_zst,
-        status=404,
-        body=requests.exceptions.HTTPError("HTTP error"),
-    )
-
-    responses.add(
-        responses.HEAD,
-        url_xz,
-        status=200,
-        headers={"ETag": "123", "Accept-Encoding": "xz", "Last-Modified": "2019-04-16"},
-    )
-
-    tmp_xz_path = tmp_path / "prova_tmp.xz"
-    mock_xz(tmp_xz_path)
-
-    with open(tmp_xz_path, "rb") as content:
-        responses.add(responses.GET, url_xz, status=200, body=content.read())
-
-    db.download(db_path)
-
-    assert db.last_modified(db_path) == datetime(2019, 4, 16)
-
-    assert os.path.exists(db_path)
-    assert os.path.exists(db_path.with_suffix(db_path.suffix + ".xz"))
-    assert os.path.exists(db_path.with_suffix(db_path.suffix + ".xz.etag"))
-
-
 @responses.activate
 def test_download_missing(tmp_path):
-    url_zst = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
-    url_xz = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.xz"
+    url = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
 
     db_path = tmp_path / "prova.json"
-    db.register(db_path, url_zst, 1)
+    db.register(db_path, url, 1)
 
     responses.add(
         responses.HEAD,
-        url_zst,
+        url,
         status=404,
         headers={"ETag": "123", "Accept-Encoding": "zstd"},
     )
 
     responses.add(
-        responses.GET,
-        url_zst,
-        status=404,
-        body=requests.exceptions.HTTPError("HTTP error"),
-    )
-
-    responses.add(
-        responses.HEAD,
-        url_xz,
-        status=404,
-        headers={"ETag": "123", "Accept-Encoding": "xz"},
-    )
-
-    responses.add(
-        responses.GET,
-        url_xz,
-        status=404,
-        body=requests.exceptions.HTTPError("HTTP error"),
+        responses.GET, url, status=404, body=requests.exceptions.HTTPError("HTTP error")
     )
 
     with pytest.raises(requests.exceptions.HTTPError):
@@ -287,16 +202,16 @@ def test_download_missing(tmp_path):
 
 @responses.activate
 def test_download_support_file_zst(tmp_path, mock_zst):
-    url_zst = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
+    url = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
     support_filename = "support.zst"
-    url = urljoin(url_zst, support_filename)
+    url_support = urljoin(url, support_filename)
 
     db_path = tmp_path / "prova.json"
-    db.register(db_path, url_zst, 1, support_files=[support_filename])
+    db.register(db_path, url, 1, support_files=[support_filename])
 
     responses.add(
         responses.HEAD,
-        url,
+        url_support,
         status=200,
         headers={"ETag": "123", "Accept-Encoding": "zstd"},
     )
@@ -305,7 +220,7 @@ def test_download_support_file_zst(tmp_path, mock_zst):
     mock_zst(tmp_zst_path)
 
     with open(tmp_zst_path, "rb") as content:
-        responses.add(responses.GET, url, status=200, body=content.read())
+        responses.add(responses.GET, url_support, status=200, body=content.read())
 
     db.download_support_file(db_path, support_filename)
 
@@ -321,63 +236,6 @@ def test_download_support_file_zst(tmp_path, mock_zst):
     )
 
 
-@responses.activate
-def test_download_support_file_xz(tmp_path, mock_xz):
-    url_zst = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
-    url_xz = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.xz"
-
-    support_filename = "support_mock.zst"
-
-    url_zst_support = urljoin(url_zst, support_filename)
-    url_xz_support = urljoin(url_xz, f"{os.path.splitext(support_filename)[0]}.xz")
-
-    db_path = tmp_path / "prova.json"
-    db.register(db_path, url_zst, 1, support_files=[support_filename])
-
-    responses.add(
-        responses.HEAD,
-        url_zst_support,
-        status=404,
-        headers={"ETag": "123", "Accept-Encoding": "zstd"},
-    )
-
-    responses.add(
-        responses.GET,
-        url_zst_support,
-        status=404,
-        body=requests.exceptions.HTTPError("HTTP error"),
-    )
-
-    responses.add(
-        responses.HEAD,
-        url_xz_support,
-        status=200,
-        headers={"ETag": "123", "Accept-Encoding": "xz"},
-    )
-
-    tmp_xz_path = tmp_path / "prova_tmp.xz"
-    mock_xz(tmp_xz_path)
-
-    with open(tmp_xz_path, "rb") as content:
-        responses.add(responses.GET, url_xz_support, status=200, body=content.read())
-
-    db.download_support_file(db_path, support_filename)
-
-    assert os.path.exists(
-        os.path.join(
-            os.path.dirname(db_path), f"{os.path.splitext(support_filename)[0]}.xz"
-        )
-    )
-    assert os.path.exists(
-        os.path.join(os.path.dirname(db_path), os.path.splitext(support_filename)[0])
-    )
-    assert os.path.exists(
-        os.path.join(
-            os.path.dirname(db_path), os.path.splitext(support_filename)[0] + ".xz.etag"
-        )
-    )
-
-
 @responses.activate
 def test_download_version(tmp_path):
     url_zst = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/prova.json.zst"
@@ -404,41 +262,23 @@ def test_download_version(tmp_path):
 
 @responses.activate
 def test_download_support_file_missing(tmp_path, capfd):
-    url_zst = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
-    url_xz = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.xz"
-
+    url = "https://index.taskcluster.net/v1/task/project.relman.bugbug.data_commits.latest/artifacts/public/commits.json.zst"
     support_filename = "support_mock.zst"
-
-    url_zst_support = urljoin(url_zst, support_filename)
-    url_xz_support = urljoin(url_xz, f"{os.path.splitext(support_filename)[0]}.xz")
+    url_support = urljoin(url, support_filename)
 
     db_path = tmp_path / "prova.json"
-    db.register(db_path, url_zst, 1, support_files=[support_filename])
+    db.register(db_path, url, 1, support_files=[support_filename])
 
     responses.add(
         responses.HEAD,
-        url_zst_support,
+        url_support,
         status=404,
         headers={"ETag": "123", "Accept-Encoding": "zstd"},
     )
 
     responses.add(
         responses.GET,
-        url_zst_support,
-        status=404,
-        body=requests.exceptions.HTTPError("HTTP error"),
-    )
-
-    responses.add(
-        responses.HEAD,
-        url_xz_support,
-        status=404,
-        headers={"ETag": "123", "Accept-Encoding": "xz"},
-    )
-
-    responses.add(
-        responses.GET,
-        url_xz_support,
+        url_support,
         status=404,
         body=requests.exceptions.HTTPError("HTTP error"),
     )
@@ -447,7 +287,7 @@ def test_download_support_file_missing(tmp_path, capfd):
 
     out, err = capfd.readouterr()
     path = os.path.join(
-        os.path.dirname(db_path), f"{os.path.splitext(support_filename)[0]}.xz"
+        os.path.dirname(db_path), f"{os.path.splitext(support_filename)[0]}.zst"
     )
     assert (
         out.split("\n")[-2]