diff --git a/server/middleware/authorization/authorization_manager.go b/server/middleware/authorization/authorization_manager.go index 3339d56..78a7d4f 100644 --- a/server/middleware/authorization/authorization_manager.go +++ b/server/middleware/authorization/authorization_manager.go @@ -55,8 +55,10 @@ func (m *AuthorizationManager) assertUserBanned() drip.StrictMiddlewareFunc { return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() - txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertUserBanned") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("AuthorizationManager.assertUserBanned") + defer segment.End() + } v := ctx.Value(drip_authentication.UserContextKey) userDetails, ok := v.(*drip_authentication.UserDetails) if !ok { @@ -83,8 +85,10 @@ func (m *AuthorizationManager) assertPublisherPermission( return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() - txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertPublisherPermission") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("AuthorizationManager.assertPublisherPermission") + defer segment.End() + } v := ctx.Value(drip_authentication.UserContextKey) userDetails, ok := v.(*drip_authentication.UserDetails) if !ok { @@ -121,8 +125,10 @@ func (m *AuthorizationManager) assertNodeBanned(extractor func(req interface{}) return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() - txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertNodeBanned") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("AuthorizationManager.assertNodeBanned") + defer segment.End() + } nodeID := extractor(request) err = m.Assertor.AssertNodeBanned(ctx, m.EntClient, nodeID) switch { @@ -145,8 +151,10 @@ func (m *AuthorizationManager) assertPublisherBanned(extractor func(req interfac return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() - txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertPublisherBanned") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("AuthorizationManager.assertPublisherBanned") + defer segment.End() + } publisherID := extractor(request) switch err = m.Assertor.AssertPublisherBanned(ctx, m.EntClient, publisherID); { @@ -172,8 +180,10 @@ func (m *AuthorizationManager) assertPersonalAccessTokenValid( return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() - txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertPersonalAccessTokenValid") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("AuthorizationManager.assertPersonalAccessTokenValid") + defer segment.End() + } pubID := extractorPublsherID(request) pat := extractorPAT(request) tokenValid, err := m.Assertor.IsPersonalAccessTokenValidForPublisher( @@ -200,8 +210,10 @@ func (m *AuthorizationManager) assertNodeBelongsToPublisher( return func(f strictecho.StrictEchoHandlerFunc, operationID string) strictecho.StrictEchoHandlerFunc { return func(c echo.Context, request interface{}) (response interface{}, err error) { ctx := c.Request().Context() - txn := m.NewRelicApp.StartTransaction("AuthorizationManager.assertNodeBelongsToPublisher") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("AuthorizationManager.assertNodeBelongsToPublisher") + defer segment.End() + } pubID := extractorPublsherID(request) nodeID := extractorNodeID(request) diff --git a/services/registry/registry_svc.go b/services/registry/registry_svc.go index 32c1345..e119824 100644 --- a/services/registry/registry_svc.go +++ b/services/registry/registry_svc.go @@ -66,6 +66,10 @@ func NewRegistryService(storageSvc storage.StorageService, pubsubService pubsub. // ListNodes retrieves a paginated list of nodes with optional filtering. func (s *RegistryService) ListNodes(ctx context.Context, client *ent.Client, page, limit int, filter *entity.NodeFilter) (*entity.ListNodesResult, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.ListNodes") + defer segment.End() + } // Ensure valid pagination parameters if page < 1 { page = 1 @@ -143,6 +147,10 @@ func (s *RegistryService) ListNodes(ctx context.Context, client *ent.Client, pag // ListPublishers queries the Publisher table with an optional user ID filter via PublisherPermission func (s *RegistryService) ListPublishers(ctx context.Context, client *ent.Client, filter *entity.PublisherFilter) ([]*ent.Publisher, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.ListPublishers") + defer segment.End() + } log.Ctx(ctx).Info().Msg("Listing publishers") query := client.Publisher.Query() @@ -165,6 +173,10 @@ func (s *RegistryService) ListPublishers(ctx context.Context, client *ent.Client } func (s *RegistryService) CreatePublisher(ctx context.Context, client *ent.Client, userId string, publisher *drip.Publisher) (*ent.Publisher, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.CreatePublisher") + defer segment.End() + } publisherValid := mapper.ValidatePublisher(publisher) if publisherValid != nil { return nil, fmt.Errorf("invalid publisher: %w", publisherValid) @@ -192,6 +204,10 @@ func (s *RegistryService) CreatePublisher(ctx context.Context, client *ent.Clien } func (s *RegistryService) UpdatePublisher(ctx context.Context, client *ent.Client, update *ent.PublisherUpdateOne) (*ent.Publisher, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.UpdatePublisher") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("updating publisher fields: %v", update.Mutation().Fields()) publisher, err := update.Save(ctx) log.Ctx(ctx).Info().Msgf("success: updated publisher: %v", publisher) @@ -203,6 +219,10 @@ func (s *RegistryService) UpdatePublisher(ctx context.Context, client *ent.Clien } func (s *RegistryService) GetPublisher(ctx context.Context, client *ent.Client, publisherID string) (*ent.Publisher, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.GetPublisher") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("getting publisher: %v", publisherID) publisher, err := client.Publisher. Query(). @@ -218,6 +238,10 @@ func (s *RegistryService) GetPublisher(ctx context.Context, client *ent.Client, } func (s *RegistryService) CreatePersonalAccessToken(ctx context.Context, client *ent.Client, publisherID, name, description string) (*ent.PersonalAccessToken, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.CreatePersonalAccessToken") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("creating personal access token for publisher: %v", publisherID) token := uuid.New().String() pat, err := client.PersonalAccessToken. @@ -235,6 +259,10 @@ func (s *RegistryService) CreatePersonalAccessToken(ctx context.Context, client } func (s *RegistryService) ListPersonalAccessTokens(ctx context.Context, client *ent.Client, publisherID string) ([]*ent.PersonalAccessToken, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.ListPersonalAccessTokens") + defer segment.End() + } pats, err := client.PersonalAccessToken.Query(). Where(personalaccesstoken.PublisherIDEQ(publisherID)). All(ctx) @@ -245,6 +273,10 @@ func (s *RegistryService) ListPersonalAccessTokens(ctx context.Context, client * } func (s *RegistryService) DeletePersonalAccessToken(ctx context.Context, client *ent.Client, tokenID uuid.UUID) error { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.DeletePersonalAccessToken") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("deleting personal access token: %v", tokenID) err := client.PersonalAccessToken. DeleteOneID(tokenID). @@ -256,6 +288,10 @@ func (s *RegistryService) DeletePersonalAccessToken(ctx context.Context, client } func (s *RegistryService) CreateNode(ctx context.Context, client *ent.Client, publisherId string, node *drip.Node) (*ent.Node, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.CreateNode") + defer segment.End() + } validNode := mapper.ValidateNode(node) if validNode != nil { return nil, fmt.Errorf("invalid node: %w", validNode) @@ -285,10 +321,11 @@ func (s *RegistryService) CreateNode(ctx context.Context, client *ent.Client, pu return createdNode, err } -func (s *RegistryService) UpdateNode( - ctx context.Context, - client *ent.Client, - updateFunc func(client *ent.Client) *ent.NodeUpdateOne) (*ent.Node, error) { +func (s *RegistryService) UpdateNode(ctx context.Context, client *ent.Client, updateFunc func(client *ent.Client) *ent.NodeUpdateOne) (*ent.Node, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.UpdateNode") + defer segment.End() + } var n *ent.Node err := db.WithTx(ctx, client, func(tx *ent.Tx) (err error) { update := updateFunc(tx.Client()) @@ -310,8 +347,11 @@ func (s *RegistryService) UpdateNode( } func (s *RegistryService) GetNode(ctx context.Context, client *ent.Client, nodeID string) (*ent.Node, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.GetNode") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.GetNode") + defer segment.End() + } + log.Ctx(ctx).Info().Msgf("getting node: %v", nodeID) node, err := client.Node.Get(ctx, nodeID) if err != nil { @@ -320,11 +360,11 @@ func (s *RegistryService) GetNode(ctx context.Context, client *ent.Client, nodeI return node, nil } -func (s *RegistryService) CreateNodeVersion( - ctx context.Context, - client *ent.Client, - publisherID, nodeID string, - nodeVersion *drip.NodeVersion) (*NodeVersionCreation, error) { +func (s *RegistryService) CreateNodeVersion(ctx context.Context, client *ent.Client, publisherID, nodeID string, nodeVersion *drip.NodeVersion) (*NodeVersionCreation, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.CreateNodeVersion") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("creating node version: %v for nodeId %v", nodeVersion, nodeID) bucketName := "comfy-registry" return db.WithTxResult(ctx, client, func(tx *ent.Tx) (*NodeVersionCreation, error) { @@ -390,10 +430,11 @@ type NodeVersionCreation struct { SignedUrl string } -func (s *RegistryService) ListNodeVersions( - ctx context.Context, client *ent.Client, filter *entity.NodeVersionFilter) (*entity.ListNodeVersionsResult, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.ListNodeVersions") - defer txn.End() +func (s *RegistryService) ListNodeVersions(ctx context.Context, client *ent.Client, filter *entity.NodeVersionFilter) (*entity.ListNodeVersionsResult, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.ListNodeVersions") + defer segment.End() + } query := client.NodeVersion.Query(). WithStorageFile(). Order(ent.Desc(nodeversion.FieldVersion)) @@ -461,6 +502,10 @@ func (s *RegistryService) ListNodeVersions( } func (s *RegistryService) AddNodeReview(ctx context.Context, client *ent.Client, nodeId, userID string, star int) (n *ent.Node, err error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.AddNodeReview") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("add review to node: %v ", nodeId) err = db.WithTx(ctx, client, func(tx *ent.Tx) error { @@ -495,8 +540,10 @@ func (s *RegistryService) AddNodeReview(ctx context.Context, client *ent.Client, } func (s *RegistryService) GetNodeVersionByVersion(ctx context.Context, client *ent.Client, nodeId, nodeVersion string) (*ent.NodeVersion, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.GetNodeVersionByVersion") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.GetNodeVersionByVersion") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("getting node version %v@%v", nodeId, nodeVersion) return client.NodeVersion. Query(). @@ -507,16 +554,20 @@ func (s *RegistryService) GetNodeVersionByVersion(ctx context.Context, client *e } func (s *RegistryService) GetNodeVersion(ctx context.Context, client *ent.Client, nodeVersionId string) (*ent.NodeVersion, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.GetNodeVersion") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.GetNodeVersion") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("getting node version %v", nodeVersionId) return client.NodeVersion. Get(ctx, uuid.MustParse(nodeVersionId)) } func (s *RegistryService) UpdateNodeVersion(ctx context.Context, client *ent.Client, update *ent.NodeVersionUpdateOne) (*ent.NodeVersion, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.UpdateNodeVersion") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.UpdateNodeVersion") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("updating node version fields: %v", update.Mutation().Fields()) return db.WithTxResult(ctx, client, func(tx *ent.Tx) (*ent.NodeVersion, error) { node, err := update.Save(ctx) @@ -534,8 +585,10 @@ func (s *RegistryService) UpdateNodeVersion(ctx context.Context, client *ent.Cli } func (s *RegistryService) RecordNodeInstallation(ctx context.Context, client *ent.Client, node *ent.Node) (*ent.Node, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.RecordNodeInstallation") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.RecordNodeInstallation") + defer segment.End() + } var n *ent.Node err := db.WithTx(ctx, client, func(tx *ent.Tx) (err error) { n, err = tx.Node.UpdateOne(node).AddTotalInstall(1).Save(ctx) @@ -553,8 +606,10 @@ func (s *RegistryService) RecordNodeInstallation(ctx context.Context, client *en } func (s *RegistryService) GetLatestNodeVersion(ctx context.Context, client *ent.Client, nodeId string) (*ent.NodeVersion, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.GetLatestNodeVersion") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.GetLatestNodeVersion") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("Getting latest version of node: %v", nodeId) nodeVersion, err := client.NodeVersion. Query(). @@ -591,6 +646,10 @@ func (s *RegistryService) MarkComfyNodeExtractionFailed( nodeVersion string, info *schema.ComfyNodeCloudBuildInfo, ) error { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.MarkComfyNodeExtractionFailed") + defer segment.End() + } u := client.NodeVersion. Update(). Where( @@ -611,8 +670,10 @@ func (s *RegistryService) CreateComfyNodes( comfyNodes map[string]drip.ComfyNode, info *schema.ComfyNodeCloudBuildInfo, ) error { - txn := s.newRelicApp.StartTransaction("RegistryService.CreateComfyNodes") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.CreateComfyNodes") + defer segment.End() + } return db.WithTx(ctx, client, func(tx *ent.Tx) error { // Query the NodeVersion with the given nodeID and nodeVersion, lock it for updates nv, err := tx.NodeVersion.Query(). @@ -699,9 +760,11 @@ func (s *RegistryService) GetComfyNode( client *ent.Client, nodeID, nodeVersion, comfyNodeName string, ) (*ent.ComfyNode, error) { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.GetComfyNode") + defer segment.End() + } // Query the NodeVersion with the given nodeID and nodeVersion, ensuring extraction status is success - txn := s.newRelicApp.StartTransaction("RegistryService.GetComfyNode") - defer txn.End() nv, err := client.NodeVersion.Query(). Where(nodeversion.VersionEQ(nodeVersion)). Where(nodeversion.NodeIDEQ(nodeID)). @@ -726,9 +789,11 @@ func (s *RegistryService) GetComfyNode( func (s *RegistryService) TriggerComfyNodesBackfill( ctx context.Context, client *ent.Client, max *int) error { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.TriggerComfyNodesBackfill") + defer segment.End() + } // Query all NodeVersions with pending comfy node extraction status - txn := s.newRelicApp.StartTransaction("RegistryService.TriggerComfyNodesBackfill") - defer txn.End() q := client.NodeVersion. Query(). WithStorageFile(). @@ -770,8 +835,10 @@ func (s *RegistryService) AssertPublisherPermissions(ctx context.Context, userID string, permissions []schema.PublisherPermissionType, ) (err error) { - txn := s.newRelicApp.StartTransaction("RegistryService.AssertPublisherPermissions") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.AssertPublisherPermissions") + defer segment.End() + } w, err := client.Publisher.Get(ctx, publisherID) if err != nil { return fmt.Errorf("fail to query publisher by id: %s %w", publisherID, err) @@ -796,8 +863,10 @@ func (s *RegistryService) IsPersonalAccessTokenValidForPublisher(ctx context.Con publisherID string, accessToken string, ) (bool, error) { - txn := s.newRelicApp.StartTransaction("RegistryService.IsPersonalAccessTokenValidForPublisher") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.IsPersonalAccessTokenValidForPublisher") + defer segment.End() + } w, err := client.Publisher.Get(ctx, publisherID) if err != nil { log.Ctx(ctx).Error().Err(err).Msgf("fail to find publisher by id: %s", publisherID) @@ -818,8 +887,10 @@ func (s *RegistryService) IsPersonalAccessTokenValidForPublisher(ctx context.Con } func (s *RegistryService) AssertNodeBelongsToPublisher(ctx context.Context, client *ent.Client, publisherID string, nodeID string) error { - txn := s.newRelicApp.StartTransaction("RegistryService.AssertNodeBelongsToPublisher") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.AssertNodeBelongsToPublisher") + defer segment.End() + } node, err := client.Node.Get(ctx, nodeID) if err != nil { return fmt.Errorf("failed to get node: %w", err) @@ -831,8 +902,10 @@ func (s *RegistryService) AssertNodeBelongsToPublisher(ctx context.Context, clie } func (s *RegistryService) AssertAccessTokenBelongsToPublisher(ctx context.Context, client *ent.Client, publisherID string, tokenId uuid.UUID) error { - txn := s.newRelicApp.StartTransaction("RegistryService.AssertAccessTokenBelongsToPublisher") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.AssertAccessTokenBelongsToPublisher") + defer segment.End() + } pat, err := client.PersonalAccessToken.Query().Where( personalaccesstoken.IDEQ(tokenId), personalaccesstoken.PublisherIDEQ(publisherID), @@ -847,8 +920,10 @@ func (s *RegistryService) AssertAccessTokenBelongsToPublisher(ctx context.Contex } func (s *RegistryService) DeletePublisher(ctx context.Context, client *ent.Client, publisherID string) error { - txn := s.newRelicApp.StartTransaction("RegistryService.DeletePublisher") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.DeletePublisher") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("deleting publisher: %v", publisherID) return db.WithTx(ctx, client, func(tx *ent.Tx) error { client = tx.Client() @@ -882,8 +957,10 @@ func (s *RegistryService) DeletePublisher(ctx context.Context, client *ent.Clien } func (s *RegistryService) DeleteNode(ctx context.Context, client *ent.Client, nodeID string) error { - txn := s.newRelicApp.StartTransaction("RegistryService.DeleteNode") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.DeleteNode") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("deleting node: %v", nodeID) db.WithTx(ctx, client, func(tx *ent.Tx) error { nv, err := tx.Client().NodeVersion.Query().Where(nodeversion.NodeID(nodeID)).All(ctx) @@ -910,8 +987,10 @@ func (s *RegistryService) DeleteNode(ctx context.Context, client *ent.Client, no } func (s *RegistryService) DeleteNodeVersion(ctx context.Context, client *ent.Client, nodeIDVersion string) error { - txn := s.newRelicApp.StartTransaction("RegistryService.DeleteNodeVersion") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.DeleteNodeVersion") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("deleting node version: %v", nodeIDVersion) db.WithTx(ctx, client, func(tx *ent.Tx) error { nv, err := tx.Client().NodeVersion.Get(ctx, uuid.MustParse(nodeIDVersion)) @@ -957,8 +1036,10 @@ func IsPermissionError(err error) bool { } func (s *RegistryService) BanPublisher(ctx context.Context, client *ent.Client, id string) error { - txn := s.newRelicApp.StartTransaction("RegistryService.BanPublisher") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.BanPublisher") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("banning publisher: %v", id) pub, err := client.Publisher.Get(ctx, id) if err != nil { @@ -1010,8 +1091,10 @@ func (s *RegistryService) BanPublisher(ctx context.Context, client *ent.Client, } func (s *RegistryService) BanNode(ctx context.Context, client *ent.Client, publisherid, id string) error { - txn := s.newRelicApp.StartTransaction("RegistryService.BanNode") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.BanNode") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("banning publisher node: %v %v", publisherid, id) return db.WithTx(ctx, client, func(tx *ent.Tx) error { @@ -1047,8 +1130,10 @@ func (s *RegistryService) BanNode(ctx context.Context, client *ent.Client, publi } func (s *RegistryService) AssertNodeBanned(ctx context.Context, client *ent.Client, nodeID string) error { - txn := s.newRelicApp.StartTransaction("RegistryService.AssertNodeBanned") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.AssertNodeBanned") + defer segment.End() + } node, err := client.Node.Get(ctx, nodeID) if ent.IsNotFound(err) { return nil @@ -1063,6 +1148,10 @@ func (s *RegistryService) AssertNodeBanned(ctx context.Context, client *ent.Clie } func (s *RegistryService) AssertPublisherBanned(ctx context.Context, client *ent.Client, publisherID string) error { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.AssertPublisherBanned") + defer segment.End() + } publisher, err := client.Publisher.Get(ctx, publisherID) if ent.IsNotFound(err) { return nil @@ -1077,8 +1166,10 @@ func (s *RegistryService) AssertPublisherBanned(ctx context.Context, client *ent } func (s *RegistryService) ReindexAllNodes(ctx context.Context, client *ent.Client) error { - txn := s.newRelicApp.StartTransaction("RegistryService.ReindexAllNodes") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.ReindexAllNodes") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("reindexing nodes") nodes, err := s.decorateNodeQueryWithLatestVersion(client.Node.Query()).All(ctx) if err != nil { @@ -1108,8 +1199,10 @@ func (s *RegistryService) ReindexAllNodes(ctx context.Context, client *ent.Clien var reindexLock = sync.Mutex{} func (s *RegistryService) ReindexAllNodesBackground(ctx context.Context, client *ent.Client) (err error) { - txn := s.newRelicApp.StartTransaction("RegistryService.ReindexAllNodesBackground") - defer txn.End() + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.ReindexAllNodesBackground") + defer segment.End() + } if !reindexLock.TryLock() { return fmt.Errorf("another reindex is in progress") } @@ -1160,6 +1253,10 @@ func (s *RegistryService) decorateNodeQueryWithLatestVersion(q *ent.NodeQuery) * func (s *RegistryService) PerformSecurityCheck( ctx context.Context, client *ent.Client, nodeVersion *ent.NodeVersion) error { + if txn := newrelic.FromContext(ctx); txn != nil { + segment := txn.StartSegment("RegistryService.PerformSecurityCheck") + defer segment.End() + } log.Ctx(ctx).Info().Msgf("Scanning node %s@%s w/ version ID: %s", nodeVersion.NodeID, nodeVersion.Version, nodeVersion.ID) @@ -1171,10 +1268,10 @@ func (s *RegistryService) PerformSecurityCheck( if err != nil { if strings.Contains(err.Error(), "404") { log.Ctx(ctx).Info().Msgf( - "Node zip file doesn’t exist %s@%s. Updating to deleted.", nodeVersion.NodeID, nodeVersion.Version) + "Node zip file doesn't exist %s@%s. Updating to deleted.", nodeVersion.NodeID, nodeVersion.Version) err := nodeVersion.Update(). SetStatus(schema.NodeVersionStatusDeleted). - SetStatusReason("Node zip file doesn’t exist").Exec(ctx) + SetStatusReason("Node zip file doesn't exist").Exec(ctx) if err != nil { log.Ctx(ctx).Error().Err(err).Msgf("failed to update node version status to active") } @@ -1183,7 +1280,7 @@ func (s *RegistryService) PerformSecurityCheck( } if issues == "" { - log.Ctx(ctx).Info().Msgf(""+ + log.Ctx(ctx).Info().Msgf( "No security issues found in node %s@%s. Updating to active.", nodeVersion.NodeID, nodeVersion.Version) err := nodeVersion.Update(). SetStatus(schema.NodeVersionStatusActive).