Retry after failure to download or failure to open files (#4609)

* add retry logic, tests when extracting tarfile fails

* fixed bug with not catching empty responses

* specify compression type

* WIP test

* more testing work

* fixed up unit test

* add changelog

* Add more comments!

* clarify why we do the json() check for None

CHANGELOG.md

@@ -21,6 +21,7 @@ Contributors:
 - Projects created using `dbt init` now have the correct `seeds` directory created (instead of `data`) ([#4588](https://github.com/dbt-labs/dbt-core/issues/4588), [#4599](https://github.com/dbt-labs/dbt-core/pull/4589))
 - Don't require a profile for dbt deps and clean commands ([#4554](https://github.com/dbt-labs/dbt-core/issues/4554), [#4610](https://github.com/dbt-labs/dbt-core/pull/4610))
 - Select modified.body works correctly when new model added([#4570](https://github.com/dbt-labs/dbt-core/issues/4570), [#4631](https://github.com/dbt-labs/dbt-core/pull/4631))
+- Fix bug in retry logic for bad response from hub and when there is a bad git tarball download. ([#4577](https://github.com/dbt-labs/dbt-core/issues/4577), [#4579](https://github.com/dbt-labs/dbt-core/issues/4579), [#4609](https://github.com/dbt-labs/dbt-core/pull/4609)) 
 
 ## dbt-core 1.0.1 (January 03, 2022)
 

core/dbt/clients/registry.py

@@ -33,7 +33,12 @@ def _get(path, registry_base_url=None):
     resp = requests.get(url, timeout=30)
     fire_event(RegistryProgressGETResponse(url=url, resp_code=resp.status_code))
     resp.raise_for_status()
-    if resp is None:
+
+    # It is unexpected for the content of the response to be None so if it is, raising this error
+    # will cause this function to retry (if called within _get_with_retries) and hopefully get
+    # a response.  This seems to happen when there's an issue with the Hub.
+    # See https://github.com/dbt-labs/dbt-core/issues/4577
+    if resp.json() is None:
         raise requests.exceptions.ContentDecodingError(
             'Request error: The response is None', response=resp
         )

core/dbt/clients/system.py

@@ -485,7 +485,7 @@ def untar_package(
 ) -> None:
     tar_path = convert_path(tar_path)
     tar_dir_name = None
-    with tarfile.open(tar_path, 'r') as tarball:
+    with tarfile.open(tar_path, 'r:gz') as tarball:
         tarball.extractall(dest_dir)
         tar_dir_name = os.path.commonprefix(tarball.getnames())
     if rename_to:

core/dbt/deps/registry.py

@@ -1,4 +1,5 @@
 import os
+import functools
 from typing import List
 
 from dbt import semver
@@ -14,6 +15,7 @@
     DependencyException,
     package_not_found,
 )
+from dbt.utils import _connection_exception_retry as connection_exception_retry
 
 
 class RegistryPackageMixin:
@@ -68,9 +70,28 @@ def install(self, project, renderer):
         system.make_directory(os.path.dirname(tar_path))
 
         download_url = metadata.downloads.tarball
-        system.download_with_retries(download_url, tar_path)
         deps_path = project.packages_install_path
         package_name = self.get_project_name(project, renderer)
+
+        download_untar_fn = functools.partial(
+            self.download_and_untar,
+            download_url,
+            tar_path,
+            deps_path,
+            package_name
+        )
+        connection_exception_retry(download_untar_fn, 5)
+
+    def download_and_untar(self, download_url, tar_path, deps_path, package_name):
+        """
+        Sometimes the download of the files fails and we want to retry.  Sometimes the
+        download appears successful but the file did not make it through as expected
+        (generally due to a github incident).  Either way we want to retry downloading
+        and untarring to see if we can get a success.  Call this within
+        `_connection_exception_retry`
+        """
+
+        system.download(download_url, tar_path)
         system.untar_package(tar_path, deps_path, package_name)
 
 

core/dbt/exceptions.py

@@ -774,6 +774,10 @@ def system_error(operation_name):
 
 
 class ConnectionException(Exception):
+    """
+    There was a problem with the connection that returned a bad response,
+    timed out, or resulted in a file that is corrupt.
+    """
     pass
 
 

core/dbt/utils.py

@@ -10,6 +10,7 @@
 import json
 import os
 import requests
+from tarfile import ReadError
 import time
 from pathlib import PosixPath, WindowsPath
 
@@ -600,14 +601,17 @@ def __contains__(self, name) -> bool:
 
 def _connection_exception_retry(fn, max_attempts: int, attempt: int = 0):
     """Attempts to run a function that makes an external call, if the call fails
-    on a connection error or timeout, it will be tried up to 5 more times.
+    on a connection error, timeout or decompression issue, it will be tried up to 5 more times.
+    See https://github.com/dbt-labs/dbt-core/issues/4579 for context on this decompression issues
+    specifically.
     """
     try:
         return fn()
     except (
         requests.exceptions.ConnectionError,
         requests.exceptions.Timeout,
         requests.exceptions.ContentDecodingError,
+        ReadError,
     ) as exc:
         if attempt <= max_attempts - 1:
             fire_event(RetryExternalCall(attempt=attempt, max=max_attempts))

test/unit/test_core_dbt_utils.py

@@ -1,4 +1,5 @@
 import requests
+import tarfile
 import unittest
 
 from dbt.exceptions import ConnectionException
@@ -22,11 +23,21 @@ def test_connection_exception_retry_success(self):
         connection_exception_retry(lambda: Counter._add_with_limited_exception(), 5)
         self.assertEqual(2, counter) # 2 = original attempt plus 1 retry
 
+    def test_connection_timeout(self):
+        Counter._reset()
+        connection_exception_retry(lambda: Counter._add_with_timeout(), 5)
+        self.assertEqual(2, counter) # 2 = original attempt plus 1 retry
+
     def test_connection_exception_retry_success_none_response(self):
         Counter._reset()
         connection_exception_retry(lambda: Counter._add_with_none_exception(), 5)
         self.assertEqual(2, counter) # 2 = original attempt returned None, plus 1 retry
 
+    def test_connection_exception_retry_success_failed_untar(self):
+        Counter._reset()
+        connection_exception_retry(lambda: Counter._add_with_untar_exception(), 5)
+        self.assertEqual(2, counter) # 2 = original attempt returned ReadError, plus 1 retry
+
 
 counter:int = 0 
 class Counter():
@@ -42,11 +53,21 @@ def _add_with_limited_exception():
         counter+=1
         if counter < 2:
             raise requests.exceptions.ConnectionError
+    def _add_with_timeout():
+        global counter
+        counter+=1
+        if counter < 2:
+            raise requests.exceptions.Timeout
     def _add_with_none_exception():
         global counter
         counter+=1
         if counter < 2:
             raise requests.exceptions.ContentDecodingError
+    def _add_with_untar_exception():
+        global counter
+        counter+=1
+        if counter < 2:
+            raise tarfile.ReadError
     def _reset():
         global counter
         counter = 0

test/unit/test_system_client.py

@@ -2,6 +2,9 @@
 import shutil
 import stat
 import unittest
+import tarfile
+import io
+from pathlib import Path
 from tempfile import mkdtemp, NamedTemporaryFile
 
 from dbt.exceptions import ExecutableError, WorkingDirectoryError
@@ -185,3 +188,63 @@ def tearDown(self):
             shutil.rmtree(self.base_dir)
         except:
             pass
+
+
+class TestUntarPackage(unittest.TestCase):
+
+    def setUp(self):
+        self.base_dir = mkdtemp()
+        self.tempdir = mkdtemp(dir=self.base_dir)
+        self.tempdest = mkdtemp(dir=self.base_dir)
+
+    def tearDown(self):
+        try:
+            shutil.rmtree(self.base_dir)
+        except:
+            pass
+
+    def test_untar_package_success(self):
+        #  set up a valid tarball to test against
+        with NamedTemporaryFile(
+            prefix='my-package.2', suffix='.tar.gz', dir=self.tempdir, delete=False
+        ) as named_tar_file:
+            tar_file_full_path = named_tar_file.name
+            with NamedTemporaryFile(
+                prefix='a', suffix='.txt', dir=self.tempdir
+            ) as file_a:
+                file_a.write(b'some text in the text file')
+                relative_file_a = os.path.basename(file_a.name)
+                with tarfile.open(fileobj=named_tar_file, mode='w:gz') as tar:
+                    tar.addfile(tarfile.TarInfo(relative_file_a), open(file_a.name))
+
+        #  now we test can test that we can untar the file successfully
+        assert tarfile.is_tarfile(tar.name)
+        dbt.clients.system.untar_package(tar_file_full_path, self.tempdest)
+        path = Path(os.path.join(self.tempdest, relative_file_a))
+        assert path.is_file()
+
+    def test_untar_package_failure(self):
+        #  create a text file then rename it as a tar (so it's invalid)
+        with NamedTemporaryFile(
+                prefix='a', suffix='.txt', dir=self.tempdir, delete=False
+            ) as file_a:
+                file_a.write(b'some text in the text file')
+                txt_file_name = file_a.name
+                file_path= os.path.dirname(txt_file_name)
+                tar_file_path = os.path.join(file_path, 'mypackage.2.tar.gz')
+        os.rename(txt_file_name, tar_file_path)
+
+        #  now that we're set up, test that untarring the file fails
+        with self.assertRaises(tarfile.ReadError) as exc:
+            dbt.clients.system.untar_package(tar_file_path, self.tempdest)
+
+    def test_untar_package_empty(self):
+        #  create a tarball with nothing in it
+        with NamedTemporaryFile(
+            prefix='my-empty-package.2', suffix='.tar.gz', dir=self.tempdir
+        ) as named_file:
+
+            #  make sure we throw an error for the empty file
+            with self.assertRaises(tarfile.ReadError) as exc:
+                dbt.clients.system.untar_package(named_file.name, self.tempdest)
+            self.assertEqual("empty file", str(exc.exception))