perf(pebble): unlink spill chunks during merge, sort resources on import

cmd/baton/sanitize.go

@@ -28,6 +28,7 @@ func sanitizeCmd() *cobra.Command {
 	cmd.Flags().String("secret-file", "", "Path to a per-c1z HMAC secret (>=32 random bytes). If unset, a fresh secret is generated and written next to --out.")
 	cmd.Flags().String("anchor", "", "RFC3339 timestamp the newest source timestamp lands on. Defaults to now.")
 	cmd.Flags().Bool("allow-unknown-annotations", false, "Pass annotations of unknown type through unchanged instead of dropping. Dangerous on real customer data.")
+	cmd.Flags().String("tmp-dir", "", "The temporary directory to use while sanitizing")
 
 	return cmd
 }
@@ -63,6 +64,10 @@ func runSanitize(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	tmpDir, err := cmd.Flags().GetString("tmp-dir")
+	if err != nil {
+		return err
+	}
 	if outPath == "" {
 		return fmt.Errorf("--out is required")
 	}
@@ -103,7 +108,11 @@ func runSanitize(cmd *cobra.Command, args []string) error {
 			zap.String("path", c1zsanitize.SecretPath(secretFile, outPath)))
 	}
 
-	src, err := openReadOnlyC1ZStore(ctx, inPath)
+	var srcOpts []dotc1z.C1ZOption
+	if tmpDir != "" {
+		srcOpts = append(srcOpts, dotc1z.WithTmpDir(tmpDir))
+	}
+	src, err := openReadOnlyC1ZStore(ctx, inPath, srcOpts...)
 	if err != nil {
 		return fmt.Errorf("open source c1z: %w", err)
 	}
@@ -136,6 +145,9 @@ func runSanitize(cmd *cobra.Command, args []string) error {
 	// These pragmas are scoped to THIS writer instance; normal connector
 	// syncs open their own store and are unaffected.
 	dstOpts := []dotc1z.C1ZOption{dotc1z.WithEngine(dstEngine)}
+	if tmpDir != "" {
+		dstOpts = append(dstOpts, dotc1z.WithTmpDir(tmpDir))
+	}
 	if dstEngine == c1zstore.EngineSQLite {
 		dstOpts = append(dstOpts,
 			dotc1z.WithPragma("journal_mode", "OFF"),
@@ -165,6 +177,7 @@ func runSanitize(cmd *cobra.Command, args []string) error {
 		Secret:                  secret,
 		TimestampAnchor:         anchor,
 		AllowUnknownAnnotations: allowUnknown,
+		TmpDir:                  tmpDir,
 	}
 
 	log.Info("c1zsanitize: starting",

cmd/baton/store.go

@@ -10,8 +10,8 @@ import (
 	"github.com/conductorone/baton-sdk/pkg/dotc1z/c1zstore"
 )
 
-func openReadOnlyC1ZStore(ctx context.Context, path string) (c1zstore.Store, error) {
-	c1zStore, err := dotc1z.NewStore(ctx, path, dotc1z.WithReadOnly(true))
+func openReadOnlyC1ZStore(ctx context.Context, path string, extraOpts ...dotc1z.C1ZOption) (c1zstore.Store, error) {
+	c1zStore, err := dotc1z.NewStore(ctx, path, append([]dotc1z.C1ZOption{dotc1z.WithReadOnly(true)}, extraOpts...)...)
 	if err != nil {
 		return nil, err
 	}

pkg/c1zsanitize/assets.go

@@ -97,12 +97,16 @@ func closeIfCloser(r io.Reader) error {
 	return nil
 }
 
+// copyAssets returns the number of assets written, which the pebble bulk
+// path stashes into the sync's stats sidecar (assets ride outside the bulk
+// import, so its ComputedStats cannot count them).
 func (s *sanitizer) copyAssets(
 	ctx context.Context,
 	src connectorstore.Reader,
 	dst connectorstore.Writer,
 	refs *assetRefSet,
-) error {
+) (int64, error) {
+	var written int64
 	ids := refs.drain()
 	for _, srcID := range ids {
 		req := v2.AssetServiceGetAssetRequest_builder{
@@ -119,12 +123,13 @@ func (s *sanitizer) copyAssets(
 			continue
 		}
 		if err := closeIfCloser(r); err != nil {
-			return fmt.Errorf("close source asset %s: %w", srcID, err)
+			return 0, fmt.Errorf("close source asset %s: %w", srcID, err)
 		}
 		dstID := s.id(srcID)
 		if err := dst.PutAsset(ctx, v2.AssetRef_builder{Id: dstID}.Build(), contentType, placeholderForContentType(contentType)); err != nil {
-			return fmt.Errorf("put dst asset %s: %w", dstID, err)
+			return 0, fmt.Errorf("put dst asset %s: %w", dstID, err)
 		}
+		written++
 	}
-	return nil
+	return written, nil
 }

pkg/c1zsanitize/bulk_sink.go

@@ -0,0 +1,109 @@
+package c1zsanitize
+
+import (
+	"context"
+	"fmt"
+
+	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
+	"github.com/conductorone/baton-sdk/pkg/dotc1z/engine/pebble"
+)
+
+// recordSink is where sanitizeSync's copy loops write the four record
+// families. It is exactly the method subset of connectorstore.Writer those
+// loops use, so a sqlite destination satisfies it as-is (upserting Put*
+// path); a pebble destination substitutes bulkImportSink so the same
+// transform loops feed the engine's bulk-import fast path instead.
+type recordSink interface {
+	PutResourceTypes(ctx context.Context, resourceTypes ...*v2.ResourceType) error
+	PutResources(ctx context.Context, resources ...*v2.Resource) error
+	PutEntitlements(ctx context.Context, entitlements ...*v2.Entitlement) error
+	PutGrants(ctx context.Context, grants ...*v2.Grant) error
+}
+
+// bulkImportSink adapts a pebble BulkSyncImport to recordSink. The bulk
+// contract (fresh sync, nothing else writes until Finish) holds on the
+// sanitize path by construction: StartNewSync marked the destination sync
+// fresh, resumable+pebble is rejected up front so every checkpoint call is
+// a no-op, and assets are copied only after finish() has ingested.
+//
+// Semantics vs. the Put* path: Put* upserts, the bulk path does not. For
+// resource types, resources, and entitlements a source sync carrying two
+// records with the same sanitized external id now fails the import
+// instead of silently last-write-wins; a valid c1z has unique external ids
+// per sync, so that only surfaces on corrupt input. Grants are keyed by
+// structural identity (entitlement + principal refs), which SQLite's
+// UNIQUE(external_id, sync_id) does not cover, so a legacy source can
+// legitimately hold two grants with distinct external ids and identical
+// refs. Those now fold at Finish with a warning — keeping the
+// earliest-discovered external id and OR-ing needs_expansion — where
+// PutGrants on pebble kept whichever arrived last. Same row count, but the
+// surviving transformed grant id can differ from the previous path's.
+//
+// Grants flow through a single shard: the sanitizer's grant loop is
+// sequential, so shard fan-out would add nothing.
+type bulkImportSink struct {
+	eng    *pebble.Engine
+	bi     *pebble.BulkSyncImport
+	shard  *pebble.BulkGrantShard
+	syncID string
+}
+
+// startBulkImportSink opens a bulk import on the destination's current
+// fresh sync. tmpDir stages spill files ("" = system temp dir). The
+// import is sized for the single grant shard this sink opens.
+func startBulkImportSink(ctx context.Context, eng *pebble.Engine, syncID string, tmpDir string) (*bulkImportSink, error) {
+	bi, err := eng.StartBulkSyncImport(ctx, syncID, tmpDir, 1)
+	if err != nil {
+		return nil, fmt.Errorf("start bulk import: %w", err)
+	}
+	shard, err := bi.NewGrantShard()
+	if err != nil {
+		bi.Abort()
+		return nil, fmt.Errorf("open grant shard: %w", err)
+	}
+	return &bulkImportSink{eng: eng, bi: bi, shard: shard, syncID: syncID}, nil
+}
+
+func (s *bulkImportSink) PutResourceTypes(ctx context.Context, resourceTypes ...*v2.ResourceType) error {
+	// copyResourceTypes writes the full set once, sorted by output id,
+	// which is the sorted-by-external-id arrival AddResourceTypes requires.
+	return s.bi.AddResourceTypes(ctx, resourceTypes...)
+}
+
+func (s *bulkImportSink) PutResources(ctx context.Context, resources ...*v2.Resource) error {
+	return s.bi.AddResources(ctx, resources...)
+}
+
+func (s *bulkImportSink) PutEntitlements(ctx context.Context, entitlements ...*v2.Entitlement) error {
+	return s.bi.AddEntitlements(ctx, entitlements...)
+}
+
+func (s *bulkImportSink) PutGrants(ctx context.Context, grants ...*v2.Grant) error {
+	return s.shard.AddGrants(ctx, grants...)
+}
+
+// finish seals the grant shard and ingests the import. After it returns
+// nil the destination sync holds every record the sink received and the
+// writer path (PutAsset, EndSync) may be used again.
+func (s *bulkImportSink) finish(ctx context.Context) error {
+	s.shard.Close()
+	return s.bi.Finish(ctx)
+}
+
+// abort discards a still-open import's staged spill files. Deferred by
+// sanitizeSync to cover error exits before finish is reached; once Finish
+// has run — success or failure — it has marked the import done and torn
+// down its own staging, and Abort is a no-op.
+func (s *bulkImportSink) abort() {
+	s.bi.Abort()
+}
+
+// stashStats hands the import's record counts (plus the asset count, which
+// rides outside the import) to the engine as the sync's stats sidecar, so
+// EndSync persists stats directly instead of re-scanning the freshly
+// ingested keyspaces.
+func (s *bulkImportSink) stashStats(assetCount int64) {
+	rec := s.bi.ComputedStats()
+	rec.SetAssets(assetCount)
+	s.eng.StashComputedSyncStats(s.syncID, rec)
+}