refactor: Do request logic

Author: cade-conklin

v3/integrations/nroci/nroci_nosql.go

@@ -3,9 +3,11 @@
 package nroci
 
 import (
+	"context"
 	"fmt"
 	"time"
 
+	"github.com/newrelic/go-agent/v3/newrelic"
 	"github.com/oracle/nosql-go-sdk/nosqldb"
 )
 
@@ -58,6 +60,10 @@ type ClientWrapper struct {
 	Client OCIClient
 }
 
+type ClientResponseWrapper[T any] struct {
+	ClientResponse T
+}
+
 func NRDefaultConfig() *ConfigWrapper {
 	cfg := nosqldb.Config{}
 	return &ConfigWrapper{
@@ -74,3 +80,54 @@ func NRCreateClient(cfg *ConfigWrapper) (*ClientWrapper, error) {
 		Client: client,
 	}, nil
 }
+
+// executeWithDatastoreSegment is a generic helper function that executes a query with a given function from the
+// OCI Client.  It takes a type parameter T as any because of the different response types that are used within the
+// OCI Client.  This function will take the transaction from the context (if it exists) and create a Datastore Segment.
+// It will then call whatever client function has been passed in.
+func executeWithDatastoreSegment[T any](
+	ctx context.Context,
+	req *nosqldb.TableRequest,
+	fn func() (T, error),
+) (*ClientResponseWrapper[T], error) {
+
+	txn := newrelic.FromContext(ctx)
+	if txn == nil {
+		return nil, fmt.Errorf("error executing DoTableRequest, no transaction")
+	}
+
+	sgmt := newrelic.DatastoreSegment{
+		StartTime:    txn.StartSegmentNow(),
+		Product:      newrelic.DatastoreOracle,
+		Collection:   req.TableName,
+		Operation:    req.Namespace,
+		DatabaseName: req.Namespace,
+	}
+
+	responseWrapper := ClientResponseWrapper[T]{}
+	res, err := fn() // call the client function
+	responseWrapper.ClientResponse = res
+	if err != nil {
+		return &responseWrapper, fmt.Errorf("error making request: %s", err.Error())
+	}
+	// 4. End segment
+	sgmt.End()
+	// 5. Return result
+	return &responseWrapper, nil
+}
+
+// Wrapper for nosqldb.Client.DoTableRequest.  Provide the ClientWrapper and Context as parameters in addition to the nosqldbTableRequest.
+// Returns a ClientResponseWrapper[*nosqldbTableResult] and error.
+func NRDoTableRequest(cw *ClientWrapper, ctx context.Context, req *nosqldb.TableRequest) (*ClientResponseWrapper[*nosqldb.TableResult], error) {
+	return executeWithDatastoreSegment(ctx, req, func() (*nosqldb.TableResult, error) {
+		return cw.Client.DoTableRequest(req)
+	})
+}
+
+// Wrapper for nosqldb.Client.DoTableRequestWait.  Provide the ClientWrapper and Context as parameters in addition to the nosqldbTableRequest,
+// timeout, and pollInterval. Returns a ClientResponseWrapper[*nosqldbTableResult] and error.
+func NRDoTableRequestAndWait(cw *ClientWrapper, ctx context.Context, req *nosqldb.TableRequest, timeout time.Duration, pollInterval time.Duration) (*ClientResponseWrapper[*nosqldb.TableResult], error) {
+	return executeWithDatastoreSegment(ctx, req, func() (*nosqldb.TableResult, error) {
+		return cw.Client.DoTableRequestAndWait(req, timeout, pollInterval)
+	})
+}