Skip to content

Commit

Permalink
fix: refactor new relic telemetry to a common method (#250)
Browse files Browse the repository at this point in the history
Signed-off-by: Kush Sharma <thekushsharma@gmail.com>
  • Loading branch information
kushsharma committed Jun 1, 2023
1 parent 8ea72c5 commit 89a517e
Show file tree
Hide file tree
Showing 15 changed files with 115 additions and 1,094 deletions.
72 changes: 14 additions & 58 deletions internal/store/postgres/flow_repository.go
Expand Up @@ -8,7 +8,6 @@ import (

"github.com/doug-martin/goqu/v9"
"github.com/google/uuid"
newrelic "github.com/newrelic/go-agent"
"github.com/odpf/salt/log"
"github.com/odpf/shield/core/authenticate"
"github.com/odpf/shield/pkg/db"
Expand Down Expand Up @@ -65,18 +64,7 @@ func (s *FlowRepository) Set(ctx context.Context, flow *authenticate.Flow) error
}

var flowModel Flow
if err = s.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_FLOWS,
Operation: "Upsert",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = s.dbc.WithTimeout(ctx, TABLE_FLOWS, "Set", func(ctx context.Context) error {
return s.dbc.QueryRowxContext(ctx, query, params...).StructScan(&flowModel)
}); err != nil {
err = checkPostgresError(err)
Expand All @@ -98,18 +86,7 @@ func (s *FlowRepository) Get(ctx context.Context, id uuid.UUID) (*authenticate.F
return nil, fmt.Errorf("%w: %s", queryErr, err)
}

if err = s.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_FLOWS,
Operation: "Get",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = s.dbc.WithTimeout(ctx, TABLE_FLOWS, "Get", func(ctx context.Context) error {
return s.dbc.QueryRowxContext(ctx, query, params...).StructScan(&flowModel)
}); err != nil {
err = checkPostgresError(err)
Expand All @@ -130,30 +107,20 @@ func (s *FlowRepository) Delete(ctx context.Context, id uuid.UUID) error {
return fmt.Errorf("%w: %s", queryErr, err)
}

return s.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_FLOWS,
Operation: "Delete",
StartTime: nrCtx.StartSegmentNow(),
return s.dbc.WithTimeout(ctx, TABLE_FLOWS,
"Delete", func(ctx context.Context) error {
result, err := s.dbc.ExecContext(ctx, query, params...)
if err != nil {
err = checkPostgresError(err)
return fmt.Errorf("%w: %s", dbErr, err)
}
defer nr.End()
}

result, err := s.dbc.ExecContext(ctx, query, params...)
if err != nil {
err = checkPostgresError(err)
return fmt.Errorf("%w: %s", dbErr, err)
}

if count, _ := result.RowsAffected(); count > 0 {
return nil
}
if count, _ := result.RowsAffected(); count > 0 {
return nil
}

return fmt.Errorf("no entry to delete")
})
return fmt.Errorf("no entry to delete")
})
}

func (s *FlowRepository) DeleteExpiredFlows(ctx context.Context) error {
Expand All @@ -167,18 +134,7 @@ func (s *FlowRepository) DeleteExpiredFlows(ctx context.Context) error {
return fmt.Errorf("%w: %s", queryErr, err)
}

return s.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_FLOWS,
Operation: "DeleteExpiredFlows",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

return s.dbc.WithTimeout(ctx, TABLE_FLOWS, "DeleteExpiredFlows", func(ctx context.Context) error {
result, err := s.dbc.ExecContext(ctx, query, params...)
if err != nil {
err = checkPostgresError(err)
Expand Down
118 changes: 9 additions & 109 deletions internal/store/postgres/group_repository.go
Expand Up @@ -13,7 +13,6 @@ import (
"github.com/odpf/shield/core/user"

"github.com/doug-martin/goqu/v9"
newrelic "github.com/newrelic/go-agent"
"github.com/odpf/shield/core/group"
"github.com/odpf/shield/core/namespace"
"github.com/odpf/shield/core/organization"
Expand Down Expand Up @@ -54,18 +53,7 @@ func (r GroupRepository) GetByID(ctx context.Context, id string) (group.Group, e
}

var groupModel Group
if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "GetByID",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "GetByID", func(ctx context.Context) error {
return r.dbc.GetContext(ctx, &groupModel, query, params...)
}); err != nil {
err = checkPostgresError(err)
Expand Down Expand Up @@ -106,18 +94,7 @@ func (r GroupRepository) GetByIDs(ctx context.Context, groupIDs []string) ([]gro

// this query will return empty array of groups if no UUID is not matched
// TODO: check and fox what to do in this scenerio
if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "GetByIDs",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "GetByIDs", func(ctx context.Context) error {
return r.dbc.SelectContext(ctx, &fetchedGroups, query, params...)
}); err != nil {
err = checkPostgresError(err)
Expand Down Expand Up @@ -169,18 +146,7 @@ func (r GroupRepository) Create(ctx context.Context, grp group.Group) (group.Gro
}

var groupModel Group
if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "Upsert",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "Upsert", func(ctx context.Context) error {
return r.dbc.QueryRowxContext(ctx, query, params...).StructScan(&groupModel)
}); err != nil {
err = checkPostgresError(err)
Expand Down Expand Up @@ -221,18 +187,7 @@ func (r GroupRepository) List(ctx context.Context, flt group.Filter) ([]group.Gr
}

var fetchedGroups []Group
if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "List",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "List", func(ctx context.Context) error {
return r.dbc.SelectContext(ctx, &fetchedGroups, query, params...)
}); err != nil {
err = checkPostgresError(err)
Expand Down Expand Up @@ -287,18 +242,7 @@ func (r GroupRepository) UpdateByID(ctx context.Context, grp group.Group) (group
}

var groupModel Group
if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "Update",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "Update", func(ctx context.Context) error {
return r.dbc.QueryRowxContext(ctx, query, params...).StructScan(&groupModel)
}); err != nil {
err = checkPostgresError(err)
Expand Down Expand Up @@ -362,18 +306,7 @@ func (r GroupRepository) ListUserGroups(ctx context.Context, userID string, role
}

var fetchedGroups []Group
if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "ListUserGroups",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "ListUserGroups", func(ctx context.Context) error {
return r.dbc.SelectContext(ctx, &fetchedGroups, query, params...)
}); err != nil {
if errors.Is(err, sql.ErrNoRows) {
Expand Down Expand Up @@ -419,18 +352,7 @@ func (r GroupRepository) ListGroupRelations(ctx context.Context, objectId string
}

var fetchedRelations []Relation
if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "ListGroupRelations",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "ListGroupRelations", func(ctx context.Context) error {
return r.dbc.SelectContext(ctx, &fetchedRelations, query, params...)
}); err != nil {
// List should return empty list and no error instead
Expand Down Expand Up @@ -461,18 +383,7 @@ func (r GroupRepository) SetState(ctx context.Context, id string, state group.St
return fmt.Errorf("%w: %s", queryErr, err)
}

if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "SetState",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "SetState", func(ctx context.Context) error {
if _, err = r.dbc.DB.ExecContext(ctx, query, params...); err != nil {
return err
}
Expand All @@ -499,18 +410,7 @@ func (r GroupRepository) Delete(ctx context.Context, id string) error {
return fmt.Errorf("%w: %s", queryErr, err)
}

if err = r.dbc.WithTimeout(ctx, func(ctx context.Context) error {
nrCtx := newrelic.FromContext(ctx)
if nrCtx != nil {
nr := newrelic.DatastoreSegment{
Product: newrelic.DatastorePostgres,
Collection: TABLE_GROUPS,
Operation: "Delete",
StartTime: nrCtx.StartSegmentNow(),
}
defer nr.End()
}

if err = r.dbc.WithTimeout(ctx, TABLE_GROUPS, "Delete", func(ctx context.Context) error {
if _, err = r.dbc.DB.ExecContext(ctx, query, params...); err != nil {
return err
}
Expand Down

0 comments on commit 89a517e

Please sign in to comment.