Support additional headers

CHANGELOG.md

@@ -1,5 +1,6 @@
 ## [Unreleased]
 
+- Add support for additional headers to all endpoints
 - Fix typo in class name `AzureBlob::ForbidenError` to `AzureBlob::ForbiddenError`
 - Fix proper URI encoding for keys with special characters like question marks
 

lib/azure_blob/client.rb

@@ -72,7 +72,7 @@ def get_blob(key, options = {})
 
       headers = {
         "x-ms-range": options[:start] && "bytes=#{options[:start]}-#{options[:end]}",
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:).get
     end
@@ -97,7 +97,7 @@ def copy_blob(key, source_key, options = {})
       headers = {
         "x-ms-copy-source": source_uri.to_s,
         "x-ms-requires-sync": "true",
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:, **options.slice(:metadata, :tags)).put
     end
@@ -116,7 +116,7 @@ def delete_blob(key, options = {})
 
       headers = {
         "x-ms-delete-snapshots": options[:delete_snapshots] || "include",
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:).delete
     end
@@ -157,7 +157,7 @@ def list_blobs(options = {})
         query[:marker] = marker
         query.reject! { |key, value| value.to_s.empty? }
         uri.query = URI.encode_www_form(**query)
-        response = Http.new(uri, signer:).get
+        response = Http.new(uri, additional_headers(options), signer:).get
       end
 
       BlobList.new(fetcher)
@@ -172,7 +172,7 @@ def list_blobs(options = {})
     def get_blob_properties(key, options = {})
       uri = generate_uri("#{container}/#{key}")
 
-      response = Http.new(uri, signer:).head
+      response = Http.new(uri, additional_headers(options), signer:).head
 
       Blob.new(response)
     end
@@ -193,10 +193,10 @@ def blob_exist?(key, options = {})
     # Takes a key (path) of the blob.
     #
     # Returns a hash of the blob's tags.
-    def get_blob_tags(key)
+    def get_blob_tags(key, options = {})
       uri = generate_uri("#{container}/#{key}")
       uri.query = URI.encode_www_form(comp: "tags")
-      response = Http.new(uri, signer:).get
+      response = Http.new(uri, additional_headers(options), signer:).get
 
       Tags.from_response(response).to_h
     end
@@ -209,7 +209,7 @@ def get_blob_tags(key)
     def get_container_properties(options = {})
       uri = generate_uri(container)
       uri.query = URI.encode_www_form(restype: "container")
-      response = Http.new(uri, signer:, raise_on_error: false).head
+      response = Http.new(uri, additional_headers(options), signer:, raise_on_error: false).head
 
       Container.new(response)
     end
@@ -229,6 +229,7 @@ def create_container(options = {})
       headers = {}
       headers[:"x-ms-blob-public-access"] = "blob" if options[:public_access]
       headers[:"x-ms-blob-public-access"] = options[:public_access] if [ "container", "blob" ].include?(options[:public_access])
+      headers.merge!(additional_headers(options))
 
       uri.query = URI.encode_www_form(restype: "container")
       response = Http.new(uri, headers, signer:).put
@@ -240,7 +241,7 @@ def create_container(options = {})
     def delete_container(options = {})
       uri = generate_uri(container)
       uri.query = URI.encode_www_form(restype: "container")
-      response = Http.new(uri, signer:).delete
+      response = Http.new(uri, additional_headers(options), signer:).delete
     end
 
     # Return a URI object to a resource in the container. Takes a path.
@@ -284,7 +285,7 @@ def create_append_blob(key, options = {})
         "Content-Type": options[:content_type],
         "Content-MD5": options[:content_md5],
         "x-ms-blob-content-disposition": options[:content_disposition],
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:, **options.slice(:metadata, :tags)).put(nil)
     end
@@ -306,7 +307,7 @@ def append_blob_block(key, content, options = {})
         "Content-Length": content.size,
         "Content-Type": options[:content_type],
         "Content-MD5": options[:content_md5],
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:).put(content)
     end
@@ -330,7 +331,7 @@ def put_blob_block(key, index, content, options = {})
         "Content-Length": content.size,
         "Content-Type": options[:content_type],
         "Content-MD5": options[:content_md5],
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:).put(content)
 
@@ -359,12 +360,17 @@ def commit_blob_blocks(key, block_ids, options = {})
         "Content-Type": options[:content_type],
         "x-ms-blob-content-md5": options[:content_md5],
         "x-ms-blob-content-disposition": options[:content_disposition],
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:, **options.slice(:metadata, :tags)).put(content)
     end
 
-    private
+  private
+
+    def additional_headers(options)
+      (options[:headers] || {}).transform_keys { |k| "x-ms-#{k}".to_sym }.
+        transform_values(&:to_s)
+    end
 
     def generate_block_id(index)
       Base64.urlsafe_encode64(index.to_s.rjust(6, "0"))
@@ -391,7 +397,7 @@ def put_blob_single(key, content, options = {})
         "Content-Type": options[:content_type],
         "x-ms-blob-content-md5": options[:content_md5],
         "x-ms-blob-content-disposition": options[:content_disposition],
-      }
+      }.merge(additional_headers(options))
 
       Http.new(uri, headers, signer:, **options.slice(:metadata, :tags)).put(content.read)
     end

test/client/test_client.rb

@@ -433,4 +433,44 @@ def test_copy_between_containers
     rescue AzureBlob::Http::FileNotFoundError
     end
   end
+
+  def test_get_blob_additional_headers
+    http_mock = Minitest::Mock.new
+    http_mock.expect :get, ""
+
+    stubbed_new = lambda do |uri, headers = {}, signer: nil, **kwargs|
+      assert_equal "bar", headers[:"x-ms-foo"]
+      http_mock
+    end
+
+    AzureBlob::Http.stub :new, stubbed_new do
+      custom_client = AzureBlob::Client.new(account_name: "foo", access_key: "bar", container: "cont")
+      custom_client.get_blob(key, headers: { foo: "bar" })
+    end
+
+    http_mock.verify
+    dummy = Minitest::Mock.new
+    dummy.expect :delete_blob, nil, [ key ]
+    @client = dummy
+  end
+
+  def test_create_append_blob_additional_headers
+    http_mock = Minitest::Mock.new
+    http_mock.expect :put, true, [ nil ]
+
+    stubbed_new = lambda do |uri, headers = {}, signer: nil, **kwargs|
+      assert_equal "bar", headers[:"x-ms-foo"]
+      http_mock
+    end
+
+    AzureBlob::Http.stub :new, stubbed_new do
+      custom_client = AzureBlob::Client.new(account_name: "foo", access_key: "bar", container: "cont")
+      custom_client.create_append_blob(key, headers: { foo: "bar" })
+    end
+
+    http_mock.verify
+    dummy = Minitest::Mock.new
+    dummy.expect :delete_blob, nil, [ key ]
+    @client = dummy
+  end
 end