Add additional headers to check API call

client/v1/docker_http_.py

@@ -52,7 +52,8 @@ def Request(transport,
             credentials,
             accepted_codes = None,
             body = None,
-            content_type = None):
+            content_type = None,
+            extra_headers = None):
   """Wrapper containing much of the boilerplate REST logic for Registry calls.
 
   Args:
@@ -62,6 +63,7 @@ def Request(transport,
     accepted_codes: the list of acceptable http status codes
     body: the body to pass into the PUT request (or None for GET)
     content_type: the mime-type of the request (or None for JSON)
+    extra_headers: optional dictionary of additional headers to include
 
   Raises:
     BadStatusException: the status codes wasn't among the acceptable set.
@@ -75,6 +77,8 @@ def Request(transport,
       'X-Docker-Token': 'true',
       'user-agent': docker_name.USER_AGENT,
   }
+  if extra_headers:
+    headers.update(extra_headers)
   resp, content = transport.request(
       url, 'PUT' if body else 'GET', body=body, headers=headers)
 

client/v2/docker_http_.py

@@ -306,7 +306,8 @@ def Request(
       accepted_codes = None,
       method = None,
       body = None,
-      content_type = None):
+      content_type = None,
+      extra_headers = None):
     """Wrapper containing much of the boilerplate REST logic for Registry calls.
 
     Args:
@@ -317,6 +318,7 @@ def Request(
       body: the body to pass into the PUT request (or None for GET)
       content_type: the mime-type of the request (or None for JSON).
               content_type is ignored when body is None.
+      extra_headers: optional dictionary of additional headers to include
 
     Raises:
       BadStateException: an unexpected internal state has been encountered.
@@ -348,6 +350,9 @@ def Request(
       if method in ('POST', 'PUT') and not body:
         headers['content-length'] = '0'
 
+      if extra_headers:
+        headers.update(extra_headers)
+
       resp, content = self._transport.request(
           url, method, body=body, headers=headers)
 

client/v2/docker_session_.py

@@ -107,7 +107,8 @@ def _manifest_exists(self, image):
         method='GET',
         accepted_codes=[
             six.moves.http_client.OK, six.moves.http_client.NOT_FOUND
-        ])
+        ],
+        extra_headers={'Databricks-Arf-Bypass-Redirect': 'true'})
 
     return resp.status == six.moves.http_client.OK  # pytype: disable=attribute-error
 

client/v2_2/docker_http_.py

@@ -335,7 +335,8 @@ def Request(self,
               method = None,
               body = None,
               content_type = None,
-              accepted_mimes = None
+              accepted_mimes = None,
+              extra_headers = None
              ):
     """Wrapper containing much of the boilerplate REST logic for Registry calls.
 
@@ -348,6 +349,7 @@ def Request(self,
       content_type: the mime-type of the request (or None for JSON).
               content_type is ignored when body is None.
       accepted_mimes: the list of acceptable mime-types
+      extra_headers: optional dictionary of additional headers to include
 
     Raises:
       BadStateException: an unexpected internal state has been encountered.
@@ -382,6 +384,9 @@ def Request(self,
       if method in ('POST', 'PUT') and not body:
         headers['content-length'] = '0'
 
+      if extra_headers:
+        headers.update(extra_headers)
+
       resp, content = self._transport.request(
           url, method, body=body, headers=headers)
 

client/v2_2/docker_session_.py

@@ -110,7 +110,8 @@ def _manifest_exists(
         accepted_codes=[
             six.moves.http_client.OK, six.moves.http_client.NOT_FOUND
         ],
-        accepted_mimes=[image.media_type()])
+        accepted_mimes=[image.media_type()],
+        extra_headers={'Databricks-Arf-Bypass-Redirect': 'true'})
 
     return resp.status == six.moves.http_client.OK  # pytype: disable=attribute-error