// Copyright (c) 2025-2026 Probo Inc . // // Permission is hereby granted, free of charge, to any person obtaining a copy // of this software and associated documentation files (the "Software"), to deal // in the Software without restriction, including without limitation the rights // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell // copies of the Software, and to permit persons to whom the Software is // furnished to do so, subject to the following conditions: // // The above copyright notice and this permission notice shall be included in // all copies or substantial portions of the Software. // // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE // SOFTWARE. package coredata import ( "context" "errors" "fmt" "maps" "time" "github.com/jackc/pgx/v5" "go.gearno.de/kit/pg" "go.probo.inc/probo/pkg/gid" "go.probo.inc/probo/pkg/iam/policy" "go.probo.inc/probo/pkg/page" ) func (r Risk) GetGeneratedDocumentID( ctx context.Context, conn pg.Querier, organizationID gid.GID, ) (*gid.GID, error) { var documentID *gid.GID err := conn.QueryRow( ctx, ` SELECT risks_document_id FROM generated_documents WHERE organization_id = @organization_id `, pgx.NamedArgs{"organization_id": organizationID}, ).Scan(&documentID) if errors.Is(err, pgx.ErrNoRows) { return nil, nil } if err != nil { return nil, fmt.Errorf("cannot get risk list document ID: %w", err) } return documentID, nil } func (r Risk) UpsertGeneratedDocumentID( ctx context.Context, conn pg.Tx, organizationID gid.GID, tenantID gid.TenantID, documentID gid.GID, ) error { now := time.Now() _, err := conn.Exec( ctx, ` INSERT INTO generated_documents ( organization_id, tenant_id, risks_document_id, created_at, updated_at ) VALUES ( @organization_id, @tenant_id, @risks_document_id, @created_at, @updated_at ) ON CONFLICT (organization_id) DO UPDATE SET risks_document_id = @risks_document_id, updated_at = @updated_at `, pgx.NamedArgs{ "organization_id": organizationID, "tenant_id": tenantID, "risks_document_id": documentID, "created_at": now, "updated_at": now, }, ) if err != nil { return fmt.Errorf("cannot upsert risk list document ID: %w", err) } return nil } func (r Risk) ClearGeneratedDocumentID( ctx context.Context, conn pg.Tx, documentIDs []gid.GID, ) error { ids := make([]string, len(documentIDs)) for i, id := range documentIDs { ids[i] = id.String() } _, err := conn.Exec( ctx, ` UPDATE generated_documents SET risks_document_id = NULL, updated_at = @now WHERE risks_document_id = ANY(@ids) `, pgx.NamedArgs{ "ids": ids, "now": time.Now(), }, ) if err != nil { return fmt.Errorf("cannot clear risk list document references: %w", err) } return nil } type ( Risk struct { ID gid.GID `db:"id"` OrganizationID gid.GID `db:"organization_id"` Name string `db:"name"` Description *string `db:"description"` Category string `db:"category"` Treatment RiskTreatment `db:"treatment"` Note string `db:"note"` OwnerID *gid.GID `db:"owner_profile_id"` InherentLikelihood int `db:"inherent_likelihood"` InherentImpact int `db:"inherent_impact"` InherentRiskScore int `db:"inherent_risk_score"` ResidualLikelihood int `db:"residual_likelihood"` ResidualImpact int `db:"residual_impact"` ResidualRiskScore int `db:"residual_risk_score"` CreatedAt time.Time `db:"created_at"` UpdatedAt time.Time `db:"updated_at"` // Ordering only OwnerFullName *string `db:"owner_full_name"` } Risks []*Risk ) func (r *Risk) CursorKey(orderBy RiskOrderField) page.CursorKey { switch orderBy { case RiskOrderFieldCreatedAt: return page.CursorKey{ID: r.ID, Value: r.CreatedAt} case RiskOrderFieldName: return page.CursorKey{ID: r.ID, Value: r.Name} case RiskOrderFieldCategory: return page.CursorKey{ID: r.ID, Value: r.Category} case RiskOrderFieldTreatment: return page.CursorKey{ID: r.ID, Value: r.Treatment} case RiskOrderFieldInherentRiskScore: return page.CursorKey{ID: r.ID, Value: r.InherentRiskScore} case RiskOrderFieldResidualRiskScore: return page.CursorKey{ID: r.ID, Value: r.ResidualRiskScore} case RiskOrderFieldOwnerFullName: return page.CursorKey{ID: r.ID, Value: r.OwnerFullName} } panic(fmt.Sprintf("unsupported order by: %s", orderBy)) } func (r *Risk) AuthorizationAttributes( ctx context.Context, conn pg.Querier, resourceIDs []gid.GID, ) (policy.AttributesByID, error) { q := `SELECT id, organization_id FROM risks WHERE id = ANY(@resource_ids::text[])` args := pgx.StrictNamedArgs{ "resource_ids": resourceIDs, } rows, err := conn.Query(ctx, q, args) if err != nil { return nil, fmt.Errorf("cannot query authorization attributes: %w", err) } defer rows.Close() attrsByID := make(policy.AttributesByID) for rows.Next() { var id, organizationID gid.GID if err := rows.Scan(&id, &organizationID); err != nil { return nil, fmt.Errorf("cannot scan authorization attributes: %w", err) } attrsByID[id] = policy.Attributes{ "organization_id": organizationID.String(), } } if err := rows.Err(); err != nil { return nil, fmt.Errorf("cannot iterate authorization attributes: %w", err) } return attrsByID, nil } func (r *Risks) CountByMeasureID( ctx context.Context, conn pg.Querier, scope Scoper, measureID gid.GID, filter *RiskFilter, ) (int, error) { q := ` WITH rsks AS ( SELECT r.id, r.tenant_id, r.search_vector FROM risks r INNER JOIN risks_measures rm ON r.id = rm.risk_id WHERE rm.measure_id = @measure_id ) SELECT COUNT(id) FROM rsks WHERE %s AND %s ` q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment()) args := pgx.NamedArgs{"measure_id": measureID} maps.Copy(args, scope.SQLArguments()) maps.Copy(args, filter.SQLArguments()) row := conn.QueryRow(ctx, q, args) var count int if err := row.Scan(&count); err != nil { return 0, fmt.Errorf("cannot scan count: %w", err) } return count, nil } func (r *Risks) LoadByMeasureID( ctx context.Context, conn pg.Querier, scope Scoper, measureID gid.GID, cursor *page.Cursor[RiskOrderField], filter *RiskFilter, ) error { q := ` WITH rsks AS ( SELECT r.id, r.tenant_id, r.organization_id, r.name, r.description, r.category, r.owner_profile_id, p.full_name as owner_full_name, r.treatment, r.note, r.inherent_likelihood, r.inherent_impact, r.inherent_risk_score, r.residual_likelihood, r.residual_impact, r.residual_risk_score, r.search_vector, r.created_at, r.updated_at FROM risks r INNER JOIN risks_measures rm ON r.id = rm.risk_id LEFT JOIN iam_membership_profiles p ON r.owner_profile_id = p.id WHERE rm.measure_id = @measure_id ) SELECT id, organization_id, name, description, category, owner_profile_id, owner_full_name, treatment, note, inherent_likelihood, inherent_impact, inherent_risk_score, residual_likelihood, residual_impact, residual_risk_score, created_at, updated_at FROM rsks WHERE %s AND %s AND %s ` q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment(), cursor.SQLFragment()) args := pgx.NamedArgs{"measure_id": measureID} maps.Copy(args, scope.SQLArguments()) maps.Copy(args, filter.SQLArguments()) maps.Copy(args, cursor.SQLArguments()) rows, err := conn.Query(ctx, q, args) if err != nil { return fmt.Errorf("cannot query risks: %w", err) } risks, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Risk]) if err != nil { return fmt.Errorf("cannot collect risks: %w", err) } *r = risks return nil } func (r *Risks) CountByOrganizationID( ctx context.Context, conn pg.Querier, scope Scoper, organizationID gid.GID, filter *RiskFilter, ) (int, error) { q := ` SELECT COUNT(id) FROM risks WHERE %s AND organization_id = @organization_id AND %s ` q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment()) args := pgx.NamedArgs{"organization_id": organizationID} maps.Copy(args, scope.SQLArguments()) maps.Copy(args, filter.SQLArguments()) row := conn.QueryRow(ctx, q, args) var count int if err := row.Scan(&count); err != nil { return 0, fmt.Errorf("cannot scan count: %w", err) } return count, nil } func (r *Risks) LoadByOrganizationID( ctx context.Context, conn pg.Querier, scope Scoper, organizationID gid.GID, cursor *page.Cursor[RiskOrderField], filter *RiskFilter, ) error { q := ` WITH rsks AS ( SELECT r.id, r.tenant_id, r.organization_id, r.name, r.description, r.owner_profile_id, p.full_name as owner_full_name, r.treatment, r.note, r.inherent_likelihood, r.inherent_impact, r.inherent_risk_score, r.residual_likelihood, r.residual_impact, r.residual_risk_score, r.category, r.search_vector, r.created_at, r.updated_at FROM risks r LEFT JOIN iam_membership_profiles p ON r.owner_profile_id = p.id WHERE r.organization_id = @organization_id ) SELECT id, organization_id, name, description, owner_profile_id, owner_full_name, treatment, note, inherent_likelihood, inherent_impact, inherent_risk_score, residual_likelihood, residual_impact, residual_risk_score, category, created_at, updated_at FROM rsks WHERE %s AND %s AND %s ` q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment(), cursor.SQLFragment()) args := pgx.NamedArgs{"organization_id": organizationID} maps.Copy(args, scope.SQLArguments()) maps.Copy(args, filter.SQLArguments()) maps.Copy(args, cursor.SQLArguments()) rows, err := conn.Query(ctx, q, args) if err != nil { return fmt.Errorf("cannot query risks: %w", err) } risks, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Risk]) if err != nil { return fmt.Errorf("cannot collect risks: %w", err) } *r = risks return nil } func (r *Risk) LoadByID( ctx context.Context, conn pg.Querier, scope Scoper, riskID gid.GID, ) error { q := ` SELECT id, organization_id, name, description, category, owner_profile_id, NULL as owner_full_name, treatment, note, inherent_likelihood, inherent_impact, inherent_risk_score, residual_likelihood, residual_impact, residual_risk_score, created_at, updated_at FROM risks WHERE %s AND id = @id LIMIT 1; ` q = fmt.Sprintf(q, scope.SQLFragment()) args := pgx.StrictNamedArgs{"id": riskID} maps.Copy(args, scope.SQLArguments()) rows, err := conn.Query(ctx, q, args) if err != nil { return fmt.Errorf("cannot query risk: %w", err) } risk, err := pgx.CollectExactlyOneRow(rows, pgx.RowToStructByName[Risk]) if err != nil { if errors.Is(err, pgx.ErrNoRows) { return ErrResourceNotFound } return fmt.Errorf("cannot collect risk: %w", err) } *r = risk return nil } func (r *Risks) LoadByIDs( ctx context.Context, conn pg.Querier, scope Scoper, riskIDs []gid.GID, ) error { q := ` SELECT id, organization_id, name, description, category, owner_profile_id, NULL as owner_full_name, treatment, note, inherent_likelihood, inherent_impact, inherent_risk_score, residual_likelihood, residual_impact, residual_risk_score, created_at, updated_at FROM risks WHERE %s AND id = ANY(@risk_ids) ` q = fmt.Sprintf(q, scope.SQLFragment()) args := pgx.StrictNamedArgs{"risk_ids": riskIDs} maps.Copy(args, scope.SQLArguments()) rows, err := conn.Query(ctx, q, args) if err != nil { return fmt.Errorf("cannot query risks: %w", err) } risks, err := pgx.CollectRows(rows, pgx.RowToAddrOfStructByName[Risk]) if err != nil { return fmt.Errorf("cannot collect risks: %w", err) } *r = risks if len(risks) != len(gid.NewSet(riskIDs...)) { return ErrResourceNotFound } return nil } func (r *Risk) Insert( ctx context.Context, conn pg.Tx, scope Scoper, ) error { q := ` INSERT INTO risks (id, tenant_id, organization_id, name, description, category, owner_profile_id, treatment, note, inherent_likelihood, inherent_impact, residual_likelihood, residual_impact, created_at, updated_at) VALUES (@id, @tenant_id, @organization_id, @name, @description, @category, @owner_profile_id, @treatment, @note, @inherent_likelihood, @inherent_impact, @residual_likelihood, @residual_impact, @created_at, @updated_at) ` args := pgx.StrictNamedArgs{ "id": r.ID, "tenant_id": scope.GetTenantID(), "organization_id": r.OrganizationID, "name": r.Name, "description": r.Description, "category": r.Category, "owner_profile_id": r.OwnerID, "treatment": r.Treatment, "note": r.Note, "inherent_likelihood": r.InherentLikelihood, "inherent_impact": r.InherentImpact, "residual_likelihood": r.ResidualLikelihood, "residual_impact": r.ResidualImpact, "created_at": r.CreatedAt, "updated_at": r.UpdatedAt, } _, err := conn.Exec(ctx, q, args) return err } func (r *Risk) Update( ctx context.Context, conn pg.Tx, scope Scoper, ) error { q := ` UPDATE risks SET name = @name, description = @description, owner_profile_id = @owner_profile_id, treatment = @treatment, inherent_likelihood = @inherent_likelihood, inherent_impact = @inherent_impact, residual_likelihood = @residual_likelihood, residual_impact = @residual_impact, category = @category, note = @note, updated_at = @updated_at WHERE %s AND id = @risk_id RETURNING inherent_risk_score, residual_risk_score ` q = fmt.Sprintf(q, scope.SQLFragment()) args := pgx.StrictNamedArgs{ "risk_id": r.ID, "name": r.Name, "description": r.Description, "category": r.Category, "owner_profile_id": r.OwnerID, "treatment": r.Treatment, "note": r.Note, "inherent_likelihood": r.InherentLikelihood, "inherent_impact": r.InherentImpact, "residual_likelihood": r.ResidualLikelihood, "residual_impact": r.ResidualImpact, "updated_at": r.UpdatedAt, } maps.Copy(args, scope.SQLArguments()) err := conn.QueryRow(ctx, q, args).Scan( &r.InherentRiskScore, &r.ResidualRiskScore, ) if err != nil { return fmt.Errorf("cannot update risk: %w", err) } return nil } func (r *Risk) Delete( ctx context.Context, conn pg.Tx, scope Scoper, riskID gid.GID, ) error { q := ` DELETE FROM risks WHERE %s AND id = @id ` q = fmt.Sprintf(q, scope.SQLFragment()) args := pgx.StrictNamedArgs{"id": riskID} maps.Copy(args, scope.SQLArguments()) _, err := conn.Exec(ctx, q, args) return err } func (r *Risks) CountByDocumentID( ctx context.Context, conn pg.Querier, scope Scoper, documentID gid.GID, filter *RiskFilter, ) (int, error) { q := ` WITH rsks AS ( SELECT r.id, r.tenant_id, r.search_vector FROM risks r INNER JOIN risks_documents rd ON r.id = rd.risk_id WHERE rd.document_id = @document_id ) SELECT COUNT(id) FROM rsks WHERE %s AND %s ` q = fmt.Sprintf(q, scope.SQLFragment(), filter.SQLFragment()) args := pgx.NamedArgs{"document_id": documentID} maps.Copy(args, scope.SQLArguments()) maps.Copy(args, filter.SQLArguments()) row := conn.QueryRow(ctx, q, args) var count int if err := row.Scan(&count); err != nil { return 0, fmt.Errorf("cannot scan count: %w", err) } return count, nil }