From c21199a1299f77ae0d7a2ae3b00d89089f6acbb7 Mon Sep 17 00:00:00 2001
From: Siddharth Gupta <siddharth.gupta@databricks.com>
Date: Tue, 6 Jan 2026 21:11:17 +0000
Subject: [PATCH] Add additional headers to check API call

---
 client/v1/docker_http_.py      | 6 +++++-
 client/v2/docker_http_.py      | 7 ++++++-
 client/v2/docker_session_.py   | 3 ++-
 client/v2_2/docker_http_.py    | 7 ++++++-
 client/v2_2/docker_session_.py | 3 ++-
 5 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/client/v1/docker_http_.py b/client/v1/docker_http_.py
index a21b68dbb..86fcc61a5 100755
--- a/client/v1/docker_http_.py
+++ b/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)
 
diff --git a/client/v2/docker_http_.py b/client/v2/docker_http_.py
index 5945531d6..61edf97c4 100755
--- a/client/v2/docker_http_.py
+++ b/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)
 
diff --git a/client/v2/docker_session_.py b/client/v2/docker_session_.py
index 070e96960..7aa80c83b 100755
--- a/client/v2/docker_session_.py
+++ b/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
 
diff --git a/client/v2_2/docker_http_.py b/client/v2_2/docker_http_.py
index b3462c76b..8afa6b74e 100755
--- a/client/v2_2/docker_http_.py
+++ b/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)
 
diff --git a/client/v2_2/docker_session_.py b/client/v2_2/docker_session_.py
index 7aff6ae35..0d3fb9d15 100755
--- a/client/v2_2/docker_session_.py
+++ b/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