feat: track external client logs

goto · Oct 21, 2024 · 6a82ca4 · 6a82ca4
1 parent e9a3775
commit 6a82ca4
Show file tree

Hide file tree

Showing 8 changed files with 89 additions and 35 deletions.
diff --git a/pkg/opentelemetry/config.go b/pkg/opentelemetry/config.go
@@ -13,5 +13,5 @@ type Config struct {
 		Endpoint string            `mapstructure:"endpoint" default:"127.0.0.1:4317"`
 	} `mapstructure:"otlp"`
 	SamplingFraction int           `mapstructure:"sampling_fraction"`
-	MetricInterval   time.Duration `mapstructure:"metric_interval"`
+	MetricInterval   time.Duration `mapstructure:"metric_interval" default:"1s"`
 }
diff --git a/plugins/providers/bigquery/client.go b/plugins/providers/bigquery/client.go
@@ -4,11 +4,14 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"net/http"
 	"strings"
 
 	bq "cloud.google.com/go/bigquery"
 	"github.com/goto/guardian/domain"
-	"github.com/goto/guardian/pkg/opentelemetry"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google"
 	bqApi "google.golang.org/api/bigquery/v2"
 	"google.golang.org/api/cloudresourcemanager/v1"
 	"google.golang.org/api/iam/v1"
@@ -23,28 +26,57 @@ type bigQueryClient struct {
 	crmService *cloudresourcemanager.Service
 }
 
-func NewBigQueryClient(projectID string, opts ...option.ClientOption) (*bigQueryClient, error) {
+func NewBigQueryClient(projectID string, credentialsJSON []byte, opts ...option.ClientOption) (*bigQueryClient, error) {
 	ctx := context.Background()
-	bqHTTPClient := opentelemetry.NewHttpClient(ctx, "BigQueryClient", opts...)
-	bqClient, err := bq.NewClient(ctx, projectID, option.WithHTTPClient(bqHTTPClient))
+	var creds *google.Credentials
+	var oauthClient func(name string) *http.Client
+	if credentialsJSON != nil {
+		var err error
+		creds, err = google.CredentialsFromJSON(ctx, credentialsJSON, cloudresourcemanager.CloudPlatformScope)
+		if err != nil {
+			return nil, fmt.Errorf("failed to obtain credentials: %w", err)
+		}
+		oauthClient = func(name string) *http.Client {
+			client := oauth2.NewClient(ctx, creds.TokenSource)
+			client.Transport = otelhttp.NewTransport(client.Transport, otelhttp.WithSpanNameFormatter(func(operation string, r *http.Request) string {
+				return fmt.Sprintf("%s %s", name, operation)
+			}))
+			return client
+		}
+	}
+
+	if credentialsJSON != nil {
+		bqOAuthClient := oauthClient("BigQuery")
+		opts = append(opts, option.WithHTTPClient(bqOAuthClient))
+	}
+	bqClient, err := bq.NewClient(ctx, projectID, opts...)
 	if err != nil {
 		return nil, err
 	}
 
-	apiHTTPClient := opentelemetry.NewHttpClient(ctx, "BQAPIClient", opts...)
-	apiClient, err := bqApi.NewService(ctx, option.WithHTTPClient(apiHTTPClient))
+	if credentialsJSON != nil {
+		bqAPIOAuthClient := oauthClient("BigQuery API")
+		opts = append(opts, option.WithHTTPClient(bqAPIOAuthClient))
+	}
+	apiService, err := bqApi.NewService(ctx, opts...)
 	if err != nil {
 		return nil, err
 	}
 
-	iamHTTPClient := opentelemetry.NewHttpClient(ctx, "IAMserviceClient", opts...)
-	iamService, err := iam.NewService(ctx, option.WithHTTPClient(iamHTTPClient))
+	if credentialsJSON != nil {
+		iamOAuthClient := oauthClient("IAM")
+		opts = append(opts, option.WithHTTPClient(iamOAuthClient))
+	}
+	iamService, err := iam.NewService(ctx, opts...)
 	if err != nil {
 		return nil, err
 	}
 
-	crmHTTPClient := opentelemetry.NewHttpClient(ctx, "CRMClient", opts...)
-	crmService, err := cloudresourcemanager.NewService(ctx, option.WithHTTPClient(crmHTTPClient))
+	if credentialsJSON != nil {
+		crmOAuthClient := oauthClient("Cloud Resource Manager")
+		opts = append(opts, option.WithHTTPClient(crmOAuthClient))
+	}
+	crmService, err := cloudresourcemanager.NewService(ctx, opts...)
 	if err != nil {
 		return nil, err
 	}
@@ -53,7 +85,7 @@ func NewBigQueryClient(projectID string, opts ...option.ClientOption) (*bigQuery
 		projectID:  projectID,
 		client:     bqClient,
 		iamService: iamService,
-		apiClient:  apiClient,
+		apiClient:  apiService,
 		crmService: crmService,
 	}, nil
 }

diff --git a/plugins/providers/bigquery/client_test.go b/plugins/providers/bigquery/client_test.go
@@ -76,7 +76,7 @@ func (s *ClientTestSuite) TestGetDatasets() {
 	}))
 	defer ts.Close()
 
-	client, err := bigquery.NewBigQueryClient(projectID, option.WithoutAuthentication(), option.WithEndpoint(ts.URL))
+	client, err := bigquery.NewBigQueryClient(projectID, nil, option.WithoutAuthentication(), option.WithEndpoint(ts.URL))
 	s.Require().NoError(err)
 
 	datasets, err := client.GetDatasets(context.Background())
@@ -145,7 +145,7 @@ func (s *ClientTestSuite) TestGetTables() {
 	}))
 	defer ts.Close()
 
-	client, err := bigquery.NewBigQueryClient(projectID, option.WithoutAuthentication(), option.WithEndpoint(ts.URL))
+	client, err := bigquery.NewBigQueryClient(projectID, nil, option.WithoutAuthentication(), option.WithEndpoint(ts.URL))
 	s.Require().NoError(err)
 
 	tables, err := client.GetTables(context.Background(), datasetID)

diff --git a/plugins/providers/bigquery/config.go b/plugins/providers/bigquery/config.go
@@ -5,13 +5,13 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
+	"log"
 	"strings"
 
 	"github.com/go-playground/validator/v10"
 	"github.com/goto/guardian/domain"
 	"github.com/goto/guardian/utils"
 	"github.com/mitchellh/mapstructure"
-	"google.golang.org/api/option"
 )
 
 const (
@@ -135,9 +135,9 @@ func (c *Config) parseAndValidate(ctx context.Context) error {
 	} else {
 		c.ProviderConfig.Credentials = credentials
 	}
-
+	log.Printf("%s", credentials.ServiceAccountKey)
 	projectID := strings.Replace(credentials.ResourceName, "projects/", "", 1)
-	client, err := NewBigQueryClient(projectID, option.WithCredentialsJSON([]byte(credentials.ServiceAccountKey)))
+	client, err := NewBigQueryClient(projectID, []byte(credentials.ServiceAccountKey))
 	if err != nil {
 		return err
 	}

diff --git a/plugins/providers/bigquery/provider.go b/plugins/providers/bigquery/provider.go
@@ -21,7 +21,6 @@ import (
 	"github.com/patrickmn/go-cache"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/api/logging/v2"
-	"google.golang.org/api/option"
 )
 
 var (
@@ -539,7 +538,7 @@ func (p *Provider) getBigQueryClient(credentials Credentials) (BigQueryClient, e
 	}
 
 	credentials.Decrypt(p.encryptor)
-	client, err := NewBigQueryClient(projectID, option.WithCredentialsJSON([]byte(credentials.ServiceAccountKey)))
+	client, err := NewBigQueryClient(projectID, []byte(credentials.ServiceAccountKey))
 	if err != nil {
 		return nil, err
 	}

diff --git a/plugins/providers/gcloudiam/client.go b/plugins/providers/gcloudiam/client.go
@@ -3,10 +3,13 @@ package gcloudiam
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"strings"
 
 	"github.com/goto/guardian/domain"
-	"github.com/goto/guardian/pkg/opentelemetry"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google"
 	"google.golang.org/api/cloudresourcemanager/v1"
 	"google.golang.org/api/iam/v1"
 	"google.golang.org/api/option"
@@ -25,17 +28,29 @@ type iamClient struct {
 
 func newIamClient(credentialsJSON []byte, resourceName string) (*iamClient, error) {
 	ctx := context.Background()
+	creds, err := google.CredentialsFromJSON(ctx, credentialsJSON)
+	if err != nil {
+		return nil, fmt.Errorf("failed to obtain credentials: %w", err)
+	}
 
-	crmHTTPClient := opentelemetry.NewHttpClient(ctx, "CloudResourceManagerClient", option.WithCredentialsJSON(credentialsJSON))
-	cloudResourceManagerService, err := cloudresourcemanager.NewService(ctx, option.WithHTTPClient(crmHTTPClient))
+	oauthClientCRM := oauth2.NewClient(ctx, creds.TokenSource)
+	oauthClientCRM.Transport = otelhttp.NewTransport(oauthClientCRM.Transport, otelhttp.WithSpanNameFormatter(func(operation string, r *http.Request) string {
+		return fmt.Sprintf("CloudResourceManagementClient %s", operation)
+	}))
+
+	cloudResourceManagerService, err := cloudresourcemanager.NewService(ctx, option.WithHTTPClient(oauthClientCRM))
 	if err != nil {
-		return nil, fmt.Errorf("failed to initialize Cloud Resource Manager service: %w", err)
+		return nil, fmt.Errorf("failed to create Cloud Resource Manager service: %w", err)
 	}
 
-	iamHTTPClient := opentelemetry.NewHttpClient(ctx, "IAMClient", option.WithCredentialsJSON(credentialsJSON))
-	iamService, err := iam.NewService(ctx, option.WithHTTPClient(iamHTTPClient))
+	oauthClientIAM := oauth2.NewClient(ctx, creds.TokenSource)
+	oauthClientIAM.Transport = otelhttp.NewTransport(oauthClientIAM.Transport, otelhttp.WithSpanNameFormatter(func(operation string, r *http.Request) string {
+		return fmt.Sprintf("IAMClient %s", operation)
+	}))
+
+	iamService, err := iam.NewService(ctx, option.WithHTTPClient(oauthClientIAM))
 	if err != nil {
-		return nil, fmt.Errorf("failed to initialize IAM service: %w", err)
+		return nil, fmt.Errorf("failed to create IAM service: %w", err)
 	}
 
 	return &iamClient{

diff --git a/plugins/providers/gcs/client.go b/plugins/providers/gcs/client.go
@@ -3,13 +3,16 @@ package gcs
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"strings"
 
 	"cloud.google.com/go/iam"
 	"cloud.google.com/go/storage"
 	"github.com/goto/guardian/domain"
-	"github.com/goto/guardian/pkg/opentelemetry"
 	"github.com/goto/guardian/utils"
+	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
+	"golang.org/x/oauth2"
+	"golang.org/x/oauth2/google"
 	"golang.org/x/sync/errgroup"
 	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
@@ -21,15 +24,22 @@ type gcsClient struct {
 }
 
 func newGCSClient(ctx context.Context, projectID string, credentialsJSON []byte) (*gcsClient, error) {
-
-	httpClient := opentelemetry.NewHttpClient(ctx, "GCSClient", nil)
-	client, err := storage.NewClient(ctx, option.WithHTTPClient(httpClient), option.WithCredentialsJSON(credentialsJSON))
+	creds, err := google.CredentialsFromJSON(ctx, credentialsJSON)
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to obtain credentials: %w", err)
 	}
 
+	client := oauth2.NewClient(ctx, creds.TokenSource)
+	client.Transport = otelhttp.NewTransport(client.Transport, otelhttp.WithSpanNameFormatter(func(operation string, r *http.Request) string {
+		return fmt.Sprintf("GCSClient %s", operation)
+	}))
+
+	clientService, err := storage.NewClient(ctx, option.WithHTTPClient(client))
+	if err != nil {
+		return nil, fmt.Errorf("failed to create GCS client: %w", err)
+	}
 	return &gcsClient{
-		client:    client,
+		client:    clientService,
 		projectID: projectID,
 	}, nil
 }

diff --git a/plugins/providers/gitlab/provider.go b/plugins/providers/gitlab/provider.go
@@ -10,7 +10,6 @@ import (
 	pv "github.com/goto/guardian/core/provider"
 	"github.com/goto/guardian/domain"
 	"github.com/goto/guardian/pkg/log"
-	"github.com/goto/guardian/pkg/opentelemetry"
 	"github.com/goto/guardian/utils"
 	"github.com/xanzy/go-gitlab"
 	"golang.org/x/sync/errgroup"
@@ -272,8 +271,7 @@ func (p *provider) getClient(pc domain.ProviderConfig) (*gitlab.Client, error) {
 		return nil, fmt.Errorf("unable to decrypt credentials: %w", err)
 	}
 
-	gitlabHTTPClient := opentelemetry.NewHttpClient(context.Background(), "GitlabClient")
-	client, err := gitlab.NewClient(creds.AccessToken, gitlab.WithBaseURL(creds.Host), gitlab.WithHTTPClient(gitlabHTTPClient))
+	client, err := gitlab.NewClient(creds.AccessToken, gitlab.WithBaseURL(creds.Host))
 	if err != nil {
 		return nil, fmt.Errorf("unable to create gitlab client: %w", err)
 	}