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,101 @@
+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. A
+// source sync carrying two records with the same sanitized external id now
+// fails the import (resource types/resources/entitlements) or folds with a
+// warning (grants) instead of silently last-write-wins. A valid c1z has
+// unique external ids per sync, so this only surfaces on corrupt input.
+//
+// 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).
+func startBulkImportSink(ctx context.Context, eng *pebble.Engine, syncID string, tmpDir string) (*bulkImportSink, error) {
+	bi, err := eng.StartBulkSyncImport(ctx, syncID, tmpDir)
+	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)
+}

pkg/c1zsanitize/expansion_parity_test.go

@@ -8,6 +8,8 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
+	"google.golang.org/protobuf/encoding/prototext"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
@@ -258,6 +260,95 @@ func entitlementCount(t *testing.T, ctx context.Context, store c1zstore.Store) i
 	return len(resp.GetList())
 }
 
+// fullRecords holds every resource-type, resource, and entitlement record
+// keyed by transformed id. These three families now ride different write
+// paths per destination engine (upserting Put* on sqlite, the non-upserting
+// bulk import on pebble), so the parity oracle pins full record identity —
+// not just cardinality and graph shape — to catch a row one path drops,
+// folds, or mangles field-by-field.
+type fullRecords struct {
+	rts  map[string]*v2.ResourceType
+	ress map[string]*v2.Resource
+	ents map[string]*v2.Entitlement
+}
+
+func fullRecordSets(t *testing.T, ctx context.Context, store c1zstore.Store) fullRecords {
+	t.Helper()
+	fr := fullRecords{
+		rts:  map[string]*v2.ResourceType{},
+		ress: map[string]*v2.Resource{},
+		ents: map[string]*v2.Entitlement{},
+	}
+	// Single-page reads: if a fixture ever outgrows the page, both arms
+	// would truncate identically and the oracle would silently narrow, so
+	// fail loudly instead of paginating.
+	rts, err := store.ListResourceTypes(ctx, v2.ResourceTypesServiceListResourceTypesRequest_builder{PageSize: 1000}.Build())
+	require.NoError(t, err)
+	require.Empty(t, rts.GetNextPageToken(), "resource-type fixture outgrew one page; oracle would silently truncate")
+	for _, rt := range rts.GetList() {
+		fr.rts[rt.GetId()] = rt
+	}
+	ress, err := store.ListResources(ctx, v2.ResourcesServiceListResourcesRequest_builder{PageSize: 1000}.Build())
+	require.NoError(t, err)
+	require.Empty(t, ress.GetNextPageToken(), "resource fixture outgrew one page; oracle would silently truncate")
+	for _, r := range ress.GetList() {
+		fr.ress[r.GetId().GetResourceType()+"/"+r.GetId().GetResource()] = r
+	}
+	ents, err := store.ListEntitlements(ctx, v2.EntitlementsServiceListEntitlementsRequest_builder{PageSize: 1000}.Build())
+	require.NoError(t, err)
+	require.Empty(t, ents.GetNextPageToken(), "entitlement fixture outgrew one page; oracle would silently truncate")
+	for _, e := range ents.GetList() {
+		fr.ents[e.GetId()] = normalizeEntitlementForParity(e)
+	}
+	return fr
+}
+
+// normalizeEntitlementForParity slims the two fields pebble's v3 schema
+// stores as identity-only refs: the embedded resource and the grantable_to
+// types are "hydrated as a stub (identity only)" on read (V3EntitlementToV2),
+// while sqlite returns the writer's full embedded copies. Both are derived
+// data — the resource's and resource types' own records are compared in full
+// separately — so slimming them costs no coverage.
+func normalizeEntitlementForParity(e *v2.Entitlement) *v2.Entitlement {
+	out, ok := proto.Clone(e).(*v2.Entitlement)
+	if !ok {
+		panic("clone changed type")
+	}
+	if r := out.GetResource(); r != nil {
+		out.SetResource(v2.Resource_builder{Id: r.GetId()}.Build())
+	}
+	if gs := out.GetGrantableTo(); len(gs) > 0 {
+		slim := make([]*v2.ResourceType, 0, len(gs))
+		for _, rt := range gs {
+			slim = append(slim, v2.ResourceType_builder{Id: rt.GetId()}.Build())
+		}
+		out.SetGrantableTo(slim)
+	}
+	return out
+}
+
+// requireRecordSetEqual asserts got holds exactly want's keys and that each
+// record is proto-equal, printing the diverging pair on failure.
+func requireRecordSetEqual[P proto.Message](t *testing.T, want, got map[string]P, arm, family string) {
+	t.Helper()
+	wantKeys := sortedKeys(want)
+	require.Equalf(t, wantKeys, sortedKeys(got), "%s %s id set", arm, family)
+	for _, id := range wantKeys {
+		require.Truef(t, proto.Equal(want[id], got[id]),
+			"%s %s %q diverged\nbase: %s\narm:  %s",
+			arm, family, id, prototext.Format(want[id]), prototext.Format(got[id]))
+	}
+}
+
+func sortedKeys[V any](m map[string]V) []string {
+	keys := make([]string, 0, len(m))
+	for k := range m {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+	return keys
+}
+
 // resourceParentChains maps each resource id to its parent resource id.
 func resourceParentChains(t *testing.T, ctx context.Context, store c1zstore.Store) map[string]string {
 	t.Helper()
@@ -388,6 +479,7 @@ func TestSanitizeCrossEnginePebbleParity(t *testing.T) {
 		parents, entOwn, refs                   map[string]string
 		sources                                 map[string][]string
 		blobs                                   map[string]expandBlob
+		records                                 fullRecords
 		assetCT                                 string
 		assetBytes                              []byte
 	}
@@ -405,6 +497,7 @@ func TestSanitizeCrossEnginePebbleParity(t *testing.T) {
 			refs:       grantRefs(t, ctx, ro),
 			sources:    grantSourcesCanonical(t, ctx, ro),
 			blobs:      pendingExpansionBlobs(t, ctx, ro),
+			records:    fullRecordSets(t, ctx, ro),
 			assetCT:    ct,
 			assetBytes: ab,
 		}
@@ -429,6 +522,9 @@ func TestSanitizeCrossEnginePebbleParity(t *testing.T) {
 		require.Equalf(t, base.refs, snaps[i].refs, "%s grant entitlement/principal refs", arms[i].name)
 		require.Equalf(t, base.sources, snaps[i].sources, "%s expansion-source edges", arms[i].name)
 		require.Equalf(t, base.blobs, snaps[i].blobs, "%s GrantExpandable blobs + needs_expansion enumeration", arms[i].name)
+		requireRecordSetEqual(t, base.records.rts, snaps[i].records.rts, arms[i].name, "resource type")
+		requireRecordSetEqual(t, base.records.ress, snaps[i].records.ress, arms[i].name, "resource")
+		requireRecordSetEqual(t, base.records.ents, snaps[i].records.ents, arms[i].name, "entitlement")
 	}
 
 	// (3) needs_expansion membership: the multi-expand and divergence-trigger