From 22024434f5e1b333f0a5d5b780b0643d1bc2805c Mon Sep 17 00:00:00 2001 From: Tomasz Janiszewski Date: Tue, 21 Oct 2025 10:22:24 +0200 Subject: [PATCH] fix(pg): do not batch with copyFrom Signed-off-by: Tomasz Janiszewski --- .../internal/store/postgres/store.go | 103 +++-- .../internal/store/postgres/store.go | 75 ++-- .../usage/store/postgres/store.go | 67 ++-- .../internal/store/postgres/store.go | 129 +++---- .../internal/store/postgres/store.go | 65 ++-- central/auth/store/postgres/store.go | 99 +++-- .../internal/store/postgres/store.go | 63 ++- .../blob/datastore/store/postgres/store.go | 65 ++-- .../internal/store/postgres/store.go | 65 ++-- .../store/clusterhealth/postgres/store.go | 73 ++-- central/clusterinit/store/postgres/store.go | 61 ++- .../store/postgres/compliance_config/store.go | 61 ++- .../internal/store/postgres/domain/store.go | 61 ++- .../internal/store/postgres/metadata/store.go | 67 ++-- .../internal/store/postgres/results/store.go | 67 ++-- .../internal/store/postgres/strings/store.go | 61 ++- .../checkresults/store/postgres/store.go | 61 ++- .../profiles/store/postgres/store.go | 61 ++- .../rules/store/postgres/store.go | 61 ++- .../scans/store/postgres/store.go | 61 ++- .../store/postgres/store.go | 61 ++- .../v2/benchmarks/store/postgres/store.go | 105 +++-- .../v2/checkresults/store/postgres/store.go | 83 ++-- .../v2/integration/store/postgres/store.go | 69 ++-- .../v2/profiles/store/postgres/store.go | 111 +++--- .../v2/remediations/store/postgres/store.go | 67 ++-- .../v2/report/store/postgres/store.go | 117 +++--- .../v2/rules/store/postgres/store.go | 109 +++--- .../scanconfigstatus/store/postgres/store.go | 67 ++-- .../store/postgres/store.go | 173 ++++----- .../v2/scans/store/postgres/store.go | 77 ++-- .../store/postgres/store.go | 67 ++-- .../v2/suites/store/postgres/store.go | 65 ++-- .../cluster/datastore/store/postgres/store.go | 81 ++-- .../image/datastore/store/postgres/store.go | 83 ++-- .../v2/datastore/store/postgres/store.go | 95 +++-- .../node/datastore/store/postgres/store.go | 85 ++--- .../health/datastore/store/postgres/store.go | 61 ++- .../internal/store/postgres/store.go | 361 ++++++++---------- .../internal/store/postgres/store.go | 73 ++-- .../internal/store/postgres/store.go | 61 ++- .../internal/store/postgres/store.go | 69 ++-- .../hash/datastore/store/postgres/store.go | 61 ++- .../datastore/store/postgres/store.go | 75 ++-- .../v2/datastore/store/postgres/store.go | 81 ++-- .../imageintegration/store/postgres/store.go | 65 ++-- .../integrationhealth/store/postgres/store.go | 61 ++- central/logimbue/store/postgres/store.go | 63 ++- .../internal/store/postgres/store.go | 71 ++-- .../networkbaseline/store/postgres/store.go | 65 ++-- .../internal/store/postgres/store.go | 61 ++- .../internal/store/postgres/store.go | 67 ++-- .../internal/store/postgres/store.go | 65 ++-- .../undodeploymentstore/postgres/store.go | 61 ++- .../internal/undostore/postgres/store.go | 61 ++- .../datastore/store/postgres/store.go | 73 ++-- .../internal/store/postgres/store.go | 63 ++- .../internal/store/postgres/store.go | 105 +++-- central/policy/store/postgres/store.go | 83 ++-- .../policycategory/store/postgres/store.go | 63 ++- .../store/postgres/store.go | 65 ++-- .../processbaseline/store/postgres/store.go | 67 ++-- .../internal/store/postgres/store.go | 65 ++-- .../processindicator/store/postgres/store.go | 87 ++--- .../store/postgres/store.go | 79 ++-- .../k8srole/internal/store/postgres/store.go | 75 ++-- .../internal/store/postgres/store.go | 115 +++--- .../reports/config/store/postgres/store.go | 107 +++--- .../datastore/store/postgres/store.go | 81 ++-- .../datastore/store/postgres/store.go | 103 +++-- .../internal/store/postgres/store.go | 69 ++-- .../store/permissionset/postgres/store.go | 63 ++- central/role/store/role/postgres/store.go | 61 ++- .../store/simpleaccessscope/postgres/store.go | 63 ++- .../secret/internal/store/postgres/store.go | 149 ++++---- .../internal/store/postgres/store.go | 73 ++-- .../internal/store/postgres/store.go | 61 ++- .../store/postgres/store.go | 63 ++- .../internal/store/postgres/store.go | 69 ++-- .../internal/store/postgres/store.go | 207 +++++----- .../internal/store/postgres/store.go | 61 ++- .../multitest/postgres/store.go | 133 +++---- .../pg-table-bindings/store.go.tpl | 66 ++-- .../pg-table-bindings/test/postgres/store.go | 81 ++-- .../testgraphtables/testchild1/store.go | 63 ++- .../testgraphtables/testchild1p4/store.go | 65 ++-- .../testgraphtables/testchild2/store.go | 67 ++-- .../testg2grandchild1/store.go | 67 ++-- .../testg3grandchild1/store.go | 63 ++- .../testgraphtables/testggrandchild1/store.go | 63 ++- .../testgraphtables/testgrandchild1/store.go | 67 ++-- .../testgraphtables/testgrandparent/store.go | 143 +++---- .../testgraphtables/testparent1/store.go | 103 +++-- .../testgraphtables/testparent2/store.go | 65 ++-- .../testgraphtables/testparent3/store.go | 65 ++-- .../testgraphtables/testparent4/store.go | 65 ++-- .../testgraphtables/testshortcircuit/store.go | 65 ++-- .../testuuidkey/postgres/store.go | 81 ++-- 98 files changed, 3627 insertions(+), 4288 deletions(-) diff --git a/central/activecomponent/datastore/internal/store/postgres/store.go b/central/activecomponent/datastore/internal/store/postgres/store.go index 3f877077a3255..e64f32ad45411 100644 --- a/central/activecomponent/datastore/internal/store/postgres/store.go +++ b/central/activecomponent/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -144,12 +143,18 @@ func copyFromActiveComponents(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -158,39 +163,29 @@ func copyFromActiveComponents(ctx context.Context, s pgSearch.Deleter, tx *postg "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - pgutils.NilOrUUID(obj.GetDeploymentId()), - obj.GetComponentId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"active_components"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + pgutils.NilOrUUID(obj.GetDeploymentId()), + obj.GetComponentId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"active_components"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -206,8 +201,6 @@ func copyFromActiveComponentsActiveContextsSlices(ctx context.Context, s pgSearc if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "active_components_id", @@ -217,27 +210,23 @@ func copyFromActiveComponentsActiveContextsSlices(ctx context.Context, s pgSearc } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - activeComponentID, - idx, - obj.GetContainerName(), - obj.GetImageId(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"active_components_active_contexts_slices"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + activeComponentID, + idx, + obj.GetContainerName(), + obj.GetImageId(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"active_components_active_contexts_slices"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/administration/events/datastore/internal/store/postgres/store.go b/central/administration/events/datastore/internal/store/postgres/store.go index 90fc04b8b8952..ac77c0695c1b2 100644 --- a/central/administration/events/datastore/internal/store/postgres/store.go +++ b/central/administration/events/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -124,12 +123,18 @@ func copyFromAdministrationEvents(ctx context.Context, s pgSearch.Deleter, tx *p if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -143,44 +148,34 @@ func copyFromAdministrationEvents(ctx context.Context, s pgSearch.Deleter, tx *p "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetType(), - obj.GetLevel(), - obj.GetDomain(), - obj.GetResource().GetType(), - obj.GetNumOccurrences(), - protocompat.NilOrTime(obj.GetLastOccurredAt()), - protocompat.NilOrTime(obj.GetCreatedAt()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"administration_events"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetType(), + obj.GetLevel(), + obj.GetDomain(), + obj.GetResource().GetType(), + obj.GetNumOccurrences(), + protocompat.NilOrTime(obj.GetLastOccurredAt()), + protocompat.NilOrTime(obj.GetCreatedAt()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"administration_events"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/administration/usage/store/postgres/store.go b/central/administration/usage/store/postgres/store.go index 40e55faf9ff03..3e02d7258ae98 100644 --- a/central/administration/usage/store/postgres/store.go +++ b/central/administration/usage/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -120,12 +119,18 @@ func copyFromSecuredUnits(ctx context.Context, s pgSearch.Deleter, tx *postgres. if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -135,40 +140,30 @@ func copyFromSecuredUnits(ctx context.Context, s pgSearch.Deleter, tx *postgres. "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - protocompat.NilOrTime(obj.GetTimestamp()), - obj.GetNumNodes(), - obj.GetNumCpuUnits(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secured_units"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + protocompat.NilOrTime(obj.GetTimestamp()), + obj.GetNumNodes(), + obj.GetNumCpuUnits(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secured_units"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/alert/datastore/internal/store/postgres/store.go b/central/alert/datastore/internal/store/postgres/store.go index 16de8706f23ee..d68dd9deb8d4d 100644 --- a/central/alert/datastore/internal/store/postgres/store.go +++ b/central/alert/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -172,12 +171,18 @@ func copyFromAlerts(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, ob if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -218,71 +223,61 @@ func copyFromAlerts(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, ob "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetPolicy().GetId(), - obj.GetPolicy().GetName(), - obj.GetPolicy().GetDescription(), - obj.GetPolicy().GetDisabled(), - obj.GetPolicy().GetCategories(), - obj.GetPolicy().GetSeverity(), - obj.GetPolicy().GetEnforcementActions(), - protocompat.NilOrTime(obj.GetPolicy().GetLastUpdated()), - obj.GetPolicy().GetSORTName(), - obj.GetPolicy().GetSORTLifecycleStage(), - obj.GetPolicy().GetSORTEnforcement(), - obj.GetLifecycleStage(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetClusterName(), - obj.GetNamespace(), - pgutils.NilOrUUID(obj.GetNamespaceId()), - pgutils.NilOrUUID(obj.GetDeployment().GetId()), - obj.GetDeployment().GetName(), - obj.GetDeployment().GetInactive(), - obj.GetImage().GetId(), - obj.GetImage().GetName().GetRegistry(), - obj.GetImage().GetName().GetRemote(), - obj.GetImage().GetName().GetTag(), - obj.GetImage().GetName().GetFullName(), - obj.GetImage().GetIdV2(), - obj.GetNode().GetId(), - obj.GetNode().GetName(), - obj.GetResource().GetResourceType(), - obj.GetResource().GetName(), - obj.GetEnforcement().GetAction(), - protocompat.NilOrTime(obj.GetTime()), - obj.GetState(), - obj.GetPlatformComponent(), - obj.GetEntityType(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"alerts"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetPolicy().GetId(), + obj.GetPolicy().GetName(), + obj.GetPolicy().GetDescription(), + obj.GetPolicy().GetDisabled(), + obj.GetPolicy().GetCategories(), + obj.GetPolicy().GetSeverity(), + obj.GetPolicy().GetEnforcementActions(), + protocompat.NilOrTime(obj.GetPolicy().GetLastUpdated()), + obj.GetPolicy().GetSORTName(), + obj.GetPolicy().GetSORTLifecycleStage(), + obj.GetPolicy().GetSORTEnforcement(), + obj.GetLifecycleStage(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetClusterName(), + obj.GetNamespace(), + pgutils.NilOrUUID(obj.GetNamespaceId()), + pgutils.NilOrUUID(obj.GetDeployment().GetId()), + obj.GetDeployment().GetName(), + obj.GetDeployment().GetInactive(), + obj.GetImage().GetId(), + obj.GetImage().GetName().GetRegistry(), + obj.GetImage().GetName().GetRemote(), + obj.GetImage().GetName().GetTag(), + obj.GetImage().GetName().GetFullName(), + obj.GetImage().GetIdV2(), + obj.GetNode().GetId(), + obj.GetNode().GetName(), + obj.GetResource().GetResourceType(), + obj.GetResource().GetName(), + obj.GetEnforcement().GetAction(), + protocompat.NilOrTime(obj.GetTime()), + obj.GetState(), + obj.GetPlatformComponent(), + obj.GetEntityType(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"alerts"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/apitoken/datastore/internal/store/postgres/store.go b/central/apitoken/datastore/internal/store/postgres/store.go index 3c35e60de3cc7..d6a6cfa35921c 100644 --- a/central/apitoken/datastore/internal/store/postgres/store.go +++ b/central/apitoken/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -118,12 +117,18 @@ func copyFromAPITokens(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -132,39 +137,29 @@ func copyFromAPITokens(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - protocompat.NilOrTime(obj.GetExpiration()), - obj.GetRevoked(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"api_tokens"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + protocompat.NilOrTime(obj.GetExpiration()), + obj.GetRevoked(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"api_tokens"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/auth/store/postgres/store.go b/central/auth/store/postgres/store.go index e22cde6a0e774..9c29ca84268d4 100644 --- a/central/auth/store/postgres/store.go +++ b/central/auth/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -139,12 +138,18 @@ func copyFromAuthMachineToMachineConfigs(ctx context.Context, s pgSearch.Deleter if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -152,38 +157,28 @@ func copyFromAuthMachineToMachineConfigs(ctx context.Context, s pgSearch.Deleter "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetIssuer(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"auth_machine_to_machine_configs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetIssuer(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"auth_machine_to_machine_configs"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -199,8 +194,6 @@ func copyFromAuthMachineToMachineConfigsMappings(ctx context.Context, s pgSearch if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "auth_machine_to_machine_configs_id", @@ -209,26 +202,22 @@ func copyFromAuthMachineToMachineConfigsMappings(ctx context.Context, s pgSearch } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(authMachineToMachineConfigID), - idx, - obj.GetRole(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"auth_machine_to_machine_configs_mappings"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(authMachineToMachineConfigID), + idx, + obj.GetRole(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"auth_machine_to_machine_configs_mappings"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/authprovider/datastore/internal/store/postgres/store.go b/central/authprovider/datastore/internal/store/postgres/store.go index 890137aa98b32..6849d3e213158 100644 --- a/central/authprovider/datastore/internal/store/postgres/store.go +++ b/central/authprovider/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -116,12 +115,18 @@ func copyFromAuthProviders(ctx context.Context, s pgSearch.Deleter, tx *postgres if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -129,38 +134,28 @@ func copyFromAuthProviders(ctx context.Context, s pgSearch.Deleter, tx *postgres "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"auth_providers"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"auth_providers"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/blob/datastore/store/postgres/store.go b/central/blob/datastore/store/postgres/store.go index d395551d20189..7d985d3cbda9d 100644 --- a/central/blob/datastore/store/postgres/store.go +++ b/central/blob/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -118,12 +117,18 @@ func copyFromBlobs(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, obj if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetName()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "name", @@ -132,39 +137,29 @@ func copyFromBlobs(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, obj "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetName(), - obj.GetLength(), - protocompat.NilOrTime(obj.GetModifiedTime()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetName()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"blobs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetName(), + obj.GetLength(), + protocompat.NilOrTime(obj.GetModifiedTime()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"blobs"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/cloudsources/datastore/internal/store/postgres/store.go b/central/cloudsources/datastore/internal/store/postgres/store.go index 9b761f49fe104..01de038bbeda4 100644 --- a/central/cloudsources/datastore/internal/store/postgres/store.go +++ b/central/cloudsources/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -126,12 +125,18 @@ func copyFromCloudSources(ctx context.Context, s pgSearch.Deleter, tx *postgres. if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -140,39 +145,29 @@ func copyFromCloudSources(ctx context.Context, s pgSearch.Deleter, tx *postgres. "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - obj.GetType(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cloud_sources"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + obj.GetType(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cloud_sources"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/cluster/store/clusterhealth/postgres/store.go b/central/cluster/store/clusterhealth/postgres/store.go index ffc89cd1cb12d..9d2137c89d6c7 100644 --- a/central/cluster/store/clusterhealth/postgres/store.go +++ b/central/cluster/store/clusterhealth/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -123,12 +122,18 @@ func copyFromClusterHealthStatuses(ctx context.Context, s pgSearch.Deleter, tx * if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -141,43 +146,33 @@ func copyFromClusterHealthStatuses(ctx context.Context, s pgSearch.Deleter, tx * "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetSensorHealthStatus(), - obj.GetCollectorHealthStatus(), - obj.GetOverallHealthStatus(), - obj.GetAdmissionControlHealthStatus(), - obj.GetScannerHealthStatus(), - protocompat.NilOrTime(obj.GetLastContact()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cluster_health_statuses"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetSensorHealthStatus(), + obj.GetCollectorHealthStatus(), + obj.GetOverallHealthStatus(), + obj.GetAdmissionControlHealthStatus(), + obj.GetScannerHealthStatus(), + protocompat.NilOrTime(obj.GetLastContact()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cluster_health_statuses"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/clusterinit/store/postgres/store.go b/central/clusterinit/store/postgres/store.go index 95a9089d611b4..a7e7a6a6eaa0a 100644 --- a/central/clusterinit/store/postgres/store.go +++ b/central/clusterinit/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -120,49 +119,45 @@ func copyFromClusterInitBundles(ctx context.Context, s pgSearch.Deleter, tx *pos if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cluster_init_bundles"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cluster_init_bundles"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/compliance/datastore/internal/store/postgres/compliance_config/store.go b/central/compliance/datastore/internal/store/postgres/compliance_config/store.go index 8d0e74d3960eb..0b3a4985c1bb3 100644 --- a/central/compliance/datastore/internal/store/postgres/compliance_config/store.go +++ b/central/compliance/datastore/internal/store/postgres/compliance_config/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromComplianceConfigs(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetStandardId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "standardid", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetStandardId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetStandardId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_configs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetStandardId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_configs"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/compliance/datastore/internal/store/postgres/domain/store.go b/central/compliance/datastore/internal/store/postgres/domain/store.go index 06410e1d3f255..799924950febf 100644 --- a/central/compliance/datastore/internal/store/postgres/domain/store.go +++ b/central/compliance/datastore/internal/store/postgres/domain/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -115,49 +114,45 @@ func copyFromComplianceDomains(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_domains"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_domains"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/compliance/datastore/internal/store/postgres/metadata/store.go b/central/compliance/datastore/internal/store/postgres/metadata/store.go index 270aa2824e056..4011dfdf52b04 100644 --- a/central/compliance/datastore/internal/store/postgres/metadata/store.go +++ b/central/compliance/datastore/internal/store/postgres/metadata/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -141,12 +140,18 @@ func copyFromComplianceRunMetadata(ctx context.Context, s pgSearch.Deleter, tx * if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetRunId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "runid", @@ -156,40 +161,30 @@ func copyFromComplianceRunMetadata(ctx context.Context, s pgSearch.Deleter, tx * "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetRunId(), - obj.GetStandardId(), - pgutils.NilOrUUID(obj.GetClusterId()), - protocompat.NilOrTime(obj.GetFinishTimestamp()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetRunId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_run_metadata"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetRunId(), + obj.GetStandardId(), + pgutils.NilOrUUID(obj.GetClusterId()), + protocompat.NilOrTime(obj.GetFinishTimestamp()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_run_metadata"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/compliance/datastore/internal/store/postgres/results/store.go b/central/compliance/datastore/internal/store/postgres/results/store.go index 53fb67e094c46..2bd5094862078 100644 --- a/central/compliance/datastore/internal/store/postgres/results/store.go +++ b/central/compliance/datastore/internal/store/postgres/results/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -141,12 +140,18 @@ func copyFromComplianceRunResults(ctx context.Context, s pgSearch.Deleter, tx *p if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetRunMetadata().GetRunId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "runmetadata_runid", @@ -156,40 +161,30 @@ func copyFromComplianceRunResults(ctx context.Context, s pgSearch.Deleter, tx *p "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetRunMetadata().GetRunId(), - obj.GetRunMetadata().GetStandardId(), - pgutils.NilOrUUID(obj.GetRunMetadata().GetClusterId()), - protocompat.NilOrTime(obj.GetRunMetadata().GetFinishTimestamp()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetRunMetadata().GetRunId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_run_results"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetRunMetadata().GetRunId(), + obj.GetRunMetadata().GetStandardId(), + pgutils.NilOrUUID(obj.GetRunMetadata().GetClusterId()), + protocompat.NilOrTime(obj.GetRunMetadata().GetFinishTimestamp()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_run_results"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/compliance/datastore/internal/store/postgres/strings/store.go b/central/compliance/datastore/internal/store/postgres/strings/store.go index 6a31666aede52..31f60af4cede3 100644 --- a/central/compliance/datastore/internal/store/postgres/strings/store.go +++ b/central/compliance/datastore/internal/store/postgres/strings/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromComplianceStrings(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_strings"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_strings"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/checkresults/store/postgres/store.go b/central/complianceoperator/checkresults/store/postgres/store.go index ba3ada0405f65..581ed40d561ac 100644 --- a/central/complianceoperator/checkresults/store/postgres/store.go +++ b/central/complianceoperator/checkresults/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromComplianceOperatorCheckResults(ctx context.Context, s pgSearch.Dele if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_check_results"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_check_results"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/profiles/store/postgres/store.go b/central/complianceoperator/profiles/store/postgres/store.go index 233ffcdff5a99..f2ebab0f56b84 100644 --- a/central/complianceoperator/profiles/store/postgres/store.go +++ b/central/complianceoperator/profiles/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromComplianceOperatorProfiles(ctx context.Context, s pgSearch.Deleter, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_profiles"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_profiles"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/rules/store/postgres/store.go b/central/complianceoperator/rules/store/postgres/store.go index fdb78bcb336ca..30ffdfb03fd47 100644 --- a/central/complianceoperator/rules/store/postgres/store.go +++ b/central/complianceoperator/rules/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromComplianceOperatorRules(ctx context.Context, s pgSearch.Deleter, tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_rules"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_rules"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/scans/store/postgres/store.go b/central/complianceoperator/scans/store/postgres/store.go index 056f19d5e08bc..3da850e2f4485 100644 --- a/central/complianceoperator/scans/store/postgres/store.go +++ b/central/complianceoperator/scans/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromComplianceOperatorScans(ctx context.Context, s pgSearch.Deleter, tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scans"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scans"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/scansettingbinding/store/postgres/store.go b/central/complianceoperator/scansettingbinding/store/postgres/store.go index 645910ce78f81..8e02b1337552f 100644 --- a/central/complianceoperator/scansettingbinding/store/postgres/store.go +++ b/central/complianceoperator/scansettingbinding/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromComplianceOperatorScanSettingBindings(ctx context.Context, s pgSear if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_setting_bindings"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_setting_bindings"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/benchmarks/store/postgres/store.go b/central/complianceoperator/v2/benchmarks/store/postgres/store.go index ffd6c9e8f9dd9..ef67ace208d6d 100644 --- a/central/complianceoperator/v2/benchmarks/store/postgres/store.go +++ b/central/complianceoperator/v2/benchmarks/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -145,12 +144,18 @@ func copyFromComplianceOperatorBenchmarkV2(ctx context.Context, s pgSearch.Delet if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -160,40 +165,30 @@ func copyFromComplianceOperatorBenchmarkV2(ctx context.Context, s pgSearch.Delet "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - obj.GetVersion(), - obj.GetShortName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_benchmark_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + obj.GetVersion(), + obj.GetShortName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_benchmark_v2"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -209,8 +204,6 @@ func copyFromComplianceOperatorBenchmarkV2Profiles(ctx context.Context, s pgSear if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "compliance_operator_benchmark_v2_id", @@ -220,27 +213,23 @@ func copyFromComplianceOperatorBenchmarkV2Profiles(ctx context.Context, s pgSear } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(complianceOperatorBenchmarkV2ID), - idx, - obj.GetProfileName(), - obj.GetProfileVersion(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_benchmark_v2_profiles"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(complianceOperatorBenchmarkV2ID), + idx, + obj.GetProfileName(), + obj.GetProfileVersion(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_benchmark_v2_profiles"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/checkresults/store/postgres/store.go b/central/complianceoperator/v2/checkresults/store/postgres/store.go index ad1bfd4c24fb8..3fcafbf846979 100644 --- a/central/complianceoperator/v2/checkresults/store/postgres/store.go +++ b/central/complianceoperator/v2/checkresults/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -149,12 +148,18 @@ func copyFromComplianceOperatorCheckResultV2(ctx context.Context, s pgSearch.Del if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -172,48 +177,38 @@ func copyFromComplianceOperatorCheckResultV2(ctx context.Context, s pgSearch.Del "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetCheckId(), - obj.GetCheckName(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetStatus(), - obj.GetSeverity(), - protocompat.NilOrTime(obj.GetCreatedTime()), - obj.GetScanConfigName(), - obj.GetRationale(), - pgutils.NilOrUUID(obj.GetScanRefId()), - pgutils.NilOrUUID(obj.GetRuleRefId()), - protocompat.NilOrTime(obj.GetLastStartedTime()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_check_result_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetCheckId(), + obj.GetCheckName(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetStatus(), + obj.GetSeverity(), + protocompat.NilOrTime(obj.GetCreatedTime()), + obj.GetScanConfigName(), + obj.GetRationale(), + pgutils.NilOrUUID(obj.GetScanRefId()), + pgutils.NilOrUUID(obj.GetRuleRefId()), + protocompat.NilOrTime(obj.GetLastStartedTime()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_check_result_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/integration/store/postgres/store.go b/central/complianceoperator/v2/integration/store/postgres/store.go index 6e999b4845812..d11ea32fa916d 100644 --- a/central/complianceoperator/v2/integration/store/postgres/store.go +++ b/central/complianceoperator/v2/integration/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -141,12 +140,18 @@ func copyFromComplianceIntegrations(ctx context.Context, s pgSearch.Deleter, tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -157,41 +162,31 @@ func copyFromComplianceIntegrations(ctx context.Context, s pgSearch.Deleter, tx "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetVersion(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetOperatorInstalled(), - obj.GetOperatorStatus(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_integrations"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetVersion(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetOperatorInstalled(), + obj.GetOperatorStatus(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_integrations"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/profiles/store/postgres/store.go b/central/complianceoperator/v2/profiles/store/postgres/store.go index 31ba22eee7328..a1e18d2a542db 100644 --- a/central/complianceoperator/v2/profiles/store/postgres/store.go +++ b/central/complianceoperator/v2/profiles/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -169,12 +168,18 @@ func copyFromComplianceOperatorProfileV2(ctx context.Context, s pgSearch.Deleter if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -188,44 +193,34 @@ func copyFromComplianceOperatorProfileV2(ctx context.Context, s pgSearch.Deleter "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetProfileId(), - obj.GetName(), - obj.GetProfileVersion(), - obj.GetProductType(), - obj.GetStandard(), - pgutils.NilOrUUID(obj.GetClusterId()), - pgutils.NilOrUUID(obj.GetProfileRefId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_profile_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetProfileId(), + obj.GetName(), + obj.GetProfileVersion(), + obj.GetProductType(), + obj.GetStandard(), + pgutils.NilOrUUID(obj.GetClusterId()), + pgutils.NilOrUUID(obj.GetProfileRefId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_profile_v2"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -241,8 +236,6 @@ func copyFromComplianceOperatorProfileV2Rules(ctx context.Context, s pgSearch.De if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "compliance_operator_profile_v2_id", @@ -251,26 +244,22 @@ func copyFromComplianceOperatorProfileV2Rules(ctx context.Context, s pgSearch.De } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - complianceOperatorProfileV2ID, - idx, - obj.GetRuleName(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_profile_v2_rules"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + complianceOperatorProfileV2ID, + idx, + obj.GetRuleName(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_profile_v2_rules"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/remediations/store/postgres/store.go b/central/complianceoperator/v2/remediations/store/postgres/store.go index 04651a51ef30d..d24db4f925324 100644 --- a/central/complianceoperator/v2/remediations/store/postgres/store.go +++ b/central/complianceoperator/v2/remediations/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -140,12 +139,18 @@ func copyFromComplianceOperatorRemediationV2(ctx context.Context, s pgSearch.Del if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -155,40 +160,30 @@ func copyFromComplianceOperatorRemediationV2(ctx context.Context, s pgSearch.Del "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - obj.GetComplianceCheckResultName(), - obj.GetClusterId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_remediation_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + obj.GetComplianceCheckResultName(), + obj.GetClusterId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_remediation_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/report/store/postgres/store.go b/central/complianceoperator/v2/report/store/postgres/store.go index bf8d83eb77c08..596816225ef31 100644 --- a/central/complianceoperator/v2/report/store/postgres/store.go +++ b/central/complianceoperator/v2/report/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -152,12 +151,18 @@ func copyFromComplianceOperatorReportSnapshotV2(ctx context.Context, s pgSearch. if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetReportId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "reportid", @@ -173,46 +178,36 @@ func copyFromComplianceOperatorReportSnapshotV2(ctx context.Context, s pgSearch. "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetReportId()), - obj.GetScanConfigurationId(), - obj.GetName(), - obj.GetReportStatus().GetRunState(), - protocompat.NilOrTime(obj.GetReportStatus().GetStartedAt()), - protocompat.NilOrTime(obj.GetReportStatus().GetCompletedAt()), - obj.GetReportStatus().GetReportRequestType(), - obj.GetReportStatus().GetReportNotificationMethod(), - obj.GetUser().GetId(), - obj.GetUser().GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetReportId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_report_snapshot_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetReportId()), + obj.GetScanConfigurationId(), + obj.GetName(), + obj.GetReportStatus().GetRunState(), + protocompat.NilOrTime(obj.GetReportStatus().GetStartedAt()), + protocompat.NilOrTime(obj.GetReportStatus().GetCompletedAt()), + obj.GetReportStatus().GetReportRequestType(), + obj.GetReportStatus().GetReportNotificationMethod(), + obj.GetUser().GetId(), + obj.GetUser().GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_report_snapshot_v2"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -228,8 +223,6 @@ func copyFromComplianceOperatorReportSnapshotV2Scans(ctx context.Context, s pgSe if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "compliance_operator_report_snapshot_v2_reportid", @@ -239,27 +232,23 @@ func copyFromComplianceOperatorReportSnapshotV2Scans(ctx context.Context, s pgSe } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(complianceOperatorReportSnapshotV2ReportId), - idx, - obj.GetScanRefId(), - protocompat.NilOrTime(obj.GetLastStartedTime()), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_report_snapshot_v2_scans"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(complianceOperatorReportSnapshotV2ReportId), + idx, + obj.GetScanRefId(), + protocompat.NilOrTime(obj.GetLastStartedTime()), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_report_snapshot_v2_scans"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/rules/store/postgres/store.go b/central/complianceoperator/v2/rules/store/postgres/store.go index 435358dd58062..2a655f5bc5110 100644 --- a/central/complianceoperator/v2/rules/store/postgres/store.go +++ b/central/complianceoperator/v2/rules/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -168,12 +167,18 @@ func copyFromComplianceOperatorRuleV2(ctx context.Context, s pgSearch.Deleter, t if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -185,42 +190,32 @@ func copyFromComplianceOperatorRuleV2(ctx context.Context, s pgSearch.Deleter, t "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetRuleType(), - obj.GetSeverity(), - pgutils.NilOrUUID(obj.GetClusterId()), - pgutils.NilOrUUID(obj.GetRuleRefId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_rule_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetRuleType(), + obj.GetSeverity(), + pgutils.NilOrUUID(obj.GetClusterId()), + pgutils.NilOrUUID(obj.GetRuleRefId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_rule_v2"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -236,8 +231,6 @@ func copyFromComplianceOperatorRuleV2Controls(ctx context.Context, s pgSearch.De if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "compliance_operator_rule_v2_id", @@ -247,27 +240,23 @@ func copyFromComplianceOperatorRuleV2Controls(ctx context.Context, s pgSearch.De } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - complianceOperatorRuleV2ID, - idx, - obj.GetStandard(), - obj.GetControl(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_rule_v2_controls"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + complianceOperatorRuleV2ID, + idx, + obj.GetStandard(), + obj.GetControl(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_rule_v2_controls"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/scanconfigurations/scanconfigstatus/store/postgres/store.go b/central/complianceoperator/v2/scanconfigurations/scanconfigstatus/store/postgres/store.go index 499310e33ad11..a70e5d1e2e43f 100644 --- a/central/complianceoperator/v2/scanconfigurations/scanconfigstatus/store/postgres/store.go +++ b/central/complianceoperator/v2/scanconfigurations/scanconfigstatus/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -141,12 +140,18 @@ func copyFromComplianceOperatorClusterScanConfigStatuses(ctx context.Context, s if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -156,40 +161,30 @@ func copyFromComplianceOperatorClusterScanConfigStatuses(ctx context.Context, s "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - pgutils.NilOrUUID(obj.GetClusterId()), - pgutils.NilOrUUID(obj.GetScanConfigId()), - protocompat.NilOrTime(obj.GetLastUpdatedTime()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_cluster_scan_config_statuses"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + pgutils.NilOrUUID(obj.GetClusterId()), + pgutils.NilOrUUID(obj.GetScanConfigId()), + protocompat.NilOrTime(obj.GetLastUpdatedTime()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_cluster_scan_config_statuses"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/scanconfigurations/store/postgres/store.go b/central/complianceoperator/v2/scanconfigurations/store/postgres/store.go index a6981fec51682..453c65ab4a233 100644 --- a/central/complianceoperator/v2/scanconfigurations/store/postgres/store.go +++ b/central/complianceoperator/v2/scanconfigurations/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -189,12 +188,18 @@ func copyFromComplianceOperatorScanConfigurationV2(ctx context.Context, s pgSear if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -203,39 +208,29 @@ func copyFromComplianceOperatorScanConfigurationV2(ctx context.Context, s pgSear "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetScanConfigName(), - obj.GetModifiedBy().GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetScanConfigName(), + obj.GetModifiedBy().GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -257,8 +252,6 @@ func copyFromComplianceOperatorScanConfigurationV2Profiles(ctx context.Context, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "compliance_operator_scan_configuration_v2_id", @@ -267,26 +260,22 @@ func copyFromComplianceOperatorScanConfigurationV2Profiles(ctx context.Context, } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(complianceOperatorScanConfigurationV2ID), - idx, - obj.GetProfileName(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2_profiles"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(complianceOperatorScanConfigurationV2ID), + idx, + obj.GetProfileName(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2_profiles"}, copyCols, inputRows); err != nil { + return err } return nil @@ -296,8 +285,6 @@ func copyFromComplianceOperatorScanConfigurationV2Clusters(ctx context.Context, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "compliance_operator_scan_configuration_v2_id", @@ -306,26 +293,22 @@ func copyFromComplianceOperatorScanConfigurationV2Clusters(ctx context.Context, } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(complianceOperatorScanConfigurationV2ID), - idx, - pgutils.NilOrUUID(obj.GetClusterId()), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2_clusters"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(complianceOperatorScanConfigurationV2ID), + idx, + pgutils.NilOrUUID(obj.GetClusterId()), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2_clusters"}, copyCols, inputRows); err != nil { + return err } return nil @@ -335,8 +318,6 @@ func copyFromComplianceOperatorScanConfigurationV2Notifiers(ctx context.Context, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "compliance_operator_scan_configuration_v2_id", @@ -345,26 +326,22 @@ func copyFromComplianceOperatorScanConfigurationV2Notifiers(ctx context.Context, } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(complianceOperatorScanConfigurationV2ID), - idx, - obj.GetId(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2_notifiers"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(complianceOperatorScanConfigurationV2ID), + idx, + obj.GetId(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_configuration_v2_notifiers"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/scans/store/postgres/store.go b/central/complianceoperator/v2/scans/store/postgres/store.go index a01098c57a96f..41af4beca928e 100644 --- a/central/complianceoperator/v2/scans/store/postgres/store.go +++ b/central/complianceoperator/v2/scans/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -146,12 +145,18 @@ func copyFromComplianceOperatorScanV2(ctx context.Context, s pgSearch.Deleter, t if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -166,45 +171,35 @@ func copyFromComplianceOperatorScanV2(ctx context.Context, s pgSearch.Deleter, t "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetScanConfigName(), - pgutils.NilOrUUID(obj.GetClusterId()), - pgutils.NilOrUUID(obj.GetProfile().GetProfileRefId()), - obj.GetStatus().GetResult(), - protocompat.NilOrTime(obj.GetLastExecutedTime()), - obj.GetScanName(), - pgutils.NilOrUUID(obj.GetScanRefId()), - protocompat.NilOrTime(obj.GetLastStartedTime()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetScanConfigName(), + pgutils.NilOrUUID(obj.GetClusterId()), + pgutils.NilOrUUID(obj.GetProfile().GetProfileRefId()), + obj.GetStatus().GetResult(), + protocompat.NilOrTime(obj.GetLastExecutedTime()), + obj.GetScanName(), + pgutils.NilOrUUID(obj.GetScanRefId()), + protocompat.NilOrTime(obj.GetLastStartedTime()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/scansettingbindings/store/postgres/store.go b/central/complianceoperator/v2/scansettingbindings/store/postgres/store.go index 29d48821d2a74..aa02048b331a4 100644 --- a/central/complianceoperator/v2/scansettingbindings/store/postgres/store.go +++ b/central/complianceoperator/v2/scansettingbindings/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -140,12 +139,18 @@ func copyFromComplianceOperatorScanSettingBindingV2(ctx context.Context, s pgSea if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -155,40 +160,30 @@ func copyFromComplianceOperatorScanSettingBindingV2(ctx context.Context, s pgSea "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetScanSettingName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_setting_binding_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetScanSettingName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_scan_setting_binding_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/complianceoperator/v2/suites/store/postgres/store.go b/central/complianceoperator/v2/suites/store/postgres/store.go index 48d6511b6ce91..6d43386d01f6c 100644 --- a/central/complianceoperator/v2/suites/store/postgres/store.go +++ b/central/complianceoperator/v2/suites/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -139,12 +138,18 @@ func copyFromComplianceOperatorSuiteV2(ctx context.Context, s pgSearch.Deleter, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -153,39 +158,29 @@ func copyFromComplianceOperatorSuiteV2(ctx context.Context, s pgSearch.Deleter, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - pgutils.NilOrUUID(obj.GetClusterId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_suite_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + pgutils.NilOrUUID(obj.GetClusterId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"compliance_operator_suite_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/cve/cluster/datastore/store/postgres/store.go b/central/cve/cluster/datastore/store/postgres/store.go index 3b688ccb86ba0..375a70c8e34a3 100644 --- a/central/cve/cluster/datastore/store/postgres/store.go +++ b/central/cve/cluster/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -126,12 +125,18 @@ func copyFromClusterCves(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -148,47 +153,37 @@ func copyFromClusterCves(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetCveBaseInfo().GetCve(), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), - obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), - obj.GetCvss(), - obj.GetSeverity(), - obj.GetImpactScore(), - obj.GetSnoozed(), - protocompat.NilOrTime(obj.GetSnoozeExpiry()), - obj.GetType(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cluster_cves"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetCveBaseInfo().GetCve(), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), + obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), + obj.GetCvss(), + obj.GetSeverity(), + obj.GetImpactScore(), + obj.GetSnoozed(), + protocompat.NilOrTime(obj.GetSnoozeExpiry()), + obj.GetType(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"cluster_cves"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/cve/image/datastore/store/postgres/store.go b/central/cve/image/datastore/store/postgres/store.go index 575500a059a59..25615b49ac2a7 100644 --- a/central/cve/image/datastore/store/postgres/store.go +++ b/central/cve/image/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -127,12 +126,18 @@ func copyFromImageCves(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -150,48 +155,38 @@ func copyFromImageCves(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetCveBaseInfo().GetCve(), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), - obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), - obj.GetOperatingSystem(), - obj.GetCvss(), - obj.GetSeverity(), - obj.GetImpactScore(), - obj.GetSnoozed(), - protocompat.NilOrTime(obj.GetSnoozeExpiry()), - obj.GetNvdcvss(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_cves"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetCveBaseInfo().GetCve(), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), + obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), + obj.GetOperatingSystem(), + obj.GetCvss(), + obj.GetSeverity(), + obj.GetImpactScore(), + obj.GetSnoozed(), + protocompat.NilOrTime(obj.GetSnoozeExpiry()), + obj.GetNvdcvss(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_cves"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/cve/image/v2/datastore/store/postgres/store.go b/central/cve/image/v2/datastore/store/postgres/store.go index 521ce3f6de80e..facb0a992dcef 100644 --- a/central/cve/image/v2/datastore/store/postgres/store.go +++ b/central/cve/image/v2/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -134,12 +133,18 @@ func copyFromImageCvesV2(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -163,54 +168,44 @@ func copyFromImageCvesV2(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetImageId(), - obj.GetCveBaseInfo().GetCve(), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), - obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), - obj.GetCvss(), - obj.GetSeverity(), - obj.GetImpactScore(), - obj.GetNvdcvss(), - protocompat.NilOrTime(obj.GetFirstImageOccurrence()), - obj.GetState(), - obj.GetIsFixable(), - obj.GetFixedBy(), - obj.GetComponentId(), - obj.GetAdvisory().GetName(), - obj.GetAdvisory().GetLink(), - pgutils.NilOrString(obj.GetImageIdV2()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_cves_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetImageId(), + obj.GetCveBaseInfo().GetCve(), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), + obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), + obj.GetCvss(), + obj.GetSeverity(), + obj.GetImpactScore(), + obj.GetNvdcvss(), + protocompat.NilOrTime(obj.GetFirstImageOccurrence()), + obj.GetState(), + obj.GetIsFixable(), + obj.GetFixedBy(), + obj.GetComponentId(), + obj.GetAdvisory().GetName(), + obj.GetAdvisory().GetLink(), + pgutils.NilOrString(obj.GetImageIdV2()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_cves_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/cve/node/datastore/store/postgres/store.go b/central/cve/node/datastore/store/postgres/store.go index 2ee7796b6ca88..5e591d5dcae1c 100644 --- a/central/cve/node/datastore/store/postgres/store.go +++ b/central/cve/node/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -128,12 +127,18 @@ func copyFromNodeCves(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -152,49 +157,39 @@ func copyFromNodeCves(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetCveBaseInfo().GetCve(), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), - protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), - obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), - obj.GetOperatingSystem(), - obj.GetCvss(), - obj.GetSeverity(), - obj.GetImpactScore(), - obj.GetSnoozed(), - protocompat.NilOrTime(obj.GetSnoozeExpiry()), - obj.GetOrphaned(), - protocompat.NilOrTime(obj.GetOrphanedTime()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"node_cves"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetCveBaseInfo().GetCve(), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetPublishedOn()), + protocompat.NilOrTime(obj.GetCveBaseInfo().GetCreatedAt()), + obj.GetCveBaseInfo().GetEpss().GetEpssProbability(), + obj.GetOperatingSystem(), + obj.GetCvss(), + obj.GetSeverity(), + obj.GetImpactScore(), + obj.GetSnoozed(), + protocompat.NilOrTime(obj.GetSnoozeExpiry()), + obj.GetOrphaned(), + protocompat.NilOrTime(obj.GetOrphanedTime()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"node_cves"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/declarativeconfig/health/datastore/store/postgres/store.go b/central/declarativeconfig/health/datastore/store/postgres/store.go index 9e91cc9219e51..f833c82ca5ce8 100644 --- a/central/declarativeconfig/health/datastore/store/postgres/store.go +++ b/central/declarativeconfig/health/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -113,49 +112,45 @@ func copyFromDeclarativeConfigHealths(ctx context.Context, s pgSearch.Deleter, t if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"declarative_config_healths"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"declarative_config_healths"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/deployment/datastore/internal/store/postgres/store.go b/central/deployment/datastore/internal/store/postgres/store.go index 5c63a7e163eb5..631a1e3acce72 100644 --- a/central/deployment/datastore/internal/store/postgres/store.go +++ b/central/deployment/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -339,12 +338,18 @@ func copyFromDeployments(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -368,54 +373,44 @@ func copyFromDeployments(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - obj.GetType(), - obj.GetNamespace(), - pgutils.NilOrUUID(obj.GetNamespaceId()), - obj.GetOrchestratorComponent(), - pgutils.EmptyOrMap(obj.GetLabels()), - pgutils.EmptyOrMap(obj.GetPodLabels()), - protocompat.NilOrTime(obj.GetCreated()), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetClusterName(), - pgutils.EmptyOrMap(obj.GetAnnotations()), - obj.GetPriority(), - obj.GetImagePullSecrets(), - obj.GetServiceAccount(), - obj.GetServiceAccountPermissionLevel(), - obj.GetRiskScore(), - obj.GetPlatformComponent(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + obj.GetType(), + obj.GetNamespace(), + pgutils.NilOrUUID(obj.GetNamespaceId()), + obj.GetOrchestratorComponent(), + pgutils.EmptyOrMap(obj.GetLabels()), + pgutils.EmptyOrMap(obj.GetPodLabels()), + protocompat.NilOrTime(obj.GetCreated()), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetClusterName(), + pgutils.EmptyOrMap(obj.GetAnnotations()), + obj.GetPriority(), + obj.GetImagePullSecrets(), + obj.GetServiceAccount(), + obj.GetServiceAccountPermissionLevel(), + obj.GetRiskScore(), + obj.GetPlatformComponent(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -434,8 +429,6 @@ func copyFromDeploymentsContainers(ctx context.Context, s pgSearch.Deleter, tx * if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "deployments_id", @@ -457,42 +450,38 @@ func copyFromDeploymentsContainers(ctx context.Context, s pgSearch.Deleter, tx * } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(deploymentID), - idx, - obj.GetImage().GetId(), - obj.GetImage().GetName().GetRegistry(), - obj.GetImage().GetName().GetRemote(), - obj.GetImage().GetName().GetTag(), - obj.GetImage().GetName().GetFullName(), - pgutils.NilOrString(obj.GetImage().GetIdV2()), - obj.GetSecurityContext().GetPrivileged(), - obj.GetSecurityContext().GetDropCapabilities(), - obj.GetSecurityContext().GetAddCapabilities(), - obj.GetSecurityContext().GetReadOnlyRootFilesystem(), - obj.GetResources().GetCpuCoresRequest(), - obj.GetResources().GetCpuCoresLimit(), - obj.GetResources().GetMemoryMbRequest(), - obj.GetResources().GetMemoryMbLimit(), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(deploymentID), + idx, + obj.GetImage().GetId(), + obj.GetImage().GetName().GetRegistry(), + obj.GetImage().GetName().GetRemote(), + obj.GetImage().GetName().GetTag(), + obj.GetImage().GetName().GetFullName(), + pgutils.NilOrString(obj.GetImage().GetIdV2()), + obj.GetSecurityContext().GetPrivileged(), + obj.GetSecurityContext().GetDropCapabilities(), + obj.GetSecurityContext().GetAddCapabilities(), + obj.GetSecurityContext().GetReadOnlyRootFilesystem(), + obj.GetResources().GetCpuCoresRequest(), + obj.GetResources().GetCpuCoresLimit(), + obj.GetResources().GetMemoryMbRequest(), + obj.GetResources().GetMemoryMbLimit(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers"}, copyCols, inputRows); err != nil { + return err } - for idx, obj := range objs { + for _, obj := range objs { if err := copyFromDeploymentsContainersEnvs(ctx, s, tx, deploymentID, idx, obj.GetConfig().GetEnv()...); err != nil { return err } @@ -511,8 +500,6 @@ func copyFromDeploymentsContainersEnvs(ctx context.Context, s pgSearch.Deleter, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "deployments_id", @@ -524,29 +511,25 @@ func copyFromDeploymentsContainersEnvs(ctx context.Context, s pgSearch.Deleter, } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(deploymentID), - deploymentContainerIdx, - idx, - obj.GetKey(), - obj.GetValue(), - obj.GetEnvVarSource(), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent + return []interface{}{ + pgutils.NilOrUUID(deploymentID), + deploymentContainerIdx, + idx, + obj.GetKey(), + obj.GetValue(), + obj.GetEnvVarSource(), + }, nil + }) - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers_envs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers_envs"}, copyCols, inputRows); err != nil { + return err } return nil @@ -556,8 +539,6 @@ func copyFromDeploymentsContainersVolumes(ctx context.Context, s pgSearch.Delete if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "deployments_id", @@ -571,31 +552,27 @@ func copyFromDeploymentsContainersVolumes(ctx context.Context, s pgSearch.Delete } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(deploymentID), - deploymentContainerIdx, - idx, - obj.GetName(), - obj.GetSource(), - obj.GetDestination(), - obj.GetReadOnly(), - obj.GetType(), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent + return []interface{}{ + pgutils.NilOrUUID(deploymentID), + deploymentContainerIdx, + idx, + obj.GetName(), + obj.GetSource(), + obj.GetDestination(), + obj.GetReadOnly(), + obj.GetType(), + }, nil + }) - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers_volumes"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers_volumes"}, copyCols, inputRows); err != nil { + return err } return nil @@ -605,8 +582,6 @@ func copyFromDeploymentsContainersSecrets(ctx context.Context, s pgSearch.Delete if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "deployments_id", @@ -617,28 +592,24 @@ func copyFromDeploymentsContainersSecrets(ctx context.Context, s pgSearch.Delete } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(deploymentID), - deploymentContainerIdx, - idx, - obj.GetName(), - obj.GetPath(), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent + return []interface{}{ + pgutils.NilOrUUID(deploymentID), + deploymentContainerIdx, + idx, + obj.GetName(), + obj.GetPath(), + }, nil + }) - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers_secrets"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_containers_secrets"}, copyCols, inputRows); err != nil { + return err } return nil @@ -648,8 +619,6 @@ func copyFromDeploymentsPorts(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "deployments_id", @@ -660,31 +629,27 @@ func copyFromDeploymentsPorts(ctx context.Context, s pgSearch.Deleter, tx *postg } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(deploymentID), - idx, - obj.GetContainerPort(), - obj.GetProtocol(), - obj.GetExposure(), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent + return []interface{}{ + pgutils.NilOrUUID(deploymentID), + idx, + obj.GetContainerPort(), + obj.GetProtocol(), + obj.GetExposure(), + }, nil + }) - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_ports"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_ports"}, copyCols, inputRows); err != nil { + return err } - for idx, obj := range objs { + for _, obj := range objs { if err := copyFromDeploymentsPortsExposureInfos(ctx, s, tx, deploymentID, idx, obj.GetExposureInfos()...); err != nil { return err } @@ -697,8 +662,6 @@ func copyFromDeploymentsPortsExposureInfos(ctx context.Context, s pgSearch.Delet if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "deployments_id", @@ -713,32 +676,28 @@ func copyFromDeploymentsPortsExposureInfos(ctx context.Context, s pgSearch.Delet } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(deploymentID), - deploymentPortIdx, - idx, - obj.GetLevel(), - obj.GetServiceName(), - obj.GetServicePort(), - obj.GetNodePort(), - obj.GetExternalIps(), - obj.GetExternalHostnames(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_ports_exposure_infos"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(deploymentID), + deploymentPortIdx, + idx, + obj.GetLevel(), + obj.GetServiceName(), + obj.GetServicePort(), + obj.GetNodePort(), + obj.GetExternalIps(), + obj.GetExternalHostnames(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"deployments_ports_exposure_infos"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/discoveredclusters/datastore/internal/store/postgres/store.go b/central/discoveredclusters/datastore/internal/store/postgres/store.go index b5c058a304b70..6ea5f01d76fbb 100644 --- a/central/discoveredclusters/datastore/internal/store/postgres/store.go +++ b/central/discoveredclusters/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -123,12 +122,18 @@ func copyFromDiscoveredClusters(ctx context.Context, s pgSearch.Deleter, tx *pos if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -141,43 +146,33 @@ func copyFromDiscoveredClusters(ctx context.Context, s pgSearch.Deleter, tx *pos "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetMetadata().GetName(), - obj.GetMetadata().GetType(), - protocompat.NilOrTime(obj.GetMetadata().GetFirstDiscoveredAt()), - obj.GetStatus(), - pgutils.NilOrUUID(obj.GetSourceId()), - protocompat.NilOrTime(obj.GetLastUpdatedAt()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"discovered_clusters"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetMetadata().GetName(), + obj.GetMetadata().GetType(), + protocompat.NilOrTime(obj.GetMetadata().GetFirstDiscoveredAt()), + obj.GetStatus(), + pgutils.NilOrUUID(obj.GetSourceId()), + protocompat.NilOrTime(obj.GetLastUpdatedAt()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"discovered_clusters"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/externalbackups/internal/store/postgres/store.go b/central/externalbackups/internal/store/postgres/store.go index 9f3ddb17ac5c0..44638919da61e 100644 --- a/central/externalbackups/internal/store/postgres/store.go +++ b/central/externalbackups/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromExternalBackups(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"external_backups"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"external_backups"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/group/datastore/internal/store/postgres/store.go b/central/group/datastore/internal/store/postgres/store.go index 142fa58ef3ba2..1a3ce8555440f 100644 --- a/central/group/datastore/internal/store/postgres/store.go +++ b/central/group/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -116,12 +115,18 @@ func copyFromGroups(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, ob if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetProps().GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "props_id", @@ -132,41 +137,31 @@ func copyFromGroups(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, ob "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetProps().GetId(), - obj.GetProps().GetAuthProviderId(), - obj.GetProps().GetKey(), - obj.GetProps().GetValue(), - obj.GetRoleName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetProps().GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"groups"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetProps().GetId(), + obj.GetProps().GetAuthProviderId(), + obj.GetProps().GetKey(), + obj.GetProps().GetValue(), + obj.GetRoleName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"groups"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/hash/datastore/store/postgres/store.go b/central/hash/datastore/store/postgres/store.go index e622f08034b39..5f0a754de48e1 100644 --- a/central/hash/datastore/store/postgres/store.go +++ b/central/hash/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromHashes(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, ob if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetClusterId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "clusterid", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetClusterId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetClusterId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"hashes"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetClusterId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"hashes"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/imagecomponent/datastore/store/postgres/store.go b/central/imagecomponent/datastore/store/postgres/store.go index b8cc70f165e79..80fda7b26b900 100644 --- a/central/imagecomponent/datastore/store/postgres/store.go +++ b/central/imagecomponent/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -122,12 +121,18 @@ func copyFromImageComponents(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -141,44 +146,34 @@ func copyFromImageComponents(ctx context.Context, s pgSearch.Deleter, tx *postgr "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetVersion(), - obj.GetPriority(), - obj.GetSource(), - obj.GetRiskScore(), - obj.GetTopCvss(), - obj.GetOperatingSystem(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_components"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetVersion(), + obj.GetPriority(), + obj.GetSource(), + obj.GetRiskScore(), + obj.GetTopCvss(), + obj.GetOperatingSystem(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_components"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/imagecomponent/v2/datastore/store/postgres/store.go b/central/imagecomponent/v2/datastore/store/postgres/store.go index 6edc3f209acf7..c03ab6c711dee 100644 --- a/central/imagecomponent/v2/datastore/store/postgres/store.go +++ b/central/imagecomponent/v2/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -126,12 +125,18 @@ func copyFromImageComponentV2(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -148,47 +153,37 @@ func copyFromImageComponentV2(ctx context.Context, s pgSearch.Deleter, tx *postg "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetVersion(), - obj.GetPriority(), - obj.GetSource(), - obj.GetRiskScore(), - obj.GetTopCvss(), - obj.GetOperatingSystem(), - obj.GetImageId(), - obj.GetLocation(), - pgutils.NilOrString(obj.GetImageIdV2()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_component_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetVersion(), + obj.GetPriority(), + obj.GetSource(), + obj.GetRiskScore(), + obj.GetTopCvss(), + obj.GetOperatingSystem(), + obj.GetImageId(), + obj.GetLocation(), + pgutils.NilOrString(obj.GetImageIdV2()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_component_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/imageintegration/store/postgres/store.go b/central/imageintegration/store/postgres/store.go index 2f01f28580d2f..748dad580babe 100644 --- a/central/imageintegration/store/postgres/store.go +++ b/central/imageintegration/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -126,12 +125,18 @@ func copyFromImageIntegrations(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -140,39 +145,29 @@ func copyFromImageIntegrations(ctx context.Context, s pgSearch.Deleter, tx *post "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - pgutils.NilOrUUID(obj.GetClusterId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_integrations"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + pgutils.NilOrUUID(obj.GetClusterId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"image_integrations"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/integrationhealth/store/postgres/store.go b/central/integrationhealth/store/postgres/store.go index 1d13e064e4a4d..a0b855b125318 100644 --- a/central/integrationhealth/store/postgres/store.go +++ b/central/integrationhealth/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromIntegrationHealths(ctx context.Context, s pgSearch.Deleter, tx *pos if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"integration_healths"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"integration_healths"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/logimbue/store/postgres/store.go b/central/logimbue/store/postgres/store.go index c23972b772428..fb214a4c2fcee 100644 --- a/central/logimbue/store/postgres/store.go +++ b/central/logimbue/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -114,12 +113,18 @@ func copyFromLogImbues(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -127,38 +132,28 @@ func copyFromLogImbues(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - protocompat.NilOrTime(obj.GetTimestamp()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"log_imbues"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + protocompat.NilOrTime(obj.GetTimestamp()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"log_imbues"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/namespace/datastore/internal/store/postgres/store.go b/central/namespace/datastore/internal/store/postgres/store.go index a6d7a9b45b526..88f6597b0f6ca 100644 --- a/central/namespace/datastore/internal/store/postgres/store.go +++ b/central/namespace/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -153,12 +152,18 @@ func copyFromNamespaces(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -170,42 +175,32 @@ func copyFromNamespaces(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetClusterName(), - pgutils.EmptyOrMap(obj.GetLabels()), - pgutils.EmptyOrMap(obj.GetAnnotations()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"namespaces"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetClusterName(), + pgutils.EmptyOrMap(obj.GetLabels()), + pgutils.EmptyOrMap(obj.GetAnnotations()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"namespaces"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/networkbaseline/store/postgres/store.go b/central/networkbaseline/store/postgres/store.go index fad41d7b0558e..e5c83ef083aad 100644 --- a/central/networkbaseline/store/postgres/store.go +++ b/central/networkbaseline/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -139,12 +138,18 @@ func copyFromNetworkBaselines(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetDeploymentId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "deploymentid", @@ -153,39 +158,29 @@ func copyFromNetworkBaselines(ctx context.Context, s pgSearch.Deleter, tx *postg "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetDeploymentId()), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetNamespace(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetDeploymentId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"network_baselines"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetDeploymentId()), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetNamespace(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"network_baselines"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/networkgraph/config/datastore/internal/store/postgres/store.go b/central/networkgraph/config/datastore/internal/store/postgres/store.go index 7e8e3287ea481..d9b814b8bcfe3 100644 --- a/central/networkgraph/config/datastore/internal/store/postgres/store.go +++ b/central/networkgraph/config/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromNetworkGraphConfigs(ctx context.Context, s pgSearch.Deleter, tx *po if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"network_graph_configs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"network_graph_configs"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/networkgraph/entity/datastore/internal/store/postgres/store.go b/central/networkgraph/entity/datastore/internal/store/postgres/store.go index f3e57ae1bbb44..e1c5ad3e26f3b 100644 --- a/central/networkgraph/entity/datastore/internal/store/postgres/store.go +++ b/central/networkgraph/entity/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -119,12 +118,18 @@ func copyFromNetworkEntities(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetInfo().GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "info_id", @@ -134,40 +139,30 @@ func copyFromNetworkEntities(ctx context.Context, s pgSearch.Deleter, tx *postgr "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetInfo().GetId(), - pgutils.NilOrCIDR(obj.GetInfo().GetExternalSource().GetCidr()), - obj.GetInfo().GetExternalSource().GetDefault(), - obj.GetInfo().GetExternalSource().GetDiscovered(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetInfo().GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"network_entities"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetInfo().GetId(), + pgutils.NilOrCIDR(obj.GetInfo().GetExternalSource().GetCidr()), + obj.GetInfo().GetExternalSource().GetDefault(), + obj.GetInfo().GetExternalSource().GetDiscovered(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"network_entities"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/networkpolicies/datastore/internal/store/postgres/store.go b/central/networkpolicies/datastore/internal/store/postgres/store.go index 12bfb97a63b54..4eb5d2515e65a 100644 --- a/central/networkpolicies/datastore/internal/store/postgres/store.go +++ b/central/networkpolicies/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -139,12 +138,18 @@ func copyFromNetworkpolicies(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -153,39 +158,29 @@ func copyFromNetworkpolicies(ctx context.Context, s pgSearch.Deleter, tx *postgr "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetNamespace(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"networkpolicies"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetNamespace(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"networkpolicies"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/networkpolicies/datastore/internal/undodeploymentstore/postgres/store.go b/central/networkpolicies/datastore/internal/undodeploymentstore/postgres/store.go index 7d6ac4d6626bc..37de03d4f30a2 100644 --- a/central/networkpolicies/datastore/internal/undodeploymentstore/postgres/store.go +++ b/central/networkpolicies/datastore/internal/undodeploymentstore/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -113,49 +112,45 @@ func copyFromNetworkpoliciesundodeployments(ctx context.Context, s pgSearch.Dele if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetDeploymentId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "deploymentid", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetDeploymentId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetDeploymentId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"networkpoliciesundodeployments"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetDeploymentId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"networkpoliciesundodeployments"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/networkpolicies/datastore/internal/undostore/postgres/store.go b/central/networkpolicies/datastore/internal/undostore/postgres/store.go index f2b622b403add..4dfcdf044174b 100644 --- a/central/networkpolicies/datastore/internal/undostore/postgres/store.go +++ b/central/networkpolicies/datastore/internal/undostore/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -113,49 +112,45 @@ func copyFromNetworkpolicyapplicationundorecords(ctx context.Context, s pgSearch if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetClusterId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "clusterid", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetClusterId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetClusterId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"networkpolicyapplicationundorecords"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetClusterId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"networkpolicyapplicationundorecords"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/nodecomponent/datastore/store/postgres/store.go b/central/nodecomponent/datastore/store/postgres/store.go index 930de90ee9a5c..2e91cf99bc4dc 100644 --- a/central/nodecomponent/datastore/store/postgres/store.go +++ b/central/nodecomponent/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -121,12 +120,18 @@ func copyFromNodeComponents(ctx context.Context, s pgSearch.Deleter, tx *postgre if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -139,43 +144,33 @@ func copyFromNodeComponents(ctx context.Context, s pgSearch.Deleter, tx *postgre "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetVersion(), - obj.GetPriority(), - obj.GetRiskScore(), - obj.GetTopCvss(), - obj.GetOperatingSystem(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"node_components"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetVersion(), + obj.GetPriority(), + obj.GetRiskScore(), + obj.GetTopCvss(), + obj.GetOperatingSystem(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"node_components"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/notifier/datastore/internal/store/postgres/store.go b/central/notifier/datastore/internal/store/postgres/store.go index 97e46cb535a3c..27d27284bd25b 100644 --- a/central/notifier/datastore/internal/store/postgres/store.go +++ b/central/notifier/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -113,12 +112,18 @@ func copyFromNotifiers(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -126,38 +131,28 @@ func copyFromNotifiers(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"notifiers"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"notifiers"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/pod/datastore/internal/store/postgres/store.go b/central/pod/datastore/internal/store/postgres/store.go index 4271ed1b8b77e..cc78f37b85eb1 100644 --- a/central/pod/datastore/internal/store/postgres/store.go +++ b/central/pod/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -175,12 +174,18 @@ func copyFromPods(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, objs if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -191,41 +196,31 @@ func copyFromPods(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, objs "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - pgutils.NilOrUUID(obj.GetDeploymentId()), - obj.GetNamespace(), - pgutils.NilOrUUID(obj.GetClusterId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"pods"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + pgutils.NilOrUUID(obj.GetDeploymentId()), + obj.GetNamespace(), + pgutils.NilOrUUID(obj.GetClusterId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"pods"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -241,8 +236,6 @@ func copyFromPodsLiveInstances(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "pods_id", @@ -251,26 +244,22 @@ func copyFromPodsLiveInstances(ctx context.Context, s pgSearch.Deleter, tx *post } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(podID), - idx, - obj.GetImageDigest(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"pods_live_instances"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(podID), + idx, + obj.GetImageDigest(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"pods_live_instances"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/policy/store/postgres/store.go b/central/policy/store/postgres/store.go index fe840a58db6fe..3ce1a752b62a7 100644 --- a/central/policy/store/postgres/store.go +++ b/central/policy/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -127,12 +126,18 @@ func copyFromPolicies(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -150,48 +155,38 @@ func copyFromPolicies(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetDescription(), - obj.GetDisabled(), - obj.GetCategories(), - obj.GetLifecycleStages(), - obj.GetSeverity(), - obj.GetEnforcementActions(), - protocompat.NilOrTime(obj.GetLastUpdated()), - obj.GetSORTName(), - obj.GetSORTLifecycleStage(), - obj.GetSORTEnforcement(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"policies"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetDescription(), + obj.GetDisabled(), + obj.GetCategories(), + obj.GetLifecycleStages(), + obj.GetSeverity(), + obj.GetEnforcementActions(), + protocompat.NilOrTime(obj.GetLastUpdated()), + obj.GetSORTName(), + obj.GetSORTLifecycleStage(), + obj.GetSORTEnforcement(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"policies"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/policycategory/store/postgres/store.go b/central/policycategory/store/postgres/store.go index 0e92c0ab564a2..252deeffcf525 100644 --- a/central/policycategory/store/postgres/store.go +++ b/central/policycategory/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -116,12 +115,18 @@ func copyFromPolicyCategories(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -129,38 +134,28 @@ func copyFromPolicyCategories(ctx context.Context, s pgSearch.Deleter, tx *postg "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"policy_categories"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"policy_categories"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/policycategoryedge/store/postgres/store.go b/central/policycategoryedge/store/postgres/store.go index 433bb4df23146..0ed5a89f8400a 100644 --- a/central/policycategoryedge/store/postgres/store.go +++ b/central/policycategoryedge/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -117,12 +116,18 @@ func copyFromPolicyCategoryEdges(ctx context.Context, s pgSearch.Deleter, tx *po if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -131,39 +136,29 @@ func copyFromPolicyCategoryEdges(ctx context.Context, s pgSearch.Deleter, tx *po "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetPolicyId(), - obj.GetCategoryId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"policy_category_edges"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetPolicyId(), + obj.GetCategoryId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"policy_category_edges"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/processbaseline/store/postgres/store.go b/central/processbaseline/store/postgres/store.go index e55015544a218..118d0695f6e8a 100644 --- a/central/processbaseline/store/postgres/store.go +++ b/central/processbaseline/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -149,12 +148,18 @@ func copyFromProcessBaselines(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -164,40 +169,30 @@ func copyFromProcessBaselines(ctx context.Context, s pgSearch.Deleter, tx *postg "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - pgutils.NilOrUUID(obj.GetKey().GetDeploymentId()), - pgutils.NilOrUUID(obj.GetKey().GetClusterId()), - obj.GetKey().GetNamespace(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"process_baselines"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + pgutils.NilOrUUID(obj.GetKey().GetDeploymentId()), + pgutils.NilOrUUID(obj.GetKey().GetClusterId()), + obj.GetKey().GetNamespace(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"process_baselines"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/processbaselineresults/datastore/internal/store/postgres/store.go b/central/processbaselineresults/datastore/internal/store/postgres/store.go index 212bf26cca89a..b2a47e07705cd 100644 --- a/central/processbaselineresults/datastore/internal/store/postgres/store.go +++ b/central/processbaselineresults/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -139,12 +138,18 @@ func copyFromProcessBaselineResults(ctx context.Context, s pgSearch.Deleter, tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetDeploymentId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "deploymentid", @@ -153,39 +158,29 @@ func copyFromProcessBaselineResults(ctx context.Context, s pgSearch.Deleter, tx "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetDeploymentId()), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetNamespace(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetDeploymentId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"process_baseline_results"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetDeploymentId()), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetNamespace(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"process_baseline_results"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/processindicator/store/postgres/store.go b/central/processindicator/store/postgres/store.go index 62a0affa77c3d..03f74fb14e609 100644 --- a/central/processindicator/store/postgres/store.go +++ b/central/processindicator/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -151,12 +150,18 @@ func copyFromProcessIndicators(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -176,50 +181,40 @@ func copyFromProcessIndicators(ctx context.Context, s pgSearch.Deleter, tx *post "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - pgutils.NilOrUUID(obj.GetDeploymentId()), - obj.GetContainerName(), - obj.GetPodId(), - pgutils.NilOrUUID(obj.GetPodUid()), - obj.GetSignal().GetContainerId(), - protocompat.NilOrTime(obj.GetSignal().GetTime()), - obj.GetSignal().GetName(), - obj.GetSignal().GetArgs(), - obj.GetSignal().GetExecFilePath(), - obj.GetSignal().GetUid(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetNamespace(), - protocompat.NilOrTime(obj.GetContainerStartTime()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"process_indicators"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + pgutils.NilOrUUID(obj.GetDeploymentId()), + obj.GetContainerName(), + obj.GetPodId(), + pgutils.NilOrUUID(obj.GetPodUid()), + obj.GetSignal().GetContainerId(), + protocompat.NilOrTime(obj.GetSignal().GetTime()), + obj.GetSignal().GetName(), + obj.GetSignal().GetArgs(), + obj.GetSignal().GetExecFilePath(), + obj.GetSignal().GetUid(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetNamespace(), + protocompat.NilOrTime(obj.GetContainerStartTime()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"process_indicators"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/processlisteningonport/store/postgres/store.go b/central/processlisteningonport/store/postgres/store.go index d55bf401391ca..5778b46610d1f 100644 --- a/central/processlisteningonport/store/postgres/store.go +++ b/central/processlisteningonport/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -147,12 +146,18 @@ func copyFromListeningEndpoints(ctx context.Context, s pgSearch.Deleter, tx *pos if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -168,46 +173,36 @@ func copyFromListeningEndpoints(ctx context.Context, s pgSearch.Deleter, tx *pos "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetPort(), - obj.GetProtocol(), - protocompat.NilOrTime(obj.GetCloseTimestamp()), - pgutils.NilOrUUID(obj.GetProcessIndicatorId()), - obj.GetClosed(), - pgutils.NilOrUUID(obj.GetDeploymentId()), - pgutils.NilOrUUID(obj.GetPodUid()), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetNamespace(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"listening_endpoints"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetPort(), + obj.GetProtocol(), + protocompat.NilOrTime(obj.GetCloseTimestamp()), + pgutils.NilOrUUID(obj.GetProcessIndicatorId()), + obj.GetClosed(), + pgutils.NilOrUUID(obj.GetDeploymentId()), + pgutils.NilOrUUID(obj.GetPodUid()), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetNamespace(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"listening_endpoints"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/rbac/k8srole/internal/store/postgres/store.go b/central/rbac/k8srole/internal/store/postgres/store.go index 51551c28e0a1c..7c89b5a754368 100644 --- a/central/rbac/k8srole/internal/store/postgres/store.go +++ b/central/rbac/k8srole/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -144,12 +143,18 @@ func copyFromK8sRoles(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -163,44 +168,34 @@ func copyFromK8sRoles(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - obj.GetNamespace(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetClusterName(), - obj.GetClusterRole(), - pgutils.EmptyOrMap(obj.GetLabels()), - pgutils.EmptyOrMap(obj.GetAnnotations()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"k8s_roles"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + obj.GetNamespace(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetClusterName(), + obj.GetClusterRole(), + pgutils.EmptyOrMap(obj.GetLabels()), + pgutils.EmptyOrMap(obj.GetAnnotations()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"k8s_roles"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/rbac/k8srolebinding/internal/store/postgres/store.go b/central/rbac/k8srolebinding/internal/store/postgres/store.go index b384766110a21..62b3e3218ad94 100644 --- a/central/rbac/k8srolebinding/internal/store/postgres/store.go +++ b/central/rbac/k8srolebinding/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -171,12 +170,18 @@ func copyFromRoleBindings(ctx context.Context, s pgSearch.Deleter, tx *postgres. if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -191,45 +196,35 @@ func copyFromRoleBindings(ctx context.Context, s pgSearch.Deleter, tx *postgres. "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - obj.GetNamespace(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetClusterName(), - obj.GetClusterRole(), - pgutils.EmptyOrMap(obj.GetLabels()), - pgutils.EmptyOrMap(obj.GetAnnotations()), - pgutils.NilOrUUID(obj.GetRoleId()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"role_bindings"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + obj.GetNamespace(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetClusterName(), + obj.GetClusterRole(), + pgutils.EmptyOrMap(obj.GetLabels()), + pgutils.EmptyOrMap(obj.GetAnnotations()), + pgutils.NilOrUUID(obj.GetRoleId()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"role_bindings"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -245,8 +240,6 @@ func copyFromRoleBindingsSubjects(ctx context.Context, s pgSearch.Deleter, tx *p if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "role_bindings_id", @@ -256,27 +249,23 @@ func copyFromRoleBindingsSubjects(ctx context.Context, s pgSearch.Deleter, tx *p } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(roleBindingID), - idx, - obj.GetKind(), - obj.GetName(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"role_bindings_subjects"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(roleBindingID), + idx, + obj.GetKind(), + obj.GetName(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"role_bindings_subjects"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/reports/config/store/postgres/store.go b/central/reports/config/store/postgres/store.go index d9c8dc3938c32..150d8439b7130 100644 --- a/central/reports/config/store/postgres/store.go +++ b/central/reports/config/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -145,12 +144,18 @@ func copyFromReportConfigurations(ctx context.Context, s pgSearch.Deleter, tx *p if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -162,42 +167,32 @@ func copyFromReportConfigurations(ctx context.Context, s pgSearch.Deleter, tx *p "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetType(), - obj.GetScopeId(), - obj.GetResourceScope().GetCollectionId(), - obj.GetCreator().GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"report_configurations"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetType(), + obj.GetScopeId(), + obj.GetResourceScope().GetCollectionId(), + obj.GetCreator().GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"report_configurations"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -213,8 +208,6 @@ func copyFromReportConfigurationsNotifiers(ctx context.Context, s pgSearch.Delet if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "report_configurations_id", @@ -223,26 +216,22 @@ func copyFromReportConfigurationsNotifiers(ctx context.Context, s pgSearch.Delet } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - reportConfigurationID, - idx, - obj.GetId(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"report_configurations_notifiers"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + reportConfigurationID, + idx, + obj.GetId(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"report_configurations_notifiers"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/reports/snapshot/datastore/store/postgres/store.go b/central/reports/snapshot/datastore/store/postgres/store.go index ff9e90c428d70..bc4642dfe2598 100644 --- a/central/reports/snapshot/datastore/store/postgres/store.go +++ b/central/reports/snapshot/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -127,12 +126,18 @@ func copyFromReportSnapshots(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetReportId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "reportid", @@ -149,47 +154,37 @@ func copyFromReportSnapshots(ctx context.Context, s pgSearch.Deleter, tx *postgr "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetReportId()), - pgutils.NilOrString(obj.GetReportConfigurationId()), - obj.GetName(), - obj.GetReportStatus().GetRunState(), - protocompat.NilOrTime(obj.GetReportStatus().GetQueuedAt()), - protocompat.NilOrTime(obj.GetReportStatus().GetCompletedAt()), - obj.GetReportStatus().GetReportRequestType(), - obj.GetReportStatus().GetReportNotificationMethod(), - obj.GetRequester().GetId(), - obj.GetRequester().GetName(), - obj.GetAreaOfConcern(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetReportId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"report_snapshots"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetReportId()), + pgutils.NilOrString(obj.GetReportConfigurationId()), + obj.GetName(), + obj.GetReportStatus().GetRunState(), + protocompat.NilOrTime(obj.GetReportStatus().GetQueuedAt()), + protocompat.NilOrTime(obj.GetReportStatus().GetCompletedAt()), + obj.GetReportStatus().GetReportRequestType(), + obj.GetReportStatus().GetReportNotificationMethod(), + obj.GetRequester().GetId(), + obj.GetRequester().GetName(), + obj.GetAreaOfConcern(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"report_snapshots"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/resourcecollection/datastore/store/postgres/store.go b/central/resourcecollection/datastore/store/postgres/store.go index fafc7104bffea..b2c1a1fdff87a 100644 --- a/central/resourcecollection/datastore/store/postgres/store.go +++ b/central/resourcecollection/datastore/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -143,12 +142,18 @@ func copyFromCollections(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -158,40 +163,30 @@ func copyFromCollections(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetCreatedBy().GetName(), - obj.GetUpdatedBy().GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"collections"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetCreatedBy().GetName(), + obj.GetUpdatedBy().GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"collections"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -207,8 +202,6 @@ func copyFromCollectionsEmbeddedCollections(ctx context.Context, s pgSearch.Dele if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "collections_id", @@ -217,26 +210,22 @@ func copyFromCollectionsEmbeddedCollections(ctx context.Context, s pgSearch.Dele } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - collectionID, - idx, - obj.GetId(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"collections_embedded_collections"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + collectionID, + idx, + obj.GetId(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"collections_embedded_collections"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/risk/datastore/internal/store/postgres/store.go b/central/risk/datastore/internal/store/postgres/store.go index 983f20874a524..e3b450a81a0ae 100644 --- a/central/risk/datastore/internal/store/postgres/store.go +++ b/central/risk/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -141,12 +140,18 @@ func copyFromRisks(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, obj if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -157,41 +162,31 @@ func copyFromRisks(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, obj "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetSubject().GetNamespace(), - pgutils.NilOrUUID(obj.GetSubject().GetClusterId()), - obj.GetSubject().GetType(), - obj.GetScore(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"risks"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetSubject().GetNamespace(), + pgutils.NilOrUUID(obj.GetSubject().GetClusterId()), + obj.GetSubject().GetType(), + obj.GetScore(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"risks"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/role/store/permissionset/postgres/store.go b/central/role/store/permissionset/postgres/store.go index 73b4b0440760c..64426a7a98091 100644 --- a/central/role/store/permissionset/postgres/store.go +++ b/central/role/store/permissionset/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -114,12 +113,18 @@ func copyFromPermissionSets(ctx context.Context, s pgSearch.Deleter, tx *postgre if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -127,38 +132,28 @@ func copyFromPermissionSets(ctx context.Context, s pgSearch.Deleter, tx *postgre "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"permission_sets"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"permission_sets"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/role/store/role/postgres/store.go b/central/role/store/role/postgres/store.go index 3b716a5405db4..dff84cef965a3 100644 --- a/central/role/store/role/postgres/store.go +++ b/central/role/store/role/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromRoles(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, obj if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetName()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "name", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetName()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"roles"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"roles"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/role/store/simpleaccessscope/postgres/store.go b/central/role/store/simpleaccessscope/postgres/store.go index f386f3044f314..27d430e7cbe02 100644 --- a/central/role/store/simpleaccessscope/postgres/store.go +++ b/central/role/store/simpleaccessscope/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -114,12 +113,18 @@ func copyFromSimpleAccessScopes(ctx context.Context, s pgSearch.Deleter, tx *pos if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -127,38 +132,28 @@ func copyFromSimpleAccessScopes(ctx context.Context, s pgSearch.Deleter, tx *pos "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"simple_access_scopes"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"simple_access_scopes"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/secret/internal/store/postgres/store.go b/central/secret/internal/store/postgres/store.go index e1bed983c5ffc..d448c518886fe 100644 --- a/central/secret/internal/store/postgres/store.go +++ b/central/secret/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -195,12 +194,18 @@ func copyFromSecrets(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, o if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -212,42 +217,32 @@ func copyFromSecrets(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, o "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetClusterName(), - obj.GetNamespace(), - protocompat.NilOrTime(obj.GetCreatedAt()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secrets"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetClusterName(), + obj.GetNamespace(), + protocompat.NilOrTime(obj.GetCreatedAt()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secrets"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -263,8 +258,6 @@ func copyFromSecretsFiles(ctx context.Context, s pgSearch.Deleter, tx *postgres. if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "secrets_id", @@ -274,30 +267,26 @@ func copyFromSecretsFiles(ctx context.Context, s pgSearch.Deleter, tx *postgres. } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(secretID), - idx, - obj.GetType(), - protocompat.NilOrTime(obj.GetCert().GetEndDate()), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secrets_files"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(secretID), + idx, + obj.GetType(), + protocompat.NilOrTime(obj.GetCert().GetEndDate()), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secrets_files"}, copyCols, inputRows); err != nil { + return err } - for idx, obj := range objs { + for _, obj := range objs { if err := copyFromSecretsFilesRegistries(ctx, s, tx, secretID, idx, obj.GetImagePullSecret().GetRegistries()...); err != nil { return err } @@ -310,8 +299,6 @@ func copyFromSecretsFilesRegistries(ctx context.Context, s pgSearch.Deleter, tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "secrets_id", @@ -321,27 +308,23 @@ func copyFromSecretsFilesRegistries(ctx context.Context, s pgSearch.Deleter, tx } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(secretID), - secretFileIdx, - idx, - obj.GetName(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secrets_files_registries"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + pgutils.NilOrUUID(secretID), + secretFileIdx, + idx, + obj.GetName(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"secrets_files_registries"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/serviceaccount/internal/store/postgres/store.go b/central/serviceaccount/internal/store/postgres/store.go index 51fca7586cc72..1a291bc221918 100644 --- a/central/serviceaccount/internal/store/postgres/store.go +++ b/central/serviceaccount/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -143,12 +142,18 @@ func copyFromServiceAccounts(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -161,43 +166,33 @@ func copyFromServiceAccounts(ctx context.Context, s pgSearch.Deleter, tx *postgr "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetName(), - obj.GetNamespace(), - obj.GetClusterName(), - pgutils.NilOrUUID(obj.GetClusterId()), - pgutils.EmptyOrMap(obj.GetLabels()), - pgutils.EmptyOrMap(obj.GetAnnotations()), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"service_accounts"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetName(), + obj.GetNamespace(), + obj.GetClusterName(), + pgutils.NilOrUUID(obj.GetClusterId()), + pgutils.EmptyOrMap(obj.GetLabels()), + pgutils.EmptyOrMap(obj.GetAnnotations()), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"service_accounts"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/serviceidentities/internal/store/postgres/store.go b/central/serviceidentities/internal/store/postgres/store.go index 320b9771e554a..a0b2b1a5156fa 100644 --- a/central/serviceidentities/internal/store/postgres/store.go +++ b/central/serviceidentities/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromServiceIdentities(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetSerialStr()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "serialstr", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetSerialStr(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetSerialStr()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"service_identities"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetSerialStr(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"service_identities"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/signatureintegration/store/postgres/store.go b/central/signatureintegration/store/postgres/store.go index fcd684013f93e..bf4fa9f50e1f3 100644 --- a/central/signatureintegration/store/postgres/store.go +++ b/central/signatureintegration/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -121,12 +120,18 @@ func copyFromSignatureIntegrations(ctx context.Context, s pgSearch.Deleter, tx * if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -134,38 +139,28 @@ func copyFromSignatureIntegrations(ctx context.Context, s pgSearch.Deleter, tx * "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"signature_integrations"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"signature_integrations"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/virtualmachine/datastore/internal/store/postgres/store.go b/central/virtualmachine/datastore/internal/store/postgres/store.go index 58921585c9ce1..dbd17b14026a6 100644 --- a/central/virtualmachine/datastore/internal/store/postgres/store.go +++ b/central/virtualmachine/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "strings" "time" @@ -141,12 +140,18 @@ func copyFromVirtualMachines(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -157,41 +162,31 @@ func copyFromVirtualMachines(ctx context.Context, s pgSearch.Deleter, tx *postgr "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetNamespace(), - obj.GetName(), - pgutils.NilOrUUID(obj.GetClusterId()), - obj.GetClusterName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"virtual_machines"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetNamespace(), + obj.GetName(), + pgutils.NilOrUUID(obj.GetClusterId()), + obj.GetClusterName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"virtual_machines"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/vulnmgmt/vulnerabilityrequest/datastore/internal/store/postgres/store.go b/central/vulnmgmt/vulnerabilityrequest/datastore/internal/store/postgres/store.go index 65b091bd770bd..8a30674e7d89b 100644 --- a/central/vulnmgmt/vulnerabilityrequest/datastore/internal/store/postgres/store.go +++ b/central/vulnmgmt/vulnerabilityrequest/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -206,12 +205,18 @@ func copyFromVulnerabilityRequests(ctx context.Context, s pgSearch.Deleter, tx * if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -236,55 +241,45 @@ func copyFromVulnerabilityRequests(ctx context.Context, s pgSearch.Deleter, tx * "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetName(), - obj.GetTargetState(), - obj.GetStatus(), - obj.GetExpired(), - obj.GetRequestor().GetName(), - protocompat.NilOrTime(obj.GetCreatedAt()), - protocompat.NilOrTime(obj.GetLastUpdated()), - obj.GetScope().GetImageScope().GetRegistry(), - obj.GetScope().GetImageScope().GetRemote(), - obj.GetScope().GetImageScope().GetTag(), - obj.GetRequesterV2().GetId(), - obj.GetRequesterV2().GetName(), - protocompat.NilOrTime(obj.GetDeferralReq().GetExpiry().GetExpiresOn()), - obj.GetDeferralReq().GetExpiry().GetExpiresWhenFixed(), - obj.GetDeferralReq().GetExpiry().GetExpiryType(), - obj.GetCves().GetCves(), - obj.GetDeferralUpdate().GetCVEs(), - obj.GetFalsePositiveUpdate().GetCVEs(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetName(), + obj.GetTargetState(), + obj.GetStatus(), + obj.GetExpired(), + obj.GetRequestor().GetName(), + protocompat.NilOrTime(obj.GetCreatedAt()), + protocompat.NilOrTime(obj.GetLastUpdated()), + obj.GetScope().GetImageScope().GetRegistry(), + obj.GetScope().GetImageScope().GetRemote(), + obj.GetScope().GetImageScope().GetTag(), + obj.GetRequesterV2().GetId(), + obj.GetRequesterV2().GetName(), + protocompat.NilOrTime(obj.GetDeferralReq().GetExpiry().GetExpiresOn()), + obj.GetDeferralReq().GetExpiry().GetExpiresWhenFixed(), + obj.GetDeferralReq().GetExpiry().GetExpiryType(), + obj.GetCves().GetCves(), + obj.GetDeferralUpdate().GetCVEs(), + obj.GetFalsePositiveUpdate().GetCVEs(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -306,8 +301,6 @@ func copyFromVulnerabilityRequestsApprovers(ctx context.Context, s pgSearch.Dele if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "vulnerability_requests_id", @@ -316,26 +309,22 @@ func copyFromVulnerabilityRequestsApprovers(ctx context.Context, s pgSearch.Dele } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - vulnerabilityRequestID, - idx, - obj.GetName(), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests_approvers"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + vulnerabilityRequestID, + idx, + obj.GetName(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests_approvers"}, copyCols, inputRows); err != nil { + return err } return nil @@ -345,8 +334,6 @@ func copyFromVulnerabilityRequestsComments(ctx context.Context, s pgSearch.Delet if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "vulnerability_requests_id", @@ -355,26 +342,22 @@ func copyFromVulnerabilityRequestsComments(ctx context.Context, s pgSearch.Delet } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - vulnerabilityRequestID, - idx, - obj.GetUser().GetName(), - }) - - idx++ + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests_comments"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + vulnerabilityRequestID, + idx, + obj.GetUser().GetName(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests_comments"}, copyCols, inputRows); err != nil { + return err } return nil @@ -384,8 +367,6 @@ func copyFromVulnerabilityRequestsApproversV2(ctx context.Context, s pgSearch.De if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "vulnerability_requests_id", @@ -395,27 +376,23 @@ func copyFromVulnerabilityRequestsApproversV2(ctx context.Context, s pgSearch.De } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - vulnerabilityRequestID, - idx, - obj.GetId(), - obj.GetName(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests_approvers_v2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + vulnerabilityRequestID, + idx, + obj.GetId(), + obj.GetName(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"vulnerability_requests_approvers_v2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/central/watchedimage/datastore/internal/store/postgres/store.go b/central/watchedimage/datastore/internal/store/postgres/store.go index ab3124c28e755..b237c362d5c41 100644 --- a/central/watchedimage/datastore/internal/store/postgres/store.go +++ b/central/watchedimage/datastore/internal/store/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -112,49 +111,45 @@ func copyFromWatchedImages(ctx context.Context, s pgSearch.Deleter, tx *postgres if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetName()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "name", "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetName(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetName()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"watched_images"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetName(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"watched_images"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/multitest/postgres/store.go b/tools/generate-helpers/pg-table-bindings/multitest/postgres/store.go index 79ad7981934b2..8a938fa1f0289 100644 --- a/tools/generate-helpers/pg-table-bindings/multitest/postgres/store.go +++ b/tools/generate-helpers/pg-table-bindings/multitest/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -160,12 +159,18 @@ func copyFromTestStructs(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetKey1()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "key1", @@ -185,50 +190,40 @@ func copyFromTestStructs(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetKey1(), - obj.GetKey2(), - obj.GetStringSlice(), - obj.GetBool(), - obj.GetUint64(), - obj.GetInt64(), - obj.GetFloat(), - pgutils.EmptyOrMap(obj.GetLabels()), - protocompat.NilOrTime(obj.GetTimestamp()), - obj.GetEnum(), - obj.GetEnums(), - obj.GetString_(), - obj.GetInt32Slice(), - obj.GetOneofnested().GetNested(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetKey1()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_structs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetKey1(), + obj.GetKey2(), + obj.GetStringSlice(), + obj.GetBool(), + obj.GetUint64(), + obj.GetInt64(), + obj.GetFloat(), + pgutils.EmptyOrMap(obj.GetLabels()), + protocompat.NilOrTime(obj.GetTimestamp()), + obj.GetEnum(), + obj.GetEnums(), + obj.GetString_(), + obj.GetInt32Slice(), + obj.GetOneofnested().GetNested(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_structs"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -244,8 +239,6 @@ func copyFromTestStructsNesteds(ctx context.Context, s pgSearch.Deleter, tx *pos if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "test_structs_key1", @@ -259,31 +252,27 @@ func copyFromTestStructsNesteds(ctx context.Context, s pgSearch.Deleter, tx *pos } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - testStructKey1, - idx, - obj.GetNested(), - obj.GetIsNested(), - obj.GetInt64(), - obj.GetNested2().GetNested2(), - obj.GetNested2().GetIsNested(), - obj.GetNested2().GetInt64(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_structs_nesteds"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + testStructKey1, + idx, + obj.GetNested(), + obj.GetIsNested(), + obj.GetInt64(), + obj.GetNested2().GetNested2(), + obj.GetNested2().GetIsNested(), + obj.GetNested2().GetInt64(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_structs_nesteds"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/store.go.tpl b/tools/generate-helpers/pg-table-bindings/store.go.tpl index 86ce988c40019..1995839b12abf 100644 --- a/tools/generate-helpers/pg-table-bindings/store.go.tpl +++ b/tools/generate-helpers/pg-table-bindings/store.go.tpl @@ -259,16 +259,23 @@ func {{ template "insertFunctionName" $schema }}(batch *pgx.Batch, obj {{$schema {{- define "copyObject"}} {{- $schema := .schema }} +{{- $singlePK := index $schema.PrimaryKeys 0 }} func {{ template "copyFunctionName" $schema }}(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx, {{ range $index, $field := $schema.FieldsReferringToParent }} {{$field.Name}} {{$field.Type}},{{end}} objs ...{{$schema.Type}}) error { if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) {{if not $schema.Parent }} - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, {{ $singlePK.Getter "obj" }}) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } {{end}} copyCols := []string { @@ -277,47 +284,32 @@ func {{ template "copyFunctionName" $schema }}(ctx context.Context, s pgSearch.D {{- end }} } - {{ $idx := false }}{{ range $field := $schema.DBColumnFields }}{{if eq "idx" ($field.Getter "obj") -}}{{ $idx = true }}{{end}}{{ end -}} - {{ if $idx }}idx := 0{{ end }} - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - {{/* If embedded, the top-level has the full serialized object */}} - {{if not $schema.Parent }} - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - {{end}} - - inputRows = append(inputRows, []interface{}{ - {{- template "insertValues" $schema }} - }) - - {{ if not $schema.Parent }} - // Add the ID to be deleted. - deletes = append(deletes, {{ range $field := $schema.PrimaryKeys }}{{$field.Getter "obj"}}, {{end}}) - {{end}} - {{- if $idx }}idx++{{ end -}} + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent {{if not $schema.Parent }} - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] {{end}} - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"{{$schema.Table|lowerCase}}"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { + + return []interface{}{ + {{- template "insertValues" $schema }} + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"{{$schema.Table|lowerCase}}"}, copyCols, inputRows); err != nil { return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] } {{if $schema.Children }} - for {{ if $idx }}idx{{else}}_{{ end }}, obj := range objs { + for _, obj := range objs { {{- range $child := $schema.Children }} if err := {{ template "copyFunctionName" $child }}(ctx, s, tx{{ range $index, $field := $schema.PrimaryKeys }}, {{$field.Getter "obj"}}{{end}}, obj.{{$child.ObjectGetter}}...); err != nil { return err diff --git a/tools/generate-helpers/pg-table-bindings/test/postgres/store.go b/tools/generate-helpers/pg-table-bindings/test/postgres/store.go index 9b6fbae61232f..04bf15bc21ca5 100644 --- a/tools/generate-helpers/pg-table-bindings/test/postgres/store.go +++ b/tools/generate-helpers/pg-table-bindings/test/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -127,12 +126,18 @@ func copyFromTestSingleKeyStructs(ctx context.Context, s pgSearch.Deleter, tx *p if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetKey()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "key", @@ -149,47 +154,37 @@ func copyFromTestSingleKeyStructs(ctx context.Context, s pgSearch.Deleter, tx *p "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetKey(), - obj.GetName(), - obj.GetStringSlice(), - obj.GetBool(), - obj.GetUint64(), - obj.GetInt64(), - obj.GetFloat(), - pgutils.EmptyOrMap(obj.GetLabels()), - protocompat.NilOrTime(obj.GetTimestamp()), - obj.GetEnum(), - obj.GetEnums(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetKey()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_single_key_structs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetKey(), + obj.GetName(), + obj.GetStringSlice(), + obj.GetBool(), + obj.GetUint64(), + obj.GetInt64(), + obj.GetFloat(), + pgutils.EmptyOrMap(obj.GetLabels()), + protocompat.NilOrTime(obj.GetTimestamp()), + obj.GetEnum(), + obj.GetEnums(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_single_key_structs"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1/store.go index 895b1b7092a39..34e90c56eae77 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -116,12 +115,18 @@ func copyFromTestChild1(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -129,38 +134,28 @@ func copyFromTestChild1(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_child1"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_child1"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1p4/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1p4/store.go index ccd283c3978e8..c1c0d1b779da5 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1p4/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild1p4/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -118,12 +117,18 @@ func copyFromTestChild1P4(ctx context.Context, s pgSearch.Deleter, tx *postgres. if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -132,39 +137,29 @@ func copyFromTestChild1P4(ctx context.Context, s pgSearch.Deleter, tx *postgres. "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - pgutils.NilOrUUID(obj.GetParentId()), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_child1_p4"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + pgutils.NilOrUUID(obj.GetParentId()), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_child1_p4"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild2/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild2/store.go index f7f1929c6f982..d3e86cd51425a 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild2/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testchild2/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -119,12 +118,18 @@ func copyFromTestChild2(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -134,40 +139,30 @@ func copyFromTestChild2(ctx context.Context, s pgSearch.Deleter, tx *postgres.Tx "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetParentId(), - obj.GetGrandparentId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_child2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetParentId(), + obj.GetGrandparentId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_child2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testg2grandchild1/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testg2grandchild1/store.go index b20ee14f8aa21..3d1f48e3cffb1 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testg2grandchild1/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testg2grandchild1/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -118,12 +117,18 @@ func copyFromTestG2GrandChild1(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -133,40 +138,30 @@ func copyFromTestG2GrandChild1(ctx context.Context, s pgSearch.Deleter, tx *post "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetParentId(), - obj.GetChildId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_g2_grand_child1"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetParentId(), + obj.GetChildId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_g2_grand_child1"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testg3grandchild1/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testg3grandchild1/store.go index f257653079f51..29ba69e18969b 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testg3grandchild1/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testg3grandchild1/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -116,12 +115,18 @@ func copyFromTestG3GrandChild1(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -129,38 +134,28 @@ func copyFromTestG3GrandChild1(ctx context.Context, s pgSearch.Deleter, tx *post "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_g3_grand_child1"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_g3_grand_child1"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testggrandchild1/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testggrandchild1/store.go index a6da5ad99e1ab..f13ed9ba69934 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testggrandchild1/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testggrandchild1/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -116,12 +115,18 @@ func copyFromTestGGrandChild1(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -129,38 +134,28 @@ func copyFromTestGGrandChild1(ctx context.Context, s pgSearch.Deleter, tx *postg "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_g_grand_child1"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_g_grand_child1"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandchild1/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandchild1/store.go index bbda17461a2dd..796a026bd6e30 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandchild1/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandchild1/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -118,12 +117,18 @@ func copyFromTestGrandChild1(ctx context.Context, s pgSearch.Deleter, tx *postgr if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -133,40 +138,30 @@ func copyFromTestGrandChild1(ctx context.Context, s pgSearch.Deleter, tx *postgr "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetParentId(), - obj.GetChildId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grand_child1"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetParentId(), + obj.GetChildId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grand_child1"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandparent/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandparent/store.go index 2986bd01aaada..05a801f27c5c6 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandparent/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testgrandparent/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -169,12 +168,18 @@ func copyFromTestGrandparents(ctx context.Context, s pgSearch.Deleter, tx *postg if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -184,40 +189,30 @@ func copyFromTestGrandparents(ctx context.Context, s pgSearch.Deleter, tx *postg "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetVal(), - obj.GetPriority(), - obj.GetRiskScore(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grandparents"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetVal(), + obj.GetPriority(), + obj.GetRiskScore(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grandparents"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -233,8 +228,6 @@ func copyFromTestGrandparentsEmbeddeds(ctx context.Context, s pgSearch.Deleter, if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "test_grandparents_id", @@ -243,29 +236,25 @@ func copyFromTestGrandparentsEmbeddeds(ctx context.Context, s pgSearch.Deleter, } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - testGrandparentID, - idx, - obj.GetVal(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grandparents_embeddeds"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + testGrandparentID, + idx, + obj.GetVal(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grandparents_embeddeds"}, copyCols, inputRows); err != nil { + return err } - for idx, obj := range objs { + for _, obj := range objs { if err := copyFromTestGrandparentsEmbeddedsEmbedded2(ctx, s, tx, testGrandparentID, idx, obj.GetEmbedded2()...); err != nil { return err } @@ -278,8 +267,6 @@ func copyFromTestGrandparentsEmbeddedsEmbedded2(ctx context.Context, s pgSearch. if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "test_grandparents_id", @@ -289,27 +276,23 @@ func copyFromTestGrandparentsEmbeddedsEmbedded2(ctx context.Context, s pgSearch. } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - testGrandparentID, - testGrandparentEmbeddedIdx, - idx, - obj.GetVal(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grandparents_embeddeds_embedded2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + testGrandparentID, + testGrandparentEmbeddedIdx, + idx, + obj.GetVal(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_grandparents_embeddeds_embedded2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent1/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent1/store.go index 54d3a6f6600ac..a3b5b20d41d24 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent1/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent1/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -143,12 +142,18 @@ func copyFromTestParent1(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -158,40 +163,30 @@ func copyFromTestParent1(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetParentId(), - obj.GetVal(), - obj.GetStringSlice(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent1"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetParentId(), + obj.GetVal(), + obj.GetStringSlice(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent1"}, copyCols, inputRows); err != nil { + return err } for _, obj := range objs { @@ -207,8 +202,6 @@ func copyFromTestParent1Childrens(ctx context.Context, s pgSearch.Deleter, tx *p if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) copyCols := []string{ "test_parent1_id", @@ -217,26 +210,22 @@ func copyFromTestParent1Childrens(ctx context.Context, s pgSearch.Deleter, tx *p } idx := 0 - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - inputRows = append(inputRows, []interface{}{ - testParent1ID, - idx, - obj.GetChildId(), - }) - - idx++ - } - - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent1_childrens"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } - // clear the input rows for the next batch - inputRows = inputRows[:0] + obj := objs[idx] + idx++ + + return []interface{}{ + testParent1ID, + idx, + obj.GetChildId(), + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent1_childrens"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent2/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent2/store.go index 3b3860f34a7f1..e298fb1ab175f 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent2/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent2/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -118,12 +117,18 @@ func copyFromTestParent2(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -132,39 +137,29 @@ func copyFromTestParent2(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetParentId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent2"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetParentId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent2"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent3/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent3/store.go index bb7c91331c9a1..aad15dda418cd 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent3/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent3/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -117,12 +116,18 @@ func copyFromTestParent3(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -131,39 +136,29 @@ func copyFromTestParent3(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetParentId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent3"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetParentId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent3"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent4/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent4/store.go index f4f58f67fb3de..056d687a4f563 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent4/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testparent4/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -118,12 +117,18 @@ func copyFromTestParent4(ctx context.Context, s pgSearch.Deleter, tx *postgres.T if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -132,39 +137,29 @@ func copyFromTestParent4(ctx context.Context, s pgSearch.Deleter, tx *postgres.T "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetId()), - obj.GetParentId(), - obj.GetVal(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent4"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetId()), + obj.GetParentId(), + obj.GetVal(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_parent4"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testgraphtables/testshortcircuit/store.go b/tools/generate-helpers/pg-table-bindings/testgraphtables/testshortcircuit/store.go index ac94fd261f9de..0513acb21a1a1 100644 --- a/tools/generate-helpers/pg-table-bindings/testgraphtables/testshortcircuit/store.go +++ b/tools/generate-helpers/pg-table-bindings/testgraphtables/testshortcircuit/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -117,12 +116,18 @@ func copyFromTestShortCircuits(ctx context.Context, s pgSearch.Deleter, tx *post if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetId()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "id", @@ -131,39 +136,29 @@ func copyFromTestShortCircuits(ctx context.Context, s pgSearch.Deleter, tx *post "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - obj.GetId(), - obj.GetChildId(), - obj.GetG2GrandchildId(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetId()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_short_circuits"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + obj.GetId(), + obj.GetChildId(), + obj.GetG2GrandchildId(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_short_circuits"}, copyCols, inputRows); err != nil { + return err } return nil diff --git a/tools/generate-helpers/pg-table-bindings/testuuidkey/postgres/store.go b/tools/generate-helpers/pg-table-bindings/testuuidkey/postgres/store.go index 549a26d076d74..a2c82aac942f1 100644 --- a/tools/generate-helpers/pg-table-bindings/testuuidkey/postgres/store.go +++ b/tools/generate-helpers/pg-table-bindings/testuuidkey/postgres/store.go @@ -4,7 +4,6 @@ package postgres import ( "context" - "slices" "time" "github.com/jackc/pgx/v5" @@ -127,12 +126,18 @@ func copyFromTestSingleUUIDKeyStructs(ctx context.Context, s pgSearch.Deleter, t if len(objs) == 0 { return nil } - batchSize := min(len(objs), pgSearch.MaxBatchSize) - inputRows := make([][]interface{}, 0, batchSize) - // This is a copy so first we must delete the rows and re-add them - // Which is essentially the desired behaviour of an upsert. - deletes := make([]string, 0, batchSize) + { + // CopyFrom does not upsert, so delete existing rows first to achieve upsert behavior. + // Parent deletion cascades to children, so only the top-level parent needs deletion. + deletes := make([]string, 0, len(objs)) + for _, obj := range objs { + deletes = append(deletes, obj.GetKey()) + } + if err := s.DeleteMany(ctx, deletes); err != nil { + return err + } + } copyCols := []string{ "key", @@ -149,47 +154,37 @@ func copyFromTestSingleUUIDKeyStructs(ctx context.Context, s pgSearch.Deleter, t "serialized", } - for objBatch := range slices.Chunk(objs, batchSize) { - for _, obj := range objBatch { - - serialized, marshalErr := obj.MarshalVT() - if marshalErr != nil { - return marshalErr - } - - inputRows = append(inputRows, []interface{}{ - pgutils.NilOrUUID(obj.GetKey()), - obj.GetName(), - obj.GetStringSlice(), - obj.GetBool(), - obj.GetUint64(), - obj.GetInt64(), - obj.GetFloat(), - pgutils.EmptyOrMap(obj.GetLabels()), - protocompat.NilOrTime(obj.GetTimestamp()), - obj.GetEnum(), - obj.GetEnums(), - serialized, - }) - - // Add the ID to be deleted. - deletes = append(deletes, obj.GetKey()) + idx := 0 + inputRows := pgx.CopyFromFunc(func() ([]any, error) { + if idx >= len(objs) { + return nil, nil } + obj := objs[idx] + idx++ - // copy does not upsert so have to delete first. parent deletion cascades so only need to - // delete for the top level parent - - if err := s.DeleteMany(ctx, deletes); err != nil { - return err + serialized, marshalErr := obj.MarshalVT() + if marshalErr != nil { + return nil, marshalErr } - // clear the inserts and vals for the next batch - deletes = deletes[:0] - if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_single_uuid_key_structs"}, copyCols, pgx.CopyFromRows(inputRows)); err != nil { - return err - } - // clear the input rows for the next batch - inputRows = inputRows[:0] + return []interface{}{ + pgutils.NilOrUUID(obj.GetKey()), + obj.GetName(), + obj.GetStringSlice(), + obj.GetBool(), + obj.GetUint64(), + obj.GetInt64(), + obj.GetFloat(), + pgutils.EmptyOrMap(obj.GetLabels()), + protocompat.NilOrTime(obj.GetTimestamp()), + obj.GetEnum(), + obj.GetEnums(), + serialized, + }, nil + }) + + if _, err := tx.CopyFrom(ctx, pgx.Identifier{"test_single_uuid_key_structs"}, copyCols, inputRows); err != nil { + return err } return nil