Fixed issue where customers could not capture response of directory a…

…nd file GetProperties calls (#21407)

sdk/storage/azdatalake/CHANGELOG.md

@@ -7,6 +7,7 @@
 ### Breaking Changes
 
 ### Bugs Fixed
+* Fixed an issue where customers could not capture the raw HTTP response of directory and file GetProperties operations.
 
 ### Other Changes
 

sdk/storage/azdatalake/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "go",
   "TagPrefix": "go/storage/azdatalake",
-  "Tag": "go/storage/azdatalake_c3c16cffab"
+  "Tag": "go/storage/azdatalake_93ff4d0fca"
 }

sdk/storage/azdatalake/directory/client.go

@@ -56,6 +56,11 @@ func NewClient(directoryURL string, cred azcore.TokenCredential, options *Client
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	blobClientOpts := blockblob.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -84,6 +89,12 @@ func NewClientWithNoCredential(directoryURL string, options *ClientOptions) (*Cl
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
+
 	blobClientOpts := blockblob.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -115,6 +126,11 @@ func NewClientWithSharedKeyCredential(directoryURL string, cred *SharedKeyCreden
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	blobClientOpts := blockblob.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -195,14 +211,15 @@ func (d *Client) NewFileClient(fileName string) (*file.Client, error) {
 	fileURL := runtime.JoinPaths(d.DFSURL(), fileName)
 	newBlobURL, fileURL := shared.GetURLs(fileURL)
 	var newBlobClient *blockblob.Client
+	clientOptions := &blockblob.ClientOptions{ClientOptions: d.getClientOptions().ClientOptions}
 	var err error
 	if d.identityCredential() != nil {
-		newBlobClient, err = blockblob.NewClient(newBlobURL, *d.identityCredential(), nil)
+		newBlobClient, err = blockblob.NewClient(newBlobURL, *d.identityCredential(), clientOptions)
 	} else if d.sharedKey() != nil {
 		blobSharedKey, _ := exported.ConvertToBlobSharedKey(d.sharedKey())
-		newBlobClient, err = blockblob.NewClientWithSharedKeyCredential(newBlobURL, blobSharedKey, nil)
+		newBlobClient, err = blockblob.NewClientWithSharedKeyCredential(newBlobURL, blobSharedKey, clientOptions)
 	} else {
-		newBlobClient, err = blockblob.NewClientWithNoCredential(newBlobURL, nil)
+		newBlobClient, err = blockblob.NewClientWithNoCredential(newBlobURL, clientOptions)
 	}
 	if err != nil {
 		return nil, exported.ConvertToDFSError(err)
@@ -230,7 +247,7 @@ func (d *Client) Delete(ctx context.Context, options *DeleteOptions) (DeleteResp
 func (d *Client) GetProperties(ctx context.Context, options *GetPropertiesOptions) (GetPropertiesResponse, error) {
 	opts := path.FormatGetPropertiesOptions(options)
 	var respFromCtx *http.Response
-	ctxWithResp := runtime.WithCaptureResponse(ctx, &respFromCtx)
+	ctxWithResp := shared.WithCaptureBlobResponse(ctx, &respFromCtx)
 	resp, err := d.blobClient().GetProperties(ctxWithResp, opts)
 	newResp := path.FormatGetPropertiesResponse(&resp, respFromCtx)
 	err = exported.ConvertToDFSError(err)

sdk/storage/azdatalake/directory/client_test.go

@@ -8,6 +8,7 @@ package directory_test
 
 import (
 	"context"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/recording"
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azdatalake/datalakeerror"
@@ -16,6 +17,7 @@ import (
 	"github.com/Azure/azure-sdk-for-go/sdk/storage/azdatalake/sas"
 	"github.com/stretchr/testify/require"
 	"github.com/stretchr/testify/suite"
+	"net/http"
 	"testing"
 	"time"
 )
@@ -2442,4 +2444,57 @@ func (s *RecordedTestSuite) TestRenameDirIfETagMatchFalse() {
 	testcommon.ValidateErrorCode(_require, err, datalakeerror.SourceConditionNotMet)
 }
 
+func (s *RecordedTestSuite) TestDirGetPropertiesResponseCapture() {
+	_require := require.New(s.T())
+	testName := s.T().Name()
+
+	filesystemName := testcommon.GenerateFileSystemName(testName)
+	fsClient, err := testcommon.GetFileSystemClient(filesystemName, s.T(), testcommon.TestAccountDatalake, nil)
+	_require.NoError(err)
+	defer testcommon.DeleteFileSystem(context.Background(), _require, fsClient)
+
+	_, err = fsClient.Create(context.Background(), nil)
+	_require.Nil(err)
+
+	dirName := testcommon.GenerateDirName(testName)
+	dirClient, err := testcommon.GetDirClient(filesystemName, dirName, s.T(), testcommon.TestAccountDatalake, nil)
+	_require.NoError(err)
+
+	resp, err := dirClient.Create(context.Background(), nil)
+	_require.Nil(err)
+	_require.NotNil(resp)
+
+	// This tests directory.NewClient
+	var respFromCtxDir *http.Response
+	ctxWithRespDir := runtime.WithCaptureResponse(context.Background(), &respFromCtxDir)
+	resp2, err := dirClient.GetProperties(ctxWithRespDir, nil)
+	_require.Nil(err)
+	_require.NotNil(resp2)
+	_require.NotNil(respFromCtxDir) // validate that the respFromCtx is actually populated
+	_require.Equal("directory", respFromCtxDir.Header.Get("x-ms-resource-type"))
+
+	// This tests filesystem.NewClient
+	dirClient = fsClient.NewDirectoryClient(dirName)
+	var respFromCtxFs *http.Response
+	ctxWithRespFs := runtime.WithCaptureResponse(context.Background(), &respFromCtxFs)
+	resp2, err = dirClient.GetProperties(ctxWithRespFs, nil)
+	_require.Nil(err)
+	_require.NotNil(resp2)
+	_require.NotNil(respFromCtxFs) // validate that the respFromCtx is actually populated
+	_require.Equal("directory", respFromCtxFs.Header.Get("x-ms-resource-type"))
+
+	// This tests service.NewClient
+	serviceClient, err := testcommon.GetServiceClient(s.T(), testcommon.TestAccountDatalake, nil)
+	_require.Nil(err)
+	fsClient = serviceClient.NewFileSystemClient(filesystemName)
+	dirClient = fsClient.NewDirectoryClient(dirName)
+	var respFromCtxService *http.Response
+	ctxWithRespService := runtime.WithCaptureResponse(context.Background(), &respFromCtxService)
+	resp2, err = dirClient.GetProperties(ctxWithRespService, nil)
+	_require.Nil(err)
+	_require.NotNil(resp2)
+	_require.NotNil(respFromCtxService) // validate that the respFromCtx is actually populated
+	_require.Equal("directory", respFromCtxService.Header.Get("x-ms-resource-type"))
+}
+
 // TODO: more tests for acls

sdk/storage/azdatalake/file/client.go

@@ -63,6 +63,11 @@ func NewClient(fileURL string, cred azcore.TokenCredential, options *ClientOptio
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	blobClientOpts := blockblob.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -91,6 +96,11 @@ func NewClientWithNoCredential(fileURL string, options *ClientOptions) (*Client,
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	blobClientOpts := blockblob.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -122,6 +132,11 @@ func NewClientWithSharedKeyCredential(fileURL string, cred *SharedKeyCredential,
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	blobClientOpts := blockblob.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -216,7 +231,7 @@ func (f *Client) Delete(ctx context.Context, options *DeleteOptions) (DeleteResp
 func (f *Client) GetProperties(ctx context.Context, options *GetPropertiesOptions) (GetPropertiesResponse, error) {
 	opts := path.FormatGetPropertiesOptions(options)
 	var respFromCtx *http.Response
-	ctxWithResp := runtime.WithCaptureResponse(ctx, &respFromCtx)
+	ctxWithResp := shared.WithCaptureBlobResponse(ctx, &respFromCtx)
 	resp, err := f.blobClient().GetProperties(ctxWithResp, opts)
 	newResp := path.FormatGetPropertiesResponse(&resp, respFromCtx)
 	err = exported.ConvertToDFSError(err)

sdk/storage/azdatalake/file/client_test.go

@@ -11,6 +11,7 @@ import (
 	"context"
 	"crypto/md5"
 	"encoding/binary"
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"hash/crc64"
 	"io"
 	"math/rand"
@@ -3885,4 +3886,80 @@ func (s *RecordedTestSuite) TestFileDownloadSmallBufferWithAccessConditions() {
 	_require.Equal(*gResp2.ContentLength, fileSize)
 }
 
+func (s *RecordedTestSuite) TestFileGetPropertiesResponseCapture() {
+	_require := require.New(s.T())
+	testName := s.T().Name()
+
+	filesystemName := testcommon.GenerateFileSystemName(testName)
+	fsClient, err := testcommon.GetFileSystemClient(filesystemName, s.T(), testcommon.TestAccountDatalake, nil)
+	_require.NoError(err)
+	defer testcommon.DeleteFileSystem(context.Background(), _require, fsClient)
+
+	_, err = fsClient.Create(context.Background(), nil)
+	_require.Nil(err)
+
+	dirName := testcommon.GenerateDirName(testName)
+	dirClient, err := testcommon.GetDirClient(filesystemName, dirName, s.T(), testcommon.TestAccountDatalake, nil)
+	_require.NoError(err)
+
+	resp, err := dirClient.Create(context.Background(), nil)
+	_require.Nil(err)
+	_require.NotNil(resp)
+
+	fileName := testcommon.GenerateFileName(testName)
+	fClient, err := testcommon.GetFileClient(filesystemName, dirName+"/"+fileName, s.T(), testcommon.TestAccountDatalake, nil)
+	_require.NoError(err)
+
+	resp, err = fClient.Create(context.Background(), nil)
+	_require.Nil(err)
+	_require.NotNil(resp)
+
+	// This tests file.NewClient
+	var respFromCtxFile *http.Response
+	ctxWithRespFile := runtime.WithCaptureResponse(context.Background(), &respFromCtxFile)
+	resp2, err := fClient.GetProperties(ctxWithRespFile, nil)
+	_require.Nil(err)
+	_require.NotNil(resp2)
+	_require.NotNil(respFromCtxFile) // validate that the respFromCtx is actually populated
+	_require.Equal("file", respFromCtxFile.Header.Get("x-ms-resource-type"))
+
+	// This tests filesystem.NewClient
+	fClient = fsClient.NewFileClient(dirName + "/" + fileName)
+	var respFromCtxFs *http.Response
+	ctxWithRespFs := runtime.WithCaptureResponse(context.Background(), &respFromCtxFs)
+	resp2, err = fClient.GetProperties(ctxWithRespFs, nil)
+	_require.Nil(err)
+	_require.NotNil(resp2)
+	_require.NotNil(respFromCtxFs) // validate that the respFromCtx is actually populated
+	_require.Equal("file", respFromCtxFs.Header.Get("x-ms-resource-type"))
+
+	// This tests service.NewClient
+	serviceClient, err := testcommon.GetServiceClient(s.T(), testcommon.TestAccountDatalake, nil)
+	_require.Nil(err)
+	fsClient = serviceClient.NewFileSystemClient(filesystemName)
+	dirClient = fsClient.NewDirectoryClient(dirName)
+	fClient, err = dirClient.NewFileClient(fileName)
+	_require.Nil(err)
+	var respFromCtxService *http.Response
+	ctxWithRespService := runtime.WithCaptureResponse(context.Background(), &respFromCtxService)
+	resp2, err = fClient.GetProperties(ctxWithRespService, nil)
+	_require.Nil(err)
+	_require.NotNil(resp2)
+	_require.NotNil(respFromCtxService) // validate that the respFromCtx is actually populated
+	_require.Equal("file", respFromCtxService.Header.Get("x-ms-resource-type"))
+
+	// This tests directory.NewClient
+	var respFromCtxDir *http.Response
+	ctxWithRespDir := runtime.WithCaptureResponse(context.Background(), &respFromCtxDir)
+	dirClient, err = testcommon.GetDirClient(filesystemName, dirName, s.T(), testcommon.TestAccountDatalake, nil)
+	_require.Nil(err)
+	fClient, err = dirClient.NewFileClient(fileName)
+	_require.Nil(err)
+	resp2, err = fClient.GetProperties(ctxWithRespDir, nil)
+	_require.Nil(err)
+	_require.NotNil(resp2)
+	_require.NotNil(respFromCtxDir) // validate that the respFromCtx is actually populated
+	_require.Equal("file", respFromCtxDir.Header.Get("x-ms-resource-type"))
+}
+
 // TODO tests all uploads/downloads with other opts

sdk/storage/azdatalake/filesystem/client.go

@@ -55,6 +55,11 @@ func NewClient(filesystemURL string, cred azcore.TokenCredential, options *Clien
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	containerClientOpts := container.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -82,6 +87,11 @@ func NewClientWithNoCredential(filesystemURL string, options *ClientOptions) (*C
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	containerClientOpts := container.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}
@@ -112,6 +122,11 @@ func NewClientWithSharedKeyCredential(filesystemURL string, cred *SharedKeyCrede
 	if options == nil {
 		options = &ClientOptions{}
 	}
+	perCallPolicies := []policy.Policy{shared.NewIncludeBlobResponsePolicy()}
+	if options.ClientOptions.PerCallPolicies != nil {
+		perCallPolicies = append(perCallPolicies, options.ClientOptions.PerCallPolicies...)
+	}
+	options.ClientOptions.PerCallPolicies = perCallPolicies
 	containerClientOpts := container.ClientOptions{
 		ClientOptions: options.ClientOptions,
 	}

sdk/storage/azdatalake/internal/shared/policy_include_blob_response.go

@@ -0,0 +1,45 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+package shared
+
+import (
+	"context"
+	"net/http"
+
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/policy"
+)
+
+// NOTE: This is duplication of the azcore includeResponsePolicy, this is necessary for some APIs like directory/file
+// 		 GetProperties. Under the hood, these APIs grab the raw response to construct the datalake properties
+//		 (owner, group, permission) in addition to the blob properties. If we use the includeResponsePolicy, this can
+//		 result in the customer not being able to retrieve the response themselves.
+
+// CtxIncludeBlobResponseKey is used as a context key for retrieving the raw response.
+type CtxIncludeBlobResponseKey struct{}
+
+type includeBlobResponsePolicy struct {
+}
+
+// NewIncludeBlobResponsePolicy creates a policy that retrieves the raw HTTP response upon request
+func NewIncludeBlobResponsePolicy() policy.Policy {
+	return &includeBlobResponsePolicy{}
+}
+
+func (p *includeBlobResponsePolicy) Do(req *policy.Request) (*http.Response, error) {
+	resp, err := req.Next()
+	if resp == nil {
+		return resp, err
+	}
+	if httpOutRaw := req.Raw().Context().Value(CtxIncludeBlobResponseKey{}); httpOutRaw != nil {
+		httpOut := httpOutRaw.(**http.Response)
+		*httpOut = resp
+	}
+	return resp, err
+}
+
+// WithCaptureBlobResponse applies the HTTP response retrieval annotation to the parent context.
+// The resp parameter will contain the HTTP response after the request has completed.
+func WithCaptureBlobResponse(parent context.Context, resp **http.Response) context.Context {
+	return context.WithValue(parent, CtxIncludeBlobResponseKey{}, resp)
+}