Added newRelic segments to gRPC API calls

core/schema/service.go

@@ -83,8 +83,6 @@ func (s *Service) checkCompatibility(ctx context.Context, nsName, schemaName, fo
 }
 
 func (s *Service) Create(ctx context.Context, nsName string, schemaName string, metadata *Metadata, data []byte) (SchemaInfo, error) {
-	endFunc := s.newrelic.StartGenericSegment(ctx, "Create Schema Info")
-	defer endFunc()
 	var scInfo SchemaInfo
 	ns, err := s.namespaceService.Get(ctx, nsName)
 	if err != nil {
@@ -114,17 +112,13 @@ func (s *Service) Create(ctx context.Context, nsName string, schemaName string,
 }
 
 func (s *Service) withMetadata(ctx context.Context, namespace, schemaName string, getData func() ([]byte, error)) (*Metadata, []byte, error) {
-	endFunc := s.newrelic.StartGenericSegment(ctx, "GetMetaData")
-	defer endFunc()
 	var data []byte
 	meta, err := s.repo.GetMetadata(ctx, namespace, schemaName)
 	if err != nil {
 		return meta, data, err
 	}
 
-	dataSegmentEndFunc := s.newrelic.StartGenericSegment(ctx, "GetData")
 	data, err = getData()
-	dataSegmentEndFunc()
 	return meta, data, err
 }
 

internal/api/namespace.go

@@ -2,7 +2,6 @@ package api
 
 import (
 	"context"
-
 	"github.com/goto/stencil/core/namespace"
 	stencilv1beta1 "github.com/goto/stencil/proto/v1beta1"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -30,23 +29,31 @@ func namespaceToProto(ns namespace.Namespace) *stencilv1beta1.Namespace {
 
 // CreateNamespace handler for creating namespace
 func (a *API) CreateNamespace(ctx context.Context, in *stencilv1beta1.CreateNamespaceRequest) (*stencilv1beta1.CreateNamespaceResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "CreateNamespace")
+	defer endFunc()
 	ns := createNamespaceRequestToNamespace(in)
 	newNamespace, err := a.namespace.Create(ctx, ns)
 	return &stencilv1beta1.CreateNamespaceResponse{Namespace: namespaceToProto(newNamespace)}, err
 }
 
 func (a *API) UpdateNamespace(ctx context.Context, in *stencilv1beta1.UpdateNamespaceRequest) (*stencilv1beta1.UpdateNamespaceResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "UpdateNamespace")
+	defer endFunc()
 	ns, err := a.namespace.Update(ctx, namespace.Namespace{ID: in.GetId(), Format: in.GetFormat().String(), Compatibility: in.GetCompatibility().String(), Description: in.GetDescription()})
 	return &stencilv1beta1.UpdateNamespaceResponse{Namespace: namespaceToProto(ns)}, err
 }
 
 func (a *API) GetNamespace(ctx context.Context, in *stencilv1beta1.GetNamespaceRequest) (*stencilv1beta1.GetNamespaceResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "GetNamespace")
+	defer endFunc()
 	namespace, err := a.namespace.Get(ctx, in.GetId())
 	return &stencilv1beta1.GetNamespaceResponse{Namespace: namespaceToProto(namespace)}, err
 }
 
 // ListNamespaces handler for returning list of available namespaces
 func (a *API) ListNamespaces(ctx context.Context, in *stencilv1beta1.ListNamespacesRequest) (*stencilv1beta1.ListNamespacesResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "ListNamespaces")
+	defer endFunc()
 	namespaces, err := a.namespace.List(ctx)
 	var nsp []*stencilv1beta1.Namespace
 	for _, n := range namespaces {
@@ -56,6 +63,8 @@ func (a *API) ListNamespaces(ctx context.Context, in *stencilv1beta1.ListNamespa
 }
 
 func (a *API) DeleteNamespace(ctx context.Context, in *stencilv1beta1.DeleteNamespaceRequest) (*stencilv1beta1.DeleteNamespaceResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "DeleteNamespace")
+	defer endFunc()
 	err := a.namespace.Delete(ctx, in.GetId())
 	message := "success"
 	if err != nil {

internal/api/schema.go

@@ -23,6 +23,8 @@ func schemaToProto(s schema.Schema) *stencilv1beta1.Schema {
 }
 
 func (a *API) CreateSchema(ctx context.Context, in *stencilv1beta1.CreateSchemaRequest) (*stencilv1beta1.CreateSchemaResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "CreateSchema")
+	defer endFunc()
 	metadata := &schema.Metadata{Format: in.GetFormat().String(), Compatibility: in.GetCompatibility().String()}
 	sc, err := a.schema.Create(ctx, in.NamespaceId, in.SchemaId, metadata, in.GetData())
 	return &stencilv1beta1.CreateSchemaResponse{
@@ -57,12 +59,16 @@ func (a *API) HTTPUpload(w http.ResponseWriter, req *http.Request, pathParams ma
 }
 
 func (a *API) CheckCompatibility(ctx context.Context, req *stencilv1beta1.CheckCompatibilityRequest) (*stencilv1beta1.CheckCompatibilityResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "CheckCompatibility")
+	defer endFunc()
 	resp := &stencilv1beta1.CheckCompatibilityResponse{}
 	err := a.schema.CheckCompatibility(ctx, req.GetNamespaceId(), req.GetSchemaId(), req.GetCompatibility().String(), req.GetData())
 	return resp, err
 }
 
 func (a *API) HTTPCheckCompatibility(w http.ResponseWriter, req *http.Request, pathParams map[string]string) error {
+	endFunc := a.newrelic.StartGenericSegment(req.Context(), "CheckCompatibility")
+	defer endFunc()
 	data, err := io.ReadAll(req.Body)
 	if err != nil {
 		return err
@@ -75,6 +81,8 @@ func (a *API) HTTPCheckCompatibility(w http.ResponseWriter, req *http.Request, p
 }
 
 func (a *API) ListSchemas(ctx context.Context, in *stencilv1beta1.ListSchemasRequest) (*stencilv1beta1.ListSchemasResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "ListSchemas")
+	defer endFunc()
 	schemas, err := a.schema.List(ctx, in.Id)
 
 	var ss []*stencilv1beta1.Schema
@@ -85,6 +93,8 @@ func (a *API) ListSchemas(ctx context.Context, in *stencilv1beta1.ListSchemasReq
 }
 
 func (a *API) GetLatestSchema(ctx context.Context, in *stencilv1beta1.GetLatestSchemaRequest) (*stencilv1beta1.GetLatestSchemaResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "GetLatestSchema")
+	defer endFunc()
 	_, data, err := a.schema.GetLatest(ctx, in.NamespaceId, in.SchemaId)
 	return &stencilv1beta1.GetLatestSchemaResponse{
 		Data: data,
@@ -101,6 +111,8 @@ func (a *API) HTTPLatestSchema(w http.ResponseWriter, req *http.Request, pathPar
 }
 
 func (a *API) GetSchema(ctx context.Context, in *stencilv1beta1.GetSchemaRequest) (*stencilv1beta1.GetSchemaResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "GetSchema")
+	defer endFunc()
 	_, data, err := a.schema.Get(ctx, in.NamespaceId, in.SchemaId, in.GetVersionId())
 	return &stencilv1beta1.GetSchemaResponse{
 		Data: data,
@@ -122,11 +134,15 @@ func (a *API) HTTPGetSchema(w http.ResponseWriter, req *http.Request, pathParams
 }
 
 func (a *API) ListVersions(ctx context.Context, in *stencilv1beta1.ListVersionsRequest) (*stencilv1beta1.ListVersionsResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "ListVersions")
+	defer endFunc()
 	versions, err := a.schema.ListVersions(ctx, in.NamespaceId, in.SchemaId)
 	return &stencilv1beta1.ListVersionsResponse{Versions: versions}, err
 }
 
 func (a *API) GetSchemaMetadata(ctx context.Context, in *stencilv1beta1.GetSchemaMetadataRequest) (*stencilv1beta1.GetSchemaMetadataResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "GetSchemaMetadata")
+	defer endFunc()
 	meta, err := a.schema.GetMetadata(ctx, in.NamespaceId, in.SchemaId)
 	return &stencilv1beta1.GetSchemaMetadataResponse{
 		Format:        stencilv1beta1.Schema_Format(stencilv1beta1.Schema_Format_value[meta.Format]),
@@ -136,6 +152,8 @@ func (a *API) GetSchemaMetadata(ctx context.Context, in *stencilv1beta1.GetSchem
 }
 
 func (a *API) UpdateSchemaMetadata(ctx context.Context, in *stencilv1beta1.UpdateSchemaMetadataRequest) (*stencilv1beta1.UpdateSchemaMetadataResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "UpdateSchemaMetadata")
+	defer endFunc()
 	meta, err := a.schema.UpdateMetadata(ctx, in.NamespaceId, in.SchemaId, &schema.Metadata{
 		Compatibility: in.Compatibility.String(),
 	})
@@ -147,6 +165,8 @@ func (a *API) UpdateSchemaMetadata(ctx context.Context, in *stencilv1beta1.Updat
 }
 
 func (a *API) DeleteSchema(ctx context.Context, in *stencilv1beta1.DeleteSchemaRequest) (*stencilv1beta1.DeleteSchemaResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "DeleteSchema")
+	defer endFunc()
 	err := a.schema.Delete(ctx, in.NamespaceId, in.SchemaId)
 	message := "success"
 	if err != nil {
@@ -158,6 +178,8 @@ func (a *API) DeleteSchema(ctx context.Context, in *stencilv1beta1.DeleteSchemaR
 }
 
 func (a *API) DeleteVersion(ctx context.Context, in *stencilv1beta1.DeleteVersionRequest) (*stencilv1beta1.DeleteVersionResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "DeleteVersion")
+	defer endFunc()
 	err := a.schema.DeleteVersion(ctx, in.NamespaceId, in.SchemaId, in.GetVersionId())
 	message := "success"
 	if err != nil {

internal/api/search.go

@@ -9,6 +9,8 @@ import (
 )
 
 func (a *API) Search(ctx context.Context, in *stencilv1beta1.SearchRequest) (*stencilv1beta1.SearchResponse, error) {
+	endFunc := a.newrelic.StartGenericSegment(ctx, "SearchSegment")
+	defer endFunc()
 	searchReq := &search.SearchRequest{
 		NamespaceID: in.GetNamespaceId(),
 		Query:       in.GetQuery(),

internal/server/server.go

@@ -40,11 +40,11 @@ func Start(cfg config.Config) {
 	ctx := context.Background()
 
 	db := postgres.NewStore(cfg.DB.ConnectionString)
-
-	namespaceRepository := postgres.NewNamespaceRepository(db)
+	newRelic := &newRelic2.NewRelic{}
+	namespaceRepository := postgres.NewNamespaceRepository(db, newRelic)
 	namespaceService := namespace.NewService(namespaceRepository)
 
-	schemaRepository := postgres.NewSchemaRepository(db)
+	schemaRepository := postgres.NewSchemaRepository(db, newRelic)
 	cache, err := ristretto.NewCache(&ristretto.Config{
 		NumCounters: 1000,
 		MaxCost:     cfg.CacheSizeInMB << 20,
@@ -53,10 +53,10 @@ func Start(cfg config.Config) {
 	if err != nil {
 		panic(err)
 	}
-	newRelic := &newRelic2.NewRelic{}
+
 	schemaService := schema.NewService(schemaRepository, provider.NewSchemaProvider(), namespaceService, cache, newRelic)
 
-	searchRepository := postgres.NewSearchRepository(db)
+	searchRepository := postgres.NewSearchRepository(db, newRelic)
 	searchService := search.NewService(searchRepository)
 
 	api := api.NewAPI(namespaceService, schemaService, searchService, newRelic)

internal/store/postgres/namespace_repository.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"context"
+	newrelic2 "github.com/goto/stencil/pkg/newrelic"
 
 	"github.com/georgysavva/scany/pgxscan"
 	"github.com/goto/stencil/core/namespace"
@@ -32,40 +33,52 @@ RETURNING *
 `
 
 type NamespaceRepository struct {
-	db *DB
+	db       *DB
+	newrelic newrelic2.Service
 }
 
-func NewNamespaceRepository(dbc *DB) *NamespaceRepository {
+func NewNamespaceRepository(dbc *DB, nr newrelic2.Service) *NamespaceRepository {
 	return &NamespaceRepository{
-		db: dbc,
+		db:       dbc,
+		newrelic: nr,
 	}
 }
 
 func (r *NamespaceRepository) Create(ctx context.Context, ns namespace.Namespace) (namespace.Namespace, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "DatabaseCall")
+	defer endFunc()
 	newNamespace := namespace.Namespace{}
 	err := pgxscan.Get(ctx, r.db, &newNamespace, namespaceInsertQuery, ns.ID, ns.Format, ns.Compatibility, ns.Description)
 	return newNamespace, wrapError(err, ns.ID)
 }
 
 func (r *NamespaceRepository) Update(ctx context.Context, ns namespace.Namespace) (namespace.Namespace, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "DatabaseCall")
+	defer endFunc()
 	newNamespace := namespace.Namespace{}
 	err := pgxscan.Get(ctx, r.db, &newNamespace, namespaceUpdateQuery, ns.ID, ns.Format, ns.Compatibility, ns.Description)
 	return newNamespace, wrapError(err, ns.ID)
 }
 
 func (r *NamespaceRepository) Get(ctx context.Context, id string) (namespace.Namespace, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "DatabaseCall")
+	defer endFunc()
 	newNamespace := namespace.Namespace{}
 	err := pgxscan.Get(ctx, r.db, &newNamespace, namespaceGetQuery, id)
 	return newNamespace, wrapError(err, id)
 }
 
 func (r *NamespaceRepository) Delete(ctx context.Context, id string) error {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "DatabaseCall")
+	defer endFunc()
 	_, err := r.db.Exec(ctx, namespaceDeleteQuery, id)
 	r.db.Exec(ctx, deleteOrphanedData)
 	return wrapError(err, id)
 }
 
 func (r *NamespaceRepository) List(ctx context.Context) ([]namespace.Namespace, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "DatabaseCall")
+	defer endFunc()
 	var namespaces []namespace.Namespace
 	err := pgxscan.Select(ctx, r.db, &namespaces, namespaceListQuery)
 	return namespaces, wrapError(err, "")

internal/store/postgres/namespace_repository_test.go

@@ -2,6 +2,7 @@ package postgres_test
 
 import (
 	"context"
+	newRelic2 "github.com/goto/stencil/pkg/newrelic"
 	"os"
 	"testing"
 
@@ -21,7 +22,8 @@ func getNamespaceStore(t *testing.T) *postgres.NamespaceRepository {
 	err := postgres.Migrate(connectionString)
 	assert.Nil(t, err)
 	dbc := postgres.NewStore(connectionString)
-	return postgres.NewNamespaceRepository(dbc)
+	newRelic := &newRelic2.NewRelic{}
+	return postgres.NewNamespaceRepository(dbc, newRelic)
 }
 
 func TestNamespace(t *testing.T) {

internal/store/postgres/schema_repository.go

@@ -2,6 +2,7 @@ package postgres
 
 import (
 	"context"
+	newrelic2 "github.com/goto/stencil/pkg/newrelic"
 
 	"github.com/georgysavva/scany/pgxscan"
 	"github.com/goto/stencil/core/schema"
@@ -10,12 +11,14 @@ import (
 )
 
 type SchemaRepository struct {
-	db *DB
+	db       *DB
+	newrelic newrelic2.Service
 }
 
-func NewSchemaRepository(dbc *DB) *SchemaRepository {
+func NewSchemaRepository(dbc *DB, nr newrelic2.Service) *SchemaRepository {
 	return &SchemaRepository{
-		db: dbc,
+		db:       dbc,
+		newrelic: nr,
 	}
 }
 
@@ -25,6 +28,8 @@ type searchData struct {
 }
 
 func (r *SchemaRepository) Create(ctx context.Context, namespace string, schemaName string, metadata *schema.Metadata, versionID string, file *schema.SchemaFile) (int32, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- Create Schema")
+	defer endFunc()
 	var version int32
 	err := r.db.BeginFunc(ctx, func(t pgx.Tx) error {
 		vErr := t.QueryRow(ctx, getSchemaVersionByID, versionID).Scan(&version)
@@ -48,6 +53,8 @@ func (r *SchemaRepository) Create(ctx context.Context, namespace string, schemaN
 }
 
 func (r *SchemaRepository) Get(ctx context.Context, namespaceId, schemaName string, versionNumber int32) ([]byte, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- Get Schema")
+	defer endFunc()
 	var versionID string
 	var data []byte
 	if err := r.db.QueryRow(ctx, getVersionIDFromSchemaNameQuery, namespaceId, schemaName, versionNumber).Scan(&versionID); err != nil {
@@ -58,6 +65,8 @@ func (r *SchemaRepository) Get(ctx context.Context, namespaceId, schemaName stri
 }
 
 func (r *SchemaRepository) GetLatestVersion(ctx context.Context, namespaceId, schemaName string) (int32, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- Get Latest Version")
+	defer endFunc()
 	var version int32
 	if err := r.db.QueryRow(ctx, getLatestVersionIDFromSchemaNameQuery, namespaceId, schemaName).Scan(&version); err != nil {
 		return version, wrapError(err, "Latest version for %s - %s", namespaceId, schemaName)
@@ -66,37 +75,49 @@ func (r *SchemaRepository) GetLatestVersion(ctx context.Context, namespaceId, sc
 }
 
 func (r *SchemaRepository) GetMetadata(ctx context.Context, namespace, sc string) (*schema.Metadata, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- Get Meta Data")
+	defer endFunc()
 	var meta schema.Metadata
 	err := pgxscan.Get(ctx, r.db, &meta, getSchemaMetaQuery, namespace, sc)
 	return &meta, wrapError(err, "meta")
 }
 
 func (r *SchemaRepository) UpdateMetadata(ctx context.Context, namespace, sc string, in *schema.Metadata) (*schema.Metadata, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- Update Meta Data")
+	defer endFunc()
 	var meta schema.Metadata
 	err := pgxscan.Get(ctx, r.db, &meta, updateSchemaMetaQuery, namespace, sc, in.Compatibility)
 	return &meta, wrapError(err, "meta")
 }
 
 func (r *SchemaRepository) List(ctx context.Context, namespaceID string) ([]schema.Schema, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- List Schemas")
+	defer endFunc()
 	var schemas []schema.Schema
 	err := pgxscan.Select(ctx, r.db, &schemas, schemaListQuery, namespaceID)
 	return schemas, wrapError(err, "List schemas")
 }
 
 func (r *SchemaRepository) Delete(ctx context.Context, ns string, sc string) error {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- Delete Schema")
+	defer endFunc()
 	_, err := r.db.Exec(ctx, deleteSchemaQuery, ns, sc)
 	// Idempotent operation to clean orphaned data.
 	r.db.Exec(ctx, deleteOrphanedData)
 	return wrapError(err, "delete schema")
 }
 
 func (r *SchemaRepository) ListVersions(ctx context.Context, ns string, sc string) ([]int32, error) {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- List Versions of schema")
+	defer endFunc()
 	var versions []int32
 	err := pgxscan.Select(ctx, r.db, &versions, listVersionsQuery, ns, sc)
 	return versions, wrapError(err, "versions")
 }
 
 func (r *SchemaRepository) DeleteVersion(ctx context.Context, ns string, sc string, version int32) error {
+	endFunc := r.newrelic.StartGenericSegment(ctx, "Database Call- Delete schema version")
+	defer endFunc()
 	_, err := r.db.Exec(ctx, deleteVersionQuery, ns, sc, version)
 	// Idempotent operation to clean orphaned data.
 	r.db.Exec(ctx, deleteOrphanedData)