Document extraction for control + subcontrol refs should be async (#2306)

Author: adelowo

.task/checksum/generate-ent-smart

@@ -1 +1 @@
-2205155eb89e6bff7f80240b5dc043
+5c4fd560f5af11c2e8c76c1dfa910647

.task/checksum/generate-openapi-smart

@@ -1 +1 @@
-76e1ae53815c369bd25a0d735abf7690
+b7aee90061ab56e45ce54858de8b8385

internal/ent/checksum/.history_schema_checksum

@@ -1 +1 @@
-71434179a5ad10caed7fabf528bcdcc9bec5ce03b425b66c66b2b1034c654a9e
+f227db62ffc3fc3f9de100b8e18ac45a2f0aa1fb53a1d1a5524c047375d10838

internal/ent/checksum/.schema_checksum

@@ -1 +1 @@
-d6ab508c9091731765993d63a4644d543db1c74f2a41c0eea915ac0b4c59f5c9
+a61c0d6804b980ca0ae94ae836db20e3e28dd791bc1f419d0df6ca0f3e96d2cd

internal/ent/hooks/documentextract.go

@@ -7,29 +7,21 @@ import (
 	"strings"
 
 	"entgo.io/ent"
+
+	"github.com/theopenlane/core/internal/ent/generated"
 	"github.com/theopenlane/core/internal/ent/generated/control"
 	"github.com/theopenlane/core/internal/ent/generated/hook"
 	"github.com/theopenlane/core/pkg/logx"
 )
 
-// addControlsMutation allows for adding to the controls edge in a mutation
-type addControlsMutation interface {
-	AddControlIDs(ids ...string)
-}
-
-// addSubcontrolsMutation allows for adding to the subcontrols edge in a mutation
-type addSubcontrolsMutation interface {
-	AddSubcontrolIDs(ids ...string)
-}
-
 // versionMutation allows for setting the revision field in a mutation
 type versionMutation interface {
 	SetRevision(revision string)
 }
 
-// HookParseAssociations is an ent hook that parses associations from a document
-// such as referenced controls and adds the necessary edges
-func HookParseAssociations() ent.Hook {
+// HookDetailsVersion is an ent hook that parses the versions from the details of a document
+// creation
+func HookDetailsVersion() ent.Hook {
 	return hook.If(func(next ent.Mutator) ent.Mutator {
 		return ent.MutateFunc(func(ctx context.Context, m ent.Mutation) (ent.Value, error) {
 			mut := m.(detailsMutation)
@@ -39,30 +31,10 @@ func HookParseAssociations() ent.Hook {
 				return next.Mutate(ctx, m)
 			}
 
-			edgeLinks := getDocumentAssociations(ctx, mut)
-
-			if edgeLinks == nil {
-				return next.Mutate(ctx, m)
-			}
-
-			if len(edgeLinks.controlIDs) > 0 {
-				conMut, ok := mut.(addControlsMutation)
-				if ok {
-					conMut.AddControlIDs(edgeLinks.controlIDs...)
-				}
-			}
-
-			if len(edgeLinks.subcontrolIDs) > 0 {
-				subconMut, ok := mut.(addSubcontrolsMutation)
-				if ok {
-					subconMut.AddSubcontrolIDs(edgeLinks.subcontrolIDs...)
-				}
-			}
-
-			if edgeLinks.version != "" {
+			if version := findVersion(details); version != "" {
 				verMut, ok := mut.(versionMutation)
 				if ok {
-					verMut.SetRevision(edgeLinks.version)
+					verMut.SetRevision(version)
 				}
 			}
 
@@ -74,31 +46,10 @@ func HookParseAssociations() ent.Hook {
 
 // edgeLinks is a struct that holds the IDs of associated entities that should be linked to the document being created or updated, such as controlIDs
 type edgeLinks struct {
-	version       string
 	controlIDs    []string
 	subcontrolIDs []string
 }
 
-// getDocumentAssociations will read text details and try to extract any associations to other entities in the system, such as referenced control IDs, asset IDs, and identity holder IDs. This is a placeholder implementation and should be replaced with actual parsing logic based on the expected format of the details text.
-func getDocumentAssociations(ctx context.Context, m detailsMutation) *edgeLinks {
-	orgControls := getOrganizationControls(ctx, m)
-
-	if orgControls == nil {
-		return nil
-	}
-
-	details, ok := m.Details()
-	if !ok || details == "" {
-		return nil
-	}
-
-	edgeLinks := findControlMatches(details, orgControls)
-
-	edgeLinks.version = findVersion(details)
-
-	return edgeLinks
-}
-
 // findVersion attempts to find the version of the document in the details by
 // looking for a line that starts with "Version:" and extracting the version number from it. It returns the version in semver format if found, or an empty string if not found.
 func findVersion(details string) string {
@@ -178,9 +129,9 @@ type controlInfo struct {
 // controlMapping contains the refCode with additional info
 type controlMapping map[string]controlInfo
 
-func getOrganizationControls(ctx context.Context, m detailsMutation) controlMapping {
+func getOrganizationControlsFromClient(ctx context.Context, client *generated.Client) controlMapping {
 	result := controlMapping{}
-	controls, err := m.Client().Control.Query().Where(
+	controls, err := client.Control.Query().Where(
 		control.IsTrustCenterControl(false),
 		control.SystemOwned(false),
 	).WithSubcontrols().All(ctx)

internal/ent/hooks/listeners_document_associations.go

@@ -0,0 +1,149 @@
+package hooks
+
+import (
+	"context"
+
+	"entgo.io/ent"
+	"github.com/samber/lo"
+
+	"github.com/theopenlane/core/internal/ent/eventqueue"
+	"github.com/theopenlane/core/internal/ent/generated"
+	"github.com/theopenlane/core/internal/workflows"
+	"github.com/theopenlane/core/pkg/gala"
+)
+
+// RegisterGalaDocumentAssociationListeners registers listeners that link
+// referenced controls to documents asynchronously after document creation.
+func RegisterGalaDocumentAssociationListeners(registry *gala.Registry) ([]gala.ListenerID, error) {
+	return gala.RegisterListeners(registry,
+		documentAssociationDefinition(generated.TypeActionPlan),
+		documentAssociationDefinition(generated.TypeInternalPolicy),
+		documentAssociationDefinition(generated.TypeProcedure),
+	)
+}
+
+func documentAssociationDefinition(schemaType string) gala.Definition[eventqueue.MutationGalaPayload] {
+	topic := eventqueue.MutationTopic(eventqueue.MutationConcernDirect, schemaType)
+
+	return gala.Definition[eventqueue.MutationGalaPayload]{
+		Topic:      topic,
+		Name:       "document.associations." + schemaType,
+		Operations: []string{ent.OpCreate.String()},
+		Handle:     handleDocumentAssociationCreated,
+	}
+}
+
+func handleDocumentAssociationCreated(ctx gala.HandlerContext, payload eventqueue.MutationGalaPayload) error {
+	ctx, client, ok := eventqueue.ClientFromHandler(ctx)
+	if !ok {
+		return nil
+	}
+
+	documentID, ok := eventqueue.MutationEntityID(payload, ctx.Envelope.Headers.Properties)
+	if !ok || documentID == "" {
+		return nil
+	}
+
+	switch payload.MutationType {
+	case generated.TypeActionPlan:
+		return parseActionPlanAssociations(ctx.Context, client, documentID)
+	case generated.TypeInternalPolicy:
+		return parseInternalPolicyAssociations(ctx.Context, client, documentID)
+	case generated.TypeProcedure:
+		return parseProcedureAssociations(ctx.Context, client, documentID)
+	default:
+		return nil
+	}
+}
+
+func parseActionPlanAssociations(ctx context.Context, client *generated.Client, documentID string) error {
+	doc, err := client.ActionPlan.Get(ctx, documentID)
+	if err != nil {
+		if generated.IsNotFound(err) {
+			return nil
+		}
+
+		return err
+	}
+
+	links := getDocumentAssociationsForDetails(ctx, client, doc.Details)
+	if links == nil || len(links.controlIDs) == 0 {
+		return nil
+	}
+
+	return client.ActionPlan.UpdateOneID(doc.ID).
+		SetRevision(doc.Revision).
+		AddControlIDs(lo.Uniq(links.controlIDs)...).
+		Exec(workflows.AllowContext(ctx))
+}
+
+func parseInternalPolicyAssociations(ctx context.Context, client *generated.Client, documentID string) error {
+	doc, err := client.InternalPolicy.Get(ctx, documentID)
+	if err != nil {
+		if generated.IsNotFound(err) {
+			return nil
+		}
+
+		return err
+	}
+
+	links := getDocumentAssociationsForDetails(ctx, client, doc.Details)
+	if links == nil || !links.hasAssociations() {
+		return nil
+	}
+
+	update := client.InternalPolicy.UpdateOneID(doc.ID).SetRevision(doc.Revision)
+	if len(links.controlIDs) > 0 {
+		update.AddControlIDs(lo.Uniq(links.controlIDs)...)
+	}
+
+	if len(links.subcontrolIDs) > 0 {
+		update.AddSubcontrolIDs(lo.Uniq(links.subcontrolIDs)...)
+	}
+
+	return update.Exec(workflows.AllowContext(ctx))
+}
+
+func parseProcedureAssociations(ctx context.Context, client *generated.Client, documentID string) error {
+	doc, err := client.Procedure.Get(ctx, documentID)
+	if err != nil {
+		if generated.IsNotFound(err) {
+			return nil
+		}
+
+		return err
+	}
+
+	links := getDocumentAssociationsForDetails(ctx, client, doc.Details)
+	if links == nil || !links.hasAssociations() {
+		return nil
+	}
+
+	update := client.Procedure.UpdateOneID(doc.ID).SetRevision(doc.Revision)
+	if len(links.controlIDs) > 0 {
+		update.AddControlIDs(lo.Uniq(links.controlIDs)...)
+	}
+
+	if len(links.subcontrolIDs) > 0 {
+		update.AddSubcontrolIDs(lo.Uniq(links.subcontrolIDs)...)
+	}
+
+	return update.Exec(workflows.AllowContext(ctx))
+}
+
+func getDocumentAssociationsForDetails(ctx context.Context, client *generated.Client, details string) *edgeLinks {
+	if details == "" {
+		return nil
+	}
+
+	orgControls := getOrganizationControlsFromClient(ctx, client)
+	if orgControls == nil {
+		return nil
+	}
+
+	return findControlMatches(details, orgControls)
+}
+
+func (e *edgeLinks) hasAssociations() bool {
+	return e != nil && (len(e.controlIDs) > 0 || len(e.subcontrolIDs) > 0)
+}

internal/ent/hooks/listeners_gala_registration_test.go

@@ -110,6 +110,22 @@ func TestRegisterGalaIdentityResolutionListeners(t *testing.T) {
 	require.False(t, registry.InterestedIn(topic, ent.OpDelete.String()))
 }
 
+func TestRegisterGalaDocumentAssociationListeners(t *testing.T) {
+	t.Parallel()
+
+	registry := gala.NewRegistry()
+
+	ids, err := RegisterGalaDocumentAssociationListeners(registry)
+	require.NoError(t, err)
+	require.Len(t, ids, 3)
+
+	for _, schemaType := range []string{entgen.TypeActionPlan, entgen.TypeInternalPolicy, entgen.TypeProcedure} {
+		topic := eventqueue.MutationTopicName(eventqueue.MutationConcernDirect, schemaType)
+		require.True(t, registry.InterestedIn(topic, ent.OpCreate.String()))
+		require.False(t, registry.InterestedIn(topic, ent.OpUpdate.String()))
+	}
+}
+
 func TestRegisterGalaNotificationListeners(t *testing.T) {
 	t.Parallel()
 

internal/ent/schema/mixin_document.go

@@ -62,7 +62,7 @@ func (d DocumentMixin) Hooks() []ent.Hook {
 			ent.OpCreate|ent.OpUpdateOne|ent.OpUpdateOne,
 		),
 		hooks.HookImportDocument(),
-		hooks.HookParseAssociations(),
+		hooks.HookDetailsVersion(),
 		hooks.HookSummarizeDetails(),
 		hook.On(
 			hooks.HookStatusApproval(),

internal/graphapi/checksum/.history_schema_checksum

@@ -1 +1 @@
-9107d46ce0a234bc6a827510c2e4ceeb7d45e80db73d6ea3839161a6785c927a
+3f04e7b3b3e5d92c7741accff77e4daefc04d658611b2518517cde2641931b7c

internal/httpserve/serveropts/gala.go

@@ -91,6 +91,7 @@ func ConfigureGala(ctx context.Context, galaApp, notificationGala *gala.Gala, db
 		{galaApp, hooks.RegisterGalaSlackListeners},
 		{galaApp, hooks.RegisterGalaVendorScoringListeners},
 		{galaApp, hooks.RegisterGalaIdentityResolutionListeners},
+		{galaApp, hooks.RegisterGalaDocumentAssociationListeners},
 		{notificationGala, hooks.RegisterGalaNotificationListeners},
 	}