feat: instrument bigtable extractor

Author: Chief-Rishab

plugins/extractors/bigtable/bigtable.go

@@ -110,6 +110,9 @@ func (e *Extractor) Init(ctx context.Context, config plugins.Config) error {
 	if err != nil {
 		return err
 	}
+
+	client = WithInstanceAdminClientMW(e.config.ProjectID)(client)
+
 	e.instanceNames, err = instanceInfoGetter(ctx, client)
 	if err != nil {
 		return err
@@ -143,6 +146,9 @@ func (e *Extractor) getTablesInfo(ctx context.Context, emit plugins.Emit) error
 		if err != nil {
 			return err
 		}
+
+		adminClient = WithAdminClientMW(e.config.ProjectID, instance)(adminClient)
+
 		tables, _ := adminClient.Tables(ctx)
 		var wg sync.WaitGroup
 		for _, table := range tables {

plugins/extractors/bigtable/middleware.go

@@ -0,0 +1,123 @@
+package bigtable
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"cloud.google.com/go/bigtable"
+	"github.com/googleapis/gax-go/v2/apierror"
+	"github.com/raystack/meteor/utils"
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+)
+
+type AdminClientMW struct {
+	tableDuration  metric.Int64Histogram
+	tablesDuration metric.Int64Histogram
+	next           AdminClient
+	attributes     []attribute.KeyValue
+}
+
+type InstanceAdminClientMW struct {
+	instancesDuration metric.Int64Histogram
+	next              InstanceAdminClient
+	attributes        []attribute.KeyValue
+}
+
+func WithAdminClientMW(projectID, instanceName string) func(AdminClient) AdminClient {
+	meter := otel.Meter("github.com/raystack/meteor/plugins/extractors/bigtable")
+
+	tablesDuration, err := meter.Int64Histogram("meteor.bigtable.client.tables.duration", metric.WithUnit("ms"))
+	handleOtelErr(err)
+
+	tableDuration, err := meter.Int64Histogram("meteor.bigtable.client.table.duration", metric.WithUnit("ms"))
+	handleOtelErr(err)
+
+	return func(next AdminClient) AdminClient {
+		return &AdminClientMW{
+			tableDuration:  tableDuration,
+			tablesDuration: tablesDuration,
+			next:           next,
+			attributes: []attribute.KeyValue{
+				attribute.String("bt.project_id", projectID),
+				attribute.String("bt.instance_name", instanceName),
+			},
+		}
+	}
+}
+
+func WithInstanceAdminClientMW(projectID string) func(InstanceAdminClient) InstanceAdminClient {
+	instancesDuration, err := otel.Meter("github.com/raystack/meteor/plugins/extractors/bigtable").
+		Int64Histogram("meteor.bigtable.client.instances.duration", metric.WithUnit("ms"))
+	handleOtelErr(err)
+
+	return func(next InstanceAdminClient) InstanceAdminClient {
+		return &InstanceAdminClientMW{
+			instancesDuration: instancesDuration,
+			next:              next,
+			attributes: []attribute.KeyValue{
+				attribute.String("bt.project_id", projectID),
+			},
+		}
+	}
+}
+
+func (o *AdminClientMW) Tables(ctx context.Context) (res []string, err error) {
+	defer func(start time.Time) {
+		attrs := o.attributes
+		if err != nil {
+			attrs = append(attrs, attribute.String("bt.error_code", apiErrReason(err)))
+		}
+		o.tablesDuration.Record(ctx,
+			time.Since(start).Milliseconds(),
+			metric.WithAttributes(attrs...))
+	}(time.Now())
+
+	return o.next.Tables(ctx)
+}
+
+func (o *AdminClientMW) TableInfo(ctx context.Context, table string) (res *bigtable.TableInfo, err error) {
+	defer func(start time.Time) {
+		attrs := append(o.attributes, attribute.String("bt.table_name", table))
+		if err != nil {
+			attrs = append(attrs, attribute.String("bt.error_code", apiErrReason(err)))
+		}
+		o.tableDuration.Record(ctx,
+			time.Since(start).Milliseconds(),
+			metric.WithAttributes(attrs...))
+	}(time.Now())
+	return o.next.TableInfo(ctx, table)
+}
+
+func (o *InstanceAdminClientMW) Instances(ctx context.Context) (res []*bigtable.InstanceInfo, err error) {
+	defer func(start time.Time) {
+		attrs := o.attributes
+		if err != nil {
+			attrs = append(o.attributes, attribute.String("bt.error_code", apiErrReason(err)))
+		}
+
+		o.instancesDuration.Record(ctx,
+			time.Since(start).Milliseconds(),
+			metric.WithAttributes(attrs...))
+	}(time.Now())
+
+	return o.next.Instances(ctx)
+}
+
+func apiErrReason(err error) string {
+	reason := utils.StatusCode(err).String()
+	var apiErr *apierror.APIError
+	if errors.As(err, &apiErr) {
+		reason = apiErr.Reason()
+	}
+
+	return reason
+}
+
+func handleOtelErr(err error) {
+	if err != nil {
+		otel.Handle(err)
+	}
+}