From 2b0d9cb7a72cb99a291f51ea384e5c1c45cab985 Mon Sep 17 00:00:00 2001
From: Piotr Fus <piotr.fus@snowflake.com>
Date: Thu, 14 Nov 2024 14:03:41 +0100
Subject: [PATCH] SNOW-1760222 Introduce S3 logging config (#1243)

---
 doc.go               |  2 ++
 s3_storage_client.go | 28 ++++++++++++++++++++++++----
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/doc.go b/doc.go
index cf095d51d..be6ad6093 100644
--- a/doc.go
+++ b/doc.go
@@ -198,6 +198,8 @@ Users can use SetLogger in driver.go to set a customized logger for gosnowflake
 In order to enable debug logging for the driver, user could use SetLogLevel("debug") in SFLogger interface
 as shown in demo code at cmd/logger.go. To redirect the logs SFlogger.SetOutput method could do the work.
 
+# If you want to define S3 client logging, override S3LoggingMode variable using configuration: https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/aws#ClientLogMode
+
 # Query tag
 
 A custom query tag can be set in the context. Each query run with this context
diff --git a/s3_storage_client.go b/s3_storage_client.go
index 7cd1e3301..c389c3127 100644
--- a/s3_storage_client.go
+++ b/s3_storage_client.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"github.com/aws/smithy-go/logging"
 	"io"
 	"net/http"
 	"os"
@@ -38,11 +39,19 @@ type s3Location struct {
 	s3Path     string
 }
 
+// S3LoggingMode allows to configure which logs should be included.
+// By default no logs are included.
+// See https://pkg.go.dev/github.com/aws/aws-sdk-go-v2/aws#ClientLogMode for allowed values.
+var S3LoggingMode aws.ClientLogMode
+
 func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAccelerateEndpoint bool) (cloudClient, error) {
 	stageCredentials := info.Creds
-	var resolver s3.EndpointResolver
+	s3Logger := logging.LoggerFunc(s3LoggingFunc)
+
+	var endpoint *string
 	if info.EndPoint != "" {
-		resolver = s3.EndpointResolverFromURL("https://" + info.EndPoint) // FIPS endpoint
+		tmp := "https://" + info.EndPoint
+		endpoint = &tmp
 	}
 
 	return s3.New(s3.Options{
@@ -51,14 +60,25 @@ func (util *snowflakeS3Client) createClient(info *execResponseStageInfo, useAcce
 			stageCredentials.AwsKeyID,
 			stageCredentials.AwsSecretKey,
 			stageCredentials.AwsToken)),
-		EndpointResolver: resolver,
-		UseAccelerate:    useAccelerateEndpoint,
+		BaseEndpoint:  endpoint,
+		UseAccelerate: useAccelerateEndpoint,
 		HTTPClient: &http.Client{
 			Transport: SnowflakeTransport,
 		},
+		ClientLogMode: S3LoggingMode,
+		Logger:        s3Logger,
 	}), nil
 }
 
+func s3LoggingFunc(classification logging.Classification, format string, v ...interface{}) {
+	switch classification {
+	case logging.Debug:
+		logger.WithField("logger", "S3").Debugf(format, v...)
+	case logging.Warn:
+		logger.WithField("logger", "S3").Warnf(format, v...)
+	}
+}
+
 type s3HeaderAPI interface {
 	HeadObject(ctx context.Context, params *s3.HeadObjectInput, optFns ...func(*s3.Options)) (*s3.HeadObjectOutput, error)
 }