diff --git a/README.md b/README.md index 9665f3a9..a58b7767 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,8 @@ comparison mode yet. The command can emit JSONL records plus Markdown and JSON s against a previous JSONL baseline. `go run ./cmd/retriever` dumps and loads live Dawgs graph databases as -manifest-based collections of compressed JSONL fragments. It supports +manifest-based collections of compressed JSONL fragments, with an optional +Parquet sidecar for dumps. It supports PostgreSQL and Neo4j, gzip and zstd compression, checksum validation before load, optional deterministic property scrubbing, and a read-throughput benchmark mode. It can also package dumps as single HPKE/ML-KEM encrypted TAR archives. @@ -87,10 +88,11 @@ archive, load, scrubbed dump, metrics verification, and benchmark examples. The same import/export functionality is available to library consumers from `github.com/specterops/dawgs/retriever`; callers provide an already-open `graph.Database`, and archive helpers support both path-based and stream-based -APIs. The package exposes CLI-matching default option constructors, structured -progress callbacks, manifest/metrics helpers, HPKE key envelope reader/writer -helpers, and typed errors for validation, compatibility, checksum, metrics, and -count mismatches. +APIs. Library dumps enable the optional sidecar with `DumpOptions.Parquet` and +receive its manifest and success-marker paths in `DumpResult`. The package +exposes CLI-matching default option constructors, structured progress callbacks, +manifest/metrics helpers, HPKE key envelope reader/writer helpers, and typed +errors for validation, compatibility, checksum, metrics, and count mismatches. PostgreSQL translates exact string property equality with a JSON string type guard and `properties ->>` extraction, so indexes created on expressions such as `properties ->> 'objectid'` and `properties ->> 'name'` can be used for selective diff --git a/cmd/retriever/README.md b/cmd/retriever/README.md index 381101ef..52061a55 100644 --- a/cmd/retriever/README.md +++ b/cmd/retriever/README.md @@ -25,6 +25,7 @@ retriever dump \ -scrub none \ -compression zstd \ -zstd-level 11 \ + -parquet \ -shard-size 100000 ``` @@ -34,9 +35,10 @@ validates that expected node and edge partitions exist. For Neo4j, `-all-graphs` means the selected Neo4j database only. Existing non-empty output directories are refused unless `-force` is supplied. -The manifest is written last as `manifest.json`; if a dump fails before that -point, the directory is intentionally left for inspection without a success -manifest. +For JSONL-only dumps, the manifest is written last as `manifest.json`. With a +Parquet sidecar, its metadata precedes `manifest.json` and `parquet/_SUCCESS` +is written last. Failed dumps intentionally leave completed fragments for +inspection without the final applicable success boundary. New dumps include a `retriever-metrics-v1` manifest section with graph metrics computed from the same node and relationship streams written to the fragments. @@ -54,6 +56,28 @@ Dump progress is emitted with `log/slog` on stderr. Notices mark output directory preparation, graph counting, scrub pre-pass work, node and relationship phase boundaries, periodic entity progress, manifest writing, and completion. +### Parquet sidecar + +Pass `-parquet` to write a strict Parquet sidecar in addition to the normal +compressed JSONL collection. JSONL remains mandatory and keeps its existing +paths, bytes, manifest, load behavior, and checksums. + +Parquet fragments use matching logical shard boundaries under +`parquet/graphs//`. Node rows contain `id`, `kinds`, and JSON-annotated +`properties`. Relationship rows contain `id`, `start_id`, `end_id`, `kind`, and +JSON-annotated `properties`; unlike JSONL v1, the Parquet sidecar retains the +source relationship ID. Columns use Zstandard compression. + +The sidecar metadata is `parquet/manifest.json` with format +`retriever-parquet-export-v1`. Publication writes that metadata first, the +existing JSONL `manifest.json` second, and `parquet/_SUCCESS` last. A requested +Parquet failure fails the dump, and `_SUCCESS` is absent unless both outputs +and their metadata completed. + +Load and verify continue to consume the JSONL collection. Encrypted archive +creation also retains its existing behavior and packages the JSONL manifest +and files; it does not include the optional Parquet sidecar. + ## Encrypted Archives Generate a recipient key pair before creating encrypted archives: diff --git a/cmd/retriever/main.go b/cmd/retriever/main.go index 124b2ae5..6784611a 100644 --- a/cmd/retriever/main.go +++ b/cmd/retriever/main.go @@ -89,6 +89,7 @@ func (s commandRuntime) runDump(ctx context.Context, args []string) error { allGraphs := flags.Bool("all-graphs", false, "Dump every graph discoverable by the selected driver.") flags.StringVar(&cfg.OutputDir, "out", "", "Output collection directory.") flags.BoolVar(&cfg.Force, "force", false, "Replace an existing non-empty output directory.") + flags.BoolVar(&cfg.Parquet, "parquet", false, "Write a Parquet sidecar export in addition to JSONL.") flags.StringVar(&archiveOut, "archive-out", "", "Optional encrypted archive output path.") flags.StringVar(&recipientPath, "recipient", "", "Recipient public key for -archive-out.") flags.StringVar(&scrubValue, "scrub", string(cfg.Scrub), "Scrub mode: none or full.") @@ -171,8 +172,12 @@ func (s commandRuntime) runDump(ctx context.Context, args []string) error { archiveLine = fmt.Sprintf("archive: %s\n", archiveOut) } + var parquetLines string + if result.ParquetManifestPath != "" { + parquetLines = fmt.Sprintf("parquet manifest: %s\nparquet success: %s\n", result.ParquetManifestPath, result.ParquetSuccessPath) + } - fmt.Fprintf(s.stdout, "dumped %d graph(s)\nmanifest: %s\n%snodes: %d\nrelationships: %d\n", len(result.Manifest.Graphs), result.ManifestPath, archiveLine, result.NodeCount, result.EdgeCount) + fmt.Fprintf(s.stdout, "dumped %d graph(s)\nmanifest: %s\n%s%snodes: %d\nrelationships: %d\n", len(result.Manifest.Graphs), result.ManifestPath, parquetLines, archiveLine, result.NodeCount, result.EdgeCount) return nil } diff --git a/go.mod b/go.mod index 1f380c05..60cb028f 100644 --- a/go.mod +++ b/go.mod @@ -15,6 +15,7 @@ require ( github.com/jackc/pgx/v5 v5.10.0 github.com/klauspost/compress v1.19.0 github.com/neo4j/neo4j-go-driver/v5 v5.28.4 + github.com/parquet-go/parquet-go v0.30.1 github.com/pashagolub/pgxmock/v5 v5.1.0 github.com/pelletier/go-toml/v2 v2.4.3 github.com/stretchr/testify v1.11.1 @@ -59,6 +60,7 @@ require ( github.com/alfatraining/structtag v1.0.0 // indirect github.com/alingse/asasalint v0.0.11 // indirect github.com/alingse/nilnesserr v0.2.0 // indirect + github.com/andybalholm/brotli v1.1.1 // indirect github.com/ashanbrown/forbidigo/v2 v2.3.0 // indirect github.com/ashanbrown/makezero/v2 v2.1.0 // indirect github.com/aymanbagabas/go-osc52/v2 v2.0.1 // indirect @@ -120,6 +122,7 @@ require ( github.com/golangci/swaggoswag v0.0.0-20250504205917-77f2aca3143e // indirect github.com/golangci/unconvert v0.0.0-20250410112200-a129a6e6413e // indirect github.com/google/go-cmp v0.7.0 // indirect + github.com/google/uuid v1.6.0 // indirect github.com/gordonklaus/ineffassign v0.2.0 // indirect github.com/gostaticanalysis/analysisutil v0.7.1 // indirect github.com/gostaticanalysis/comment v1.5.0 // indirect @@ -172,6 +175,9 @@ require ( github.com/nishanths/exhaustive v0.12.0 // indirect github.com/nishanths/predeclared v0.2.2 // indirect github.com/nunnatsa/ginkgolinter v0.23.0 // indirect + github.com/parquet-go/bitpack v1.0.0 // indirect + github.com/parquet-go/jsonlite v1.0.0 // indirect + github.com/pierrec/lz4/v4 v4.1.21 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/prometheus/client_golang v1.22.0 // indirect @@ -212,6 +218,7 @@ require ( github.com/timonwong/loggercheck v0.11.0 // indirect github.com/tomarrell/wrapcheck/v2 v2.12.0 // indirect github.com/tommy-muehle/go-mnd/v2 v2.5.1 // indirect + github.com/twpayne/go-geom v1.6.1 // indirect github.com/ultraware/funlen v0.2.0 // indirect github.com/ultraware/whitespace v0.2.0 // indirect github.com/uudashr/gocognit v1.2.1 // indirect diff --git a/go.sum b/go.sum index bfc4bb5c..9cf8e446 100644 --- a/go.sum +++ b/go.sum @@ -31,6 +31,8 @@ github.com/Antonboom/testifylint v1.6.4/go.mod h1:YO33FROXX2OoUfwjz8g+gUxQXio5i9 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU= github.com/BurntSushi/toml v1.6.0 h1:dRaEfpa2VI55EwlIW72hMRHdWouJeRF7TPYhI+AUQjk= github.com/BurntSushi/toml v1.6.0/go.mod h1:ukJfTF/6rtPPRCnwkur4qwRxa8vTRFBF0uk2lLoLwho= +github.com/DATA-DOG/go-sqlmock v1.5.2 h1:OcvFkGmslmlZibjAjaHm3L//6LiuBgolP7OputlJIzU= +github.com/DATA-DOG/go-sqlmock v1.5.2/go.mod h1:88MAG/4G7SMwSE3CeA0ZKzrT5CiOU3OJ+JlNzwDqpNU= github.com/Djarvur/go-err113 v0.1.1 h1:eHfopDqXRwAi+YmCUas75ZE0+hoBHJ2GQNLYRSxao4g= github.com/Djarvur/go-err113 v0.1.1/go.mod h1:IaWJdYFLg76t2ihfflPZnM1LIQszWOsFDh2hhhAVF6k= github.com/Masterminds/semver/v3 v3.1.1/go.mod h1:VPu/7SZ7ePZ3QOrcuXROw5FAcLl4a0cBrbBpGY/8hQs= @@ -60,6 +62,8 @@ github.com/alingse/asasalint v0.0.11 h1:SFwnQXJ49Kx/1GghOFz1XGqHYKp21Kq1nHad/0WQ github.com/alingse/asasalint v0.0.11/go.mod h1:nCaoMhw7a9kSJObvQyVzNTPBDbNpdocqrSP7t/cW5+I= github.com/alingse/nilnesserr v0.2.0 h1:raLem5KG7EFVb4UIDAXgrv3N2JIaffeKNtcEXkEWd/w= github.com/alingse/nilnesserr v0.2.0/go.mod h1:1xJPrXonEtX7wyTq8Dytns5P2hNzoWymVUIaKm4HNFg= +github.com/andybalholm/brotli v1.1.1 h1:PR2pgnyFznKEugtsUo0xLdDop5SKXd5Qf5ysW+7XdTA= +github.com/andybalholm/brotli v1.1.1/go.mod h1:05ib4cKhjx3OQYUY22hTVd34Bc8upXjOLL2rKwwZBoA= github.com/antlr4-go/antlr/v4 v4.13.1 h1:SqQKkuVZ+zWkMMNkjy5FZe5mr5WURWnlpmOuzYWrPrQ= github.com/antlr4-go/antlr/v4 v4.13.1/go.mod h1:GKmUxMtwp6ZgGwZSva4eWPC5mS6vUAmOABFgjdkM7Nw= github.com/ashanbrown/forbidigo/v2 v2.3.0 h1:OZZDOchCgsX5gvToVtEBoV2UWbFfI6RKQTir2UZzSxo= @@ -443,10 +447,18 @@ github.com/otiai10/curr v0.0.0-20150429015615-9b4961190c95/go.mod h1:9qAhocn7zKJ github.com/otiai10/curr v1.0.0/go.mod h1:LskTG5wDwr8Rs+nNQ+1LlxRjAtTZZjtJW4rMXl6j4vs= github.com/otiai10/mint v1.3.0/go.mod h1:F5AjcsTsWUqX+Na9fpHb52P8pcRX2CI6A3ctIT91xUo= github.com/otiai10/mint v1.3.1/go.mod h1:/yxELlJQ0ufhjUwhshSj+wFjZ78CnZ48/1wtmBH1OTc= +github.com/parquet-go/bitpack v1.0.0 h1:AUqzlKzPPXf2bCdjfj4sTeacrUwsT7NlcYDMUQxPcQA= +github.com/parquet-go/bitpack v1.0.0/go.mod h1:XnVk9TH+O40eOOmvpAVZ7K2ocQFrQwysLMnc6M/8lgs= +github.com/parquet-go/jsonlite v1.0.0 h1:87QNdi56wOfsE5bdgas0vRzHPxfJgzrXGml1zZdd7VU= +github.com/parquet-go/jsonlite v1.0.0/go.mod h1:nDjpkpL4EOtqs6NQugUsi0Rleq9sW/OtC1NnZEnxzF0= +github.com/parquet-go/parquet-go v0.30.1 h1:Oy6ganNrAdFiVwy7wNmWagfPTWA2X9Z3tVHBc7JtuX8= +github.com/parquet-go/parquet-go v0.30.1/go.mod h1:navtkAYr2LGoJVp141oXPlO/sxLvaOe3la2JEoD8+rg= github.com/pashagolub/pgxmock/v5 v5.1.0 h1:NZ4pl82b335sEGIbD/+tk2fVIgVs3yNWr1R42ukpUvU= github.com/pashagolub/pgxmock/v5 v5.1.0/go.mod h1:8IJct22b7+EuqecVmYb9aKiENJLLqTsbjFHXH/znAEg= github.com/pelletier/go-toml/v2 v2.4.3 h1:GTRvJQutkOSftxIFD5xw9aepkYNuPWmVJpffdDPYVpY= github.com/pelletier/go-toml/v2 v2.4.3/go.mod h1:2gIqNv+qfxSVS7cM2xJQKtLSTLUE9V8t9Stt+h56mCY= +github.com/pierrec/lz4/v4 v4.1.21 h1:yOVMLb6qSIDP67pl/5F7RepeKYu/VmTyEXvuMI5d9mQ= +github.com/pierrec/lz4/v4 v4.1.21/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4= github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4= github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0= @@ -568,6 +580,8 @@ github.com/tomarrell/wrapcheck/v2 v2.12.0 h1:H/qQ1aNWz/eeIhxKAFvkfIA+N7YDvq6TWVF github.com/tomarrell/wrapcheck/v2 v2.12.0/go.mod h1:AQhQuZd0p7b6rfW+vUwHm5OMCGgp63moQ9Qr/0BpIWo= github.com/tommy-muehle/go-mnd/v2 v2.5.1 h1:NowYhSdyE/1zwK9QCLeRb6USWdoif80Ie+v+yU8u1Zw= github.com/tommy-muehle/go-mnd/v2 v2.5.1/go.mod h1:WsUAkMJMYww6l/ufffCD3m+P7LEvr8TnZn9lwVDlgzw= +github.com/twpayne/go-geom v1.6.1 h1:iLE+Opv0Ihm/ABIcvQFGIiFBXd76oBIar9drAwHFhR4= +github.com/twpayne/go-geom v1.6.1/go.mod h1:Kr+Nly6BswFsKM5sd31YaoWS5PeDDH2NftJTK7Gd028= github.com/ultraware/funlen v0.2.0 h1:gCHmCn+d2/1SemTdYMiKLAHFYxTYz7z9VIDRaTGyLkI= github.com/ultraware/funlen v0.2.0/go.mod h1:ZE0q4TsJ8T1SQcjmkhN/w+MceuatI6pBFSxxyteHIJA= github.com/ultraware/whitespace v0.2.0 h1:TYowo2m9Nfj1baEQBjuHzvMRbp19i+RCcRYrSWoFa+g= @@ -580,6 +594,8 @@ github.com/xen0n/gosmopolitan v1.3.0 h1:zAZI1zefvo7gcpbCOrPSHJZJYA9ZgLfJqtKzZ5pH github.com/xen0n/gosmopolitan v1.3.0/go.mod h1:rckfr5T6o4lBtM1ga7mLGKZmLxswUoH1zxHgNXOsEt4= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no= github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM= +github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= +github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yagipy/maintidx v1.0.0 h1:h5NvIsCz+nRDapQ0exNv4aJ0yXSI0420omVANTv3GJM= github.com/yagipy/maintidx v1.0.0/go.mod h1:0qNf/I/CCZXSMhsRsrEPDZ+DkekpKLXAJfsTACwgXLk= github.com/yeya24/promlinter v0.3.0 h1:JVDbMp08lVCP7Y6NP3qHroGAO6z2yGKQtS5JsjqtoFs= diff --git a/retriever/archive_tar_test.go b/retriever/archive_tar_test.go index 121c6f70..96d507ee 100644 --- a/retriever/archive_tar_test.go +++ b/retriever/archive_tar_test.go @@ -54,6 +54,25 @@ func TestCollectionTarDeterministicAndSorted(t *testing.T) { } } +func TestCollectionTarIgnoresFilesAbsentFromManifest(t *testing.T) { + dir := writeArchiveFixture(t) + orphanPath := filepath.Join(dir, "graphs", "secret-graph", "orphan.jsonl.gz") + if err := os.WriteFile(orphanPath, []byte("orphan"), 0o600); err != nil { + t.Fatalf("write orphan fragment: %v", err) + } + + var buffer bytes.Buffer + if err := writeCollectionTar(&buffer, dir); err != nil { + t.Fatalf("write tar: %v", err) + } + + for _, name := range tarEntryNames(t, buffer.Bytes()) { + if name == "graphs/secret-graph/orphan.jsonl.gz" { + t.Fatalf("archive included unmanifested file %q", name) + } + } +} + func TestCollectionTarStableMetadata(t *testing.T) { dir := writeArchiveFixture(t) var buffer bytes.Buffer @@ -194,23 +213,27 @@ func writeArchiveFixture(t *testing.T) string { Compression: CompressionGzip, ZstdLevel: DefaultZstdLevel, } - nodeEntry, err := writeNodeFragment(dir, "secret-graph", 1, options, []FragmentNode{{ + nodeSummary := shardSummary{ + ID: shardID{Graph: "secret-graph", Phase: PhaseNodes, Number: 1}, + Rows: 1, + } + nodeMetadata := writeTestJSONLShard(t, newJSONLNodeSink(options), nodeSummary, []normalizedNode{{ ID: "1", Kinds: []string{"User"}, Properties: map[string]any{"name": "alice"}, - }}, nil) - if err != nil { - t.Fatalf("write node fragment: %v", err) - } + }}) + nodeEntry := newJSONLFileManifest(nodeSummary, nodeMetadata) - edgeEntry, err := writeEdgeFragment(dir, "secret-graph", 1, options, []FragmentEdge{{ + edgeSummary := shardSummary{ + ID: shardID{Graph: "secret-graph", Phase: PhaseEdges, Number: 1}, + Rows: 1, + } + edgeMetadata := writeTestJSONLShard(t, newJSONLEdgeSink(options), edgeSummary, []normalizedEdge{{ StartID: "1", EndID: "1", Kind: "MemberOf", - }}, nil) - if err != nil { - t.Fatalf("write edge fragment: %v", err) - } + }}) + edgeEntry := newJSONLFileManifest(edgeSummary, edgeMetadata) value := newValidTestManifest(1) value.Schema.Graphs = []GraphSchemaMetadata{{ diff --git a/retriever/compression.go b/retriever/compression.go index ee22140b..ff14c94c 100644 --- a/retriever/compression.go +++ b/retriever/compression.go @@ -4,6 +4,7 @@ import ( "bufio" "bytes" "compress/gzip" + "context" "crypto/sha256" "encoding/hex" "encoding/json" @@ -27,18 +28,45 @@ type countingWriter struct { } type compressedJSONLinesWriter struct { - path string - tempPath string - file *os.File + artifact stagedWorkspaceFile compressor io.WriteCloser encoder *json.Encoder compressedCounter *countingWriter uncompressedCounter *countingWriter hasher hash.Hash count int - closed bool + state compressedJSONLinesWriterState } +type compressedJSONLinesWriterState uint8 + +const ( + compressedWriterOpen compressedJSONLinesWriterState = iota + compressedWriterPrepared + compressedWriterAborted +) + +type preparedCompressedJSONLinesFragment struct { + artifact stagedWorkspaceFile + metadata compressedJSONLinesMetadata + state preparedCompressedJSONLinesState +} + +type compressedJSONLinesMetadata struct { + Rows int + CompressedBytes int64 + UncompressedBytes int64 + SHA256 string +} + +type preparedCompressedJSONLinesState uint8 + +const ( + compressedFragmentPrepared preparedCompressedJSONLinesState = iota + compressedFragmentCommitted + compressedFragmentAborted +) + func (s *countingWriter) Write(p []byte) (int, error) { n, err := s.writer.Write(p) s.count += int64(n) @@ -97,26 +125,23 @@ func newDecompressionReader(reader io.Reader, codec CompressionCodec) (io.ReadCl } func newCompressedJSONLinesWriter(path string, codec CompressionCodec, zstdLevel int) (*compressedJSONLinesWriter, error) { - if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { - return nil, fmt.Errorf("create fragment directory: %w", err) - } + workspace := newLocalCollectionWorkspace(filepath.Dir(path), false) + return newCompressedJSONLinesWriterInWorkspace(context.Background(), workspace, filepath.Base(path), codec, zstdLevel) +} - tempPath := path + ".tmp" - file, err := os.OpenFile(tempPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) +func newCompressedJSONLinesWriterInWorkspace(ctx context.Context, workspace collectionWorkspace, relativePath string, codec CompressionCodec, zstdLevel int) (*compressedJSONLinesWriter, error) { + artifact, err := workspace.Stage(ctx, relativePath) if err != nil { - return nil, fmt.Errorf("open fragment temp file: %w", err) + return nil, err } hasher := sha256.New() compressedCounter := &countingWriter{ - writer: io.MultiWriter(file, hasher), + writer: io.MultiWriter(artifact, hasher), } compressor, err := newCompressionWriter(compressedCounter, codec, zstdLevel) if err != nil { - _ = file.Close() - _ = os.Remove(tempPath) - - return nil, err + return nil, cleanupOnError(err, artifact.Abort) } uncompressedCounter := &countingWriter{ @@ -126,9 +151,7 @@ func newCompressedJSONLinesWriter(path string, codec CompressionCodec, zstdLevel encoder.SetEscapeHTML(false) return &compressedJSONLinesWriter{ - path: path, - tempPath: tempPath, - file: file, + artifact: artifact, compressor: compressor, encoder: encoder, compressedCounter: compressedCounter, @@ -138,8 +161,8 @@ func newCompressedJSONLinesWriter(path string, codec CompressionCodec, zstdLevel } func (s *compressedJSONLinesWriter) Write(value any) error { - if s.closed { - return fmt.Errorf("write closed JSONL fragment") + if s.state != compressedWriterOpen { + return fmt.Errorf("write JSONL fragment after prepare or abort") } if err := s.encoder.Encode(value); err != nil { @@ -155,65 +178,99 @@ func (s *compressedJSONLinesWriter) Count() int { return s.count } -func (s *compressedJSONLinesWriter) Close() (FileManifest, error) { - if s.closed { - return FileManifest{}, fmt.Errorf("close JSONL fragment more than once") +func (s *compressedJSONLinesWriter) Prepare() (*preparedCompressedJSONLinesFragment, error) { + if s.state != compressedWriterOpen { + return nil, fmt.Errorf("prepare JSONL fragment more than once or after abort") } - s.closed = true + s.state = compressedWriterPrepared if err := s.compressor.Close(); err != nil { - _ = s.file.Close() - _ = os.Remove(s.tempPath) + s.state = compressedWriterAborted + return nil, cleanupOnError(fmt.Errorf("finish compressed fragment: %w", err), s.artifact.Abort) + } - return FileManifest{}, fmt.Errorf("finish compressed fragment: %w", err) + if err := s.artifact.Close(); err != nil { + s.state = compressedWriterAborted + return nil, cleanupOnError(fmt.Errorf("close fragment file: %w", err), s.artifact.Abort) } - if err := s.file.Close(); err != nil { - _ = os.Remove(s.tempPath) + return &preparedCompressedJSONLinesFragment{ + artifact: s.artifact, + metadata: compressedJSONLinesMetadata{ + Rows: s.count, + CompressedBytes: s.compressedCounter.count, + UncompressedBytes: s.uncompressedCounter.count, + SHA256: hex.EncodeToString(s.hasher.Sum(nil)), + }, + state: compressedFragmentPrepared, + }, nil +} - return FileManifest{}, fmt.Errorf("close fragment file: %w", err) +func (s *compressedJSONLinesWriter) Abort() error { + switch s.state { + case compressedWriterOpen: + s.state = compressedWriterAborted + return collectErrors(s.compressor.Close(), s.artifact.Abort()) + case compressedWriterAborted: + return nil + default: + return fmt.Errorf("abort JSONL writer after prepare") } +} - if err := os.Rename(s.tempPath, s.path); err != nil { - _ = os.Remove(s.tempPath) +func (s *preparedCompressedJSONLinesFragment) Metadata() compressedJSONLinesMetadata { + return s.metadata +} - return FileManifest{}, fmt.Errorf("rename fragment: %w", err) +func (s *preparedCompressedJSONLinesFragment) Commit(ctx context.Context) error { + if s.state != compressedFragmentPrepared { + return fmt.Errorf("commit JSONL fragment that is not prepared") + } + if err := ctx.Err(); err != nil { + return err } - return FileManifest{ - Count: s.count, - CompressedBytes: s.compressedCounter.count, - UncompressedBytes: s.uncompressedCounter.count, - SHA256: hex.EncodeToString(s.hasher.Sum(nil)), - }, nil + if err := s.artifact.Commit(ctx); err != nil { + s.state = compressedFragmentAborted + return cleanupOnError(err, s.artifact.Abort) + } + + s.state = compressedFragmentCommitted + return nil } -func (s *compressedJSONLinesWriter) Abort() { - if s.closed { - return +func (s *preparedCompressedJSONLinesFragment) Abort() error { + switch s.state { + case compressedFragmentPrepared: + s.state = compressedFragmentAborted + return s.artifact.Abort() + case compressedFragmentAborted: + return nil + default: + return fmt.Errorf("abort committed JSONL fragment") } - s.closed = true - - _ = s.compressor.Close() - _ = s.file.Close() - _ = os.Remove(s.tempPath) } -func writeCompressedJSONLines[T any](path string, codec CompressionCodec, zstdLevel int, records []T) (FileManifest, error) { +func writeCompressedJSONLines[T any](path string, codec CompressionCodec, zstdLevel int, records []T) (compressedJSONLinesMetadata, error) { writer, err := newCompressedJSONLinesWriter(path, codec, zstdLevel) if err != nil { - return FileManifest{}, err + return compressedJSONLinesMetadata{}, err } for _, record := range records { if err := writer.Write(record); err != nil { - writer.Abort() - - return FileManifest{}, err + return compressedJSONLinesMetadata{}, cleanupOnError(err, writer.Abort) } } - return writer.Close() + prepared, err := writer.Prepare() + if err != nil { + return compressedJSONLinesMetadata{}, err + } + if err := prepared.Commit(context.Background()); err != nil { + return compressedJSONLinesMetadata{}, cleanupOnError(err, prepared.Abort) + } + return prepared.Metadata(), nil } func readCompressedJSONLines[T any](path string, codec CompressionCodec, handle func(T) error) (int, error) { diff --git a/retriever/compression_test.go b/retriever/compression_test.go index f8011c4a..04f16cc6 100644 --- a/retriever/compression_test.go +++ b/retriever/compression_test.go @@ -2,6 +2,7 @@ package retriever import ( "bytes" + "context" "errors" "os" "path/filepath" @@ -30,7 +31,7 @@ func TestCompressedJSONLinesRoundTrip(t *testing.T) { t.Fatalf("write compressed JSONL: %v", err) } - if entry.Count != len(records) || entry.SHA256 == "" { + if entry.Rows != len(records) || entry.SHA256 == "" { t.Fatalf("unexpected fragment metadata: %+v", entry) } @@ -139,7 +140,7 @@ func TestEmptyJSONLinesFragmentRoundTrip(t *testing.T) { t.Fatalf("read empty fragment: %v", err) } - if entry.Count != 0 || entry.UncompressedBytes != 0 || count != 0 { + if entry.Rows != 0 || entry.UncompressedBytes != 0 || count != 0 { t.Fatalf("unexpected empty fragment metadata: entry=%+v count=%d", entry, count) } } @@ -262,7 +263,9 @@ func TestCompressedJSONLinesWriterAbort(t *testing.T) { if err := writer.Write(FragmentNode{ID: "1"}); err != nil { t.Fatalf("write record: %v", err) } - writer.Abort() + if err := writer.Abort(); err != nil { + t.Fatalf("abort writer: %v", err) + } if _, err := os.Stat(path); !os.IsNotExist(err) { t.Fatalf("expected final path to be absent, got %v", err) @@ -275,6 +278,87 @@ func TestCompressedJSONLinesWriterAbort(t *testing.T) { } } +func TestCompressedJSONLinesWriterPrepareCommitLifecycle(t *testing.T) { + path := filepath.Join(t.TempDir(), "fragment.gz") + writer, err := newCompressedJSONLinesWriter(path, CompressionGzip, DefaultZstdLevel) + if err != nil { + t.Fatalf("open writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "1"}); err != nil { + t.Fatalf("write record: %v", err) + } + + prepared, err := writer.Prepare() + if err != nil { + t.Fatalf("prepare writer: %v", err) + } + if prepared.Metadata().Rows != 1 || prepared.Metadata().SHA256 == "" { + t.Fatalf("prepared metadata = %+v", prepared.Metadata()) + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("final path exists before commit: %v", err) + } + if _, err := os.Stat(path + ".tmp"); err != nil { + t.Fatalf("staged path missing after prepare: %v", err) + } + if err := writer.Write(FragmentNode{ID: "2"}); err == nil { + t.Fatalf("expected write after prepare to fail") + } + if _, err := writer.Prepare(); err == nil { + t.Fatalf("expected double prepare to fail") + } + if err := writer.Abort(); err == nil { + t.Fatalf("prepared fragment must be aborted through prepared handle") + } + + if err := prepared.Commit(context.Background()); err != nil { + t.Fatalf("commit fragment: %v", err) + } + if _, err := os.Stat(path); err != nil { + t.Fatalf("committed path missing: %v", err) + } + if _, err := os.Stat(path + ".tmp"); !os.IsNotExist(err) { + t.Fatalf("staged path remains after commit: %v", err) + } + if err := prepared.Commit(context.Background()); err == nil { + t.Fatalf("expected double commit to fail") + } + if err := prepared.Abort(); err == nil { + t.Fatalf("expected abort after commit to fail") + } +} + +func TestCompressedJSONLinesPreparedFragmentAbort(t *testing.T) { + path := filepath.Join(t.TempDir(), "fragment.gz") + writer, err := newCompressedJSONLinesWriter(path, CompressionGzip, DefaultZstdLevel) + if err != nil { + t.Fatalf("open writer: %v", err) + } + if err := writer.Write(FragmentNode{ID: "1"}); err != nil { + t.Fatalf("write record: %v", err) + } + prepared, err := writer.Prepare() + if err != nil { + t.Fatalf("prepare writer: %v", err) + } + + if err := prepared.Abort(); err != nil { + t.Fatalf("abort prepared fragment: %v", err) + } + if err := prepared.Abort(); err != nil { + t.Fatalf("repeat prepared abort: %v", err) + } + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("final path exists after abort: %v", err) + } + if _, err := os.Stat(path + ".tmp"); !os.IsNotExist(err) { + t.Fatalf("staged path exists after abort: %v", err) + } + if err := prepared.Commit(context.Background()); err == nil { + t.Fatalf("expected commit after abort to fail") + } +} + func writeCompressedPayload(t *testing.T, path string, codec CompressionCodec, payload string) { t.Helper() diff --git a/retriever/doc.go b/retriever/doc.go index e654c180..e73f985b 100644 --- a/retriever/doc.go +++ b/retriever/doc.go @@ -1,5 +1,6 @@ // Package retriever exports and imports Dawgs graph databases using the -// manifest-based retriever collection format. +// manifest-based retriever collection format. Dumps may additionally emit a +// Parquet sidecar while retaining JSONL as the loadable collection. // // The primary database operations accept an already-open graph.Database. The // package also exposes the collection manifest and JSONL record structs for diff --git a/retriever/dump.go b/retriever/dump.go index 3ae3194a..56b7c3c9 100644 --- a/retriever/dump.go +++ b/retriever/dump.go @@ -4,29 +4,33 @@ import ( "context" "fmt" "log/slog" - "os" - "path" - "path/filepath" - "sort" "time" "github.com/specterops/dawgs/graph" - "github.com/specterops/dawgs/query" ) type DumpResult struct { - Manifest Manifest - ManifestPath string - NodeCount int64 - EdgeCount int64 + Manifest Manifest + ManifestPath string + ParquetManifestPath string + ParquetSuccessPath string + NodeCount int64 + EdgeCount int64 } -type graphEntitySnapshot struct { - NodeCount int64 - EdgeCount int64 +type dumpOverrides struct { + workspace collectionWorkspace + publisher collectionPublisher + parquet parquetPublisher + nodeOutput shardOutput[normalizedNode] + edgeOutput shardOutput[normalizedEdge] } func Dump(ctx context.Context, db graph.Database, driverName string, targets []GraphTarget, options DumpOptions) (DumpResult, error) { + return runDump(ctx, newDatabaseGraphSource(db), driverName, targets, options, dumpOverrides{}) +} + +func runDump(ctx context.Context, source graphSource, driverName string, targets []GraphTarget, options DumpOptions, overrides dumpOverrides) (DumpResult, error) { if err := options.validate(); err != nil { return DumpResult{}, err } @@ -43,6 +47,7 @@ func Dump(ctx context.Context, db graph.Database, driverName string, targets []G slog.Int("batch_size", options.BatchSize), slog.Int("shard_size", options.ShardSize), slog.String("compression", string(options.Compression)), + slog.Bool("parquet", options.Parquet), slog.String("scrub", string(options.Scrub)), ) options.Progress.emit(ProgressEvent{ @@ -54,32 +59,24 @@ func Dump(ctx context.Context, db graph.Database, driverName string, targets []G BatchSize: options.BatchSize, ShardSize: options.ShardSize, Compression: options.Compression, + Parquet: options.Parquet, Scrub: options.Scrub, }) - var ( - activeScrubber *scrubber - - scrubInfo = ScrubMetadata{ - Mode: ScrubNone, - NodeActionCounts: map[string]int{}, - EdgeActionCounts: map[string]int{}, - } - ) - if options.Scrub == ScrubFull { - nextScrubber, err := newScrubber(options.ScrubConfig, options.Salt) - if err != nil { - return DumpResult{}, err - } - activeScrubber = nextScrubber - scrubInfo = activeScrubber.metadata() + transform, err := newTransformSession(options) + if err != nil { + return DumpResult{}, err } slog.Info("retriever dump preparing output directory", slog.String("output_dir", options.OutputDir), slog.Bool("force", options.Force), ) - if err := prepareOutputDirectory(options.OutputDir, options.Force); err != nil { + workspace := overrides.workspace + if workspace == nil { + workspace = newLocalCollectionWorkspace(options.OutputDir, options.Force) + } + if err := workspace.Prepare(ctx); err != nil { return DumpResult{}, err } @@ -92,8 +89,31 @@ func Dump(ctx context.Context, db graph.Database, driverName string, targets []G OutputDir: options.OutputDir, }) - nextManifest := newManifest(driverName, options.Compression, options.ZstdLevel, scrubInfo, len(targets)) - nextMetrics := newMetricsManifest(len(targets)) + publisher := overrides.publisher + if publisher == nil { + publisher = newJSONLCollectionPublisher(workspace, driverName, options, transform.Metadata(), len(targets)) + } + parquetPublisher := overrides.parquet + if options.Parquet && parquetPublisher == nil { + parquetPublisher = newParquetCollectionPublisher(workspace, len(targets)) + } + nodeOutput := overrides.nodeOutput + if nodeOutput == nil { + sinks := []shardSink[normalizedNode]{newJSONLShardSink(newJSONLNodeSinkInWorkspace(options, workspace))} + if options.Parquet { + sinks = append(sinks, newParquetShardSink[normalizedNode](newParquetNodeSinkInWorkspace(workspace))) + } + nodeOutput = newShardSinkSet(sinks...) + } + edgeOutput := overrides.edgeOutput + if edgeOutput == nil { + sinks := []shardSink[normalizedEdge]{newJSONLShardSink(newJSONLEdgeSinkInWorkspace(options, workspace))} + if options.Parquet { + sinks = append(sinks, newParquetShardSink[normalizedEdge](newParquetEdgeSinkInWorkspace(workspace))) + } + edgeOutput = newShardSinkSet(sinks...) + } + var totalNodes, totalEdges int64 for targetIndex, target := range targets { @@ -112,17 +132,15 @@ func Dump(ctx context.Context, db graph.Database, driverName string, targets []G OutputDir: options.OutputDir, }) - graphEntry, schemaEntry, metricsEntry, err := dumpGraph(ctx, db, target, options, activeScrubber) + graphEntry, schemaEntry, metricsEntry, err := dumpGraph(ctx, source, target, options, transform, nodeOutput, edgeOutput, parquetPublisher) if err != nil { return DumpResult{}, err } - nextManifest.Graphs = append(nextManifest.Graphs, graphEntry) - nextManifest.Schema.Graphs = append(nextManifest.Schema.Graphs, schemaEntry) - nextMetrics.Graphs = append(nextMetrics.Graphs, metricsEntry) - - addActionCounts(nextManifest.Scrub.NodeActionCounts, graphEntry.NodeActionCounts) - addActionCounts(nextManifest.Scrub.EdgeActionCounts, graphEntry.EdgeActionCounts) + publisher.AddGraph(graphEntry, schemaEntry, metricsEntry) + if parquetPublisher != nil { + parquetPublisher.AddGraph(graphEntry.Name, graphEntry.NodeCount, graphEntry.EdgeCount) + } totalNodes += graphEntry.NodeCount totalEdges += graphEntry.EdgeCount @@ -148,23 +166,20 @@ func Dump(ctx context.Context, db graph.Database, driverName string, targets []G }) } - nextManifest.Metrics = &nextMetrics - slog.Info("retriever dump writing manifest", slog.String("output_dir", options.OutputDir), slog.Int64("node_count", totalNodes), slog.Int64("edge_count", totalEdges), ) - if err := writeManifest(options.OutputDir, nextManifest); err != nil { + publication, parquetPublication, err := publishDumpOutputs(ctx, publisher, parquetPublisher) + if err != nil { return DumpResult{}, err } - manifestPath := filepath.Join(options.OutputDir, manifestFileName) - slog.Info("retriever dump completed", slog.String("driver", driverName), slog.Int("graph_count", len(targets)), - slog.String("manifest", manifestPath), + slog.String("manifest", publication.Path), slog.Int64("node_count", totalNodes), slog.Int64("edge_count", totalEdges), slog.Duration("wall_elapsed", time.Since(startedAt)), @@ -181,46 +196,16 @@ func Dump(ctx context.Context, db graph.Database, driverName string, targets []G }) return DumpResult{ - Manifest: nextManifest, - ManifestPath: manifestPath, - NodeCount: totalNodes, - EdgeCount: totalEdges, + Manifest: publication.Manifest, + ManifestPath: publication.Path, + ParquetManifestPath: parquetPublication.ManifestPath, + ParquetSuccessPath: parquetPublication.SuccessPath, + NodeCount: totalNodes, + EdgeCount: totalEdges, }, nil } -func prepareOutputDirectory(outputDir string, force bool) error { - info, err := os.Stat(outputDir) - if err == nil { - if !info.IsDir() { - return fmt.Errorf("output path %q exists and is not a directory", outputDir) - } - - entries, err := os.ReadDir(outputDir) - if err != nil { - return fmt.Errorf("read output directory: %w", err) - } - - if len(entries) > 0 { - if !force { - return fmt.Errorf("output directory %q is not empty; pass -force to replace it", outputDir) - } - - if err := os.RemoveAll(outputDir); err != nil { - return fmt.Errorf("replace output directory: %w", err) - } - } - } else if !os.IsNotExist(err) { - return fmt.Errorf("inspect output directory: %w", err) - } - - if err := os.MkdirAll(outputDir, 0o755); err != nil { - return fmt.Errorf("create output directory: %w", err) - } - - return nil -} - -func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, options DumpOptions, activeScrubber *scrubber) (GraphManifest, GraphSchemaMetadata, GraphMetrics, error) { +func dumpGraph(ctx context.Context, source graphSource, target GraphTarget, options DumpOptions, transform transformSession, nodeOutput shardOutput[normalizedNode], edgeOutput shardOutput[normalizedEdge], parquetPublisher parquetPublisher) (GraphManifest, GraphSchemaMetadata, GraphMetrics, error) { targetGraph := graph.Graph{ Name: target.Name, } @@ -229,7 +214,7 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio slog.Info("retriever dump counting graph entities", slog.String("graph", target.Name), ) - entitySnapshot, err := countGraphEntitySnapshot(ctx, db, targetGraph) + entitySnapshot, err := source.Inventory(ctx, targetGraph) if err != nil { return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err } @@ -249,7 +234,7 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio Elapsed: time.Since(countStartedAt), }) - if activeScrubber != nil { + if transform.NeedsPreparation() { scrubStartedAt := time.Now() slog.Info("retriever dump scrub pre-pass started", slog.String("graph", target.Name), @@ -265,7 +250,7 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio BatchSize: options.BatchSize, }) - observedNodes, err := collectScrubRegistry(ctx, db, targetGraph, options.BatchSize, activeScrubber, entitySnapshot, options.Progress, options.ProgressInterval) + observedNodes, err := prepareTransformSession(ctx, source, targetGraph, options.BatchSize, transform, entitySnapshot, options.Progress, options.ProgressInterval) if err != nil { return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err } @@ -288,16 +273,11 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio } graphEntry := GraphManifest{ - Name: target.Name, - NodeCount: entitySnapshot.NodeCount, - EdgeCount: entitySnapshot.EdgeCount, - NodeActionCounts: map[string]int{}, - EdgeActionCounts: map[string]int{}, + Name: target.Name, + NodeCount: entitySnapshot.NodeCount, + EdgeCount: entitySnapshot.EdgeCount, } - - nodeKinds := map[string]struct{}{} - edgeKinds := map[string]struct{}{} - metricsBuilder := newMetricsBuilder(target.Name, entitySnapshot.NodeCount) + observer := newGraphObserver(target.Name, entitySnapshot.NodeCount) nodeStartedAt := time.Now() slog.Info("retriever dump node phase started", @@ -316,7 +296,20 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio ShardSize: options.ShardSize, }) - nodeFiles, err := dumpNodePhase(ctx, db, targetGraph, options, activeScrubber, nodeKinds, graphEntry.NodeActionCounts, entitySnapshot, metricsBuilder) + nodeFiles, err := dumpEntityPhase( + ctx, + source.Nodes(targetGraph, entitySnapshot.NodeCount, options.BatchSize), + dumpPhaseConfig{ + Graph: targetGraph.Name, + Phase: PhaseNodes, + Planned: entitySnapshot.NodeCount, + Options: options, + }, + transform.TransformNodes, + observer.ObserveNode, + nodeOutput, + parquetPublisher, + ) if err != nil { return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err } @@ -358,7 +351,20 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio ShardSize: options.ShardSize, }) - edgeFiles, err := dumpEdgePhase(ctx, db, targetGraph, options, activeScrubber, edgeKinds, graphEntry.EdgeActionCounts, entitySnapshot, metricsBuilder) + edgeFiles, err := dumpEntityPhase( + ctx, + source.Edges(targetGraph, entitySnapshot.EdgeCount, options.BatchSize), + dumpPhaseConfig{ + Graph: targetGraph.Name, + Phase: PhaseEdges, + Planned: entitySnapshot.EdgeCount, + Options: options, + }, + transform.TransformEdges, + observer.ObserveEdge, + edgeOutput, + parquetPublisher, + ) if err != nil { return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, err } @@ -383,29 +389,31 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio EntitiesPerSecond: perSecond(fileTotal(edgeFiles), time.Since(edgeStartedAt)), }) - if fileTotal(nodeFiles) != entitySnapshot.NodeCount { + if observer.nodeCount != entitySnapshot.NodeCount { return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, EntityCountMismatchError{ Operation: OperationDump, Graph: target.Name, Phase: PhaseNodes, Expected: entitySnapshot.NodeCount, - Actual: fileTotal(nodeFiles), - Message: fmt.Sprintf("dumped %d nodes for graph %q but counted %d at scan start; source graph changed during dump or the ID scan was inconsistent", fileTotal(nodeFiles), target.Name, entitySnapshot.NodeCount), + Actual: observer.nodeCount, + Message: fmt.Sprintf("dumped %d nodes for graph %q but counted %d at scan start; source graph changed during dump or the ID scan was inconsistent", observer.nodeCount, target.Name, entitySnapshot.NodeCount), } } - if fileTotal(edgeFiles) != entitySnapshot.EdgeCount { + if observer.edgeCount != entitySnapshot.EdgeCount { return GraphManifest{}, GraphSchemaMetadata{}, GraphMetrics{}, EntityCountMismatchError{ Operation: OperationDump, Graph: target.Name, Phase: PhaseEdges, Expected: entitySnapshot.EdgeCount, - Actual: fileTotal(edgeFiles), - Message: fmt.Sprintf("dumped %d relationships for graph %q but counted %d at scan start; source graph changed during dump or the ID scan was inconsistent", fileTotal(edgeFiles), target.Name, entitySnapshot.EdgeCount), + Actual: observer.edgeCount, + Message: fmt.Sprintf("dumped %d relationships for graph %q but counted %d at scan start; source graph changed during dump or the ID scan was inconsistent", observer.edgeCount, target.Name, entitySnapshot.EdgeCount), } } - metricsEntry := metricsBuilder.finalize() + graphEntry.NodeActionCounts = observer.nodeActionCounts + graphEntry.EdgeActionCounts = observer.edgeActionCounts + metricsEntry := observer.Metrics() slog.Info("retriever dump metrics fingerprint computed", slog.String("graph", target.Name), @@ -414,19 +422,15 @@ func dumpGraph(ctx context.Context, db graph.Database, target GraphTarget, optio slog.Int64("edge_count", metricsEntry.EdgeCount), ) - schemaEntry := GraphSchemaMetadata{ - Name: target.Name, - NodeKinds: stringsFromKindSet(nodeKinds), - EdgeKinds: stringsFromKindSet(edgeKinds), - } + schemaEntry := observer.Schema() return graphEntry, schemaEntry, metricsEntry, nil } -func collectScrubRegistry(ctx context.Context, db graph.Database, targetGraph graph.Graph, batchSize int, activeScrubber *scrubber, entitySnapshot graphEntitySnapshot, progress ProgressFunc, progressInterval int64) (int64, error) { - processed, err := scanDatabaseNodesWithProgressInterval(ctx, db, targetGraph, entitySnapshot.NodeCount, batchSize, progressInterval, func(nodes []*graph.Node) error { +func prepareTransformSession(ctx context.Context, source graphSource, targetGraph graph.Graph, batchSize int, transform transformSession, entitySnapshot graphEntitySnapshot, progress ProgressFunc, progressInterval int64) (int64, error) { + processed, err := runFaucetWithProgress(ctx, source.Nodes(targetGraph, entitySnapshot.NodeCount, batchSize), entitySnapshot.NodeCount, progressInterval, func(nodes []*graph.Node) error { for _, node := range nodes { - activeScrubber.observeNode(node.Properties.MapOrEmpty()) + transform.PrepareNode(node) } return nil }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { @@ -439,317 +443,71 @@ func collectScrubRegistry(ctx context.Context, db graph.Database, targetGraph gr return processed, nil } -func countGraphEntities(ctx context.Context, db graph.Database, targetGraph graph.Graph) (int64, int64, error) { - entitySnapshot, err := countGraphEntitySnapshot(ctx, db, targetGraph) - if err != nil { - return 0, 0, err - } - - return entitySnapshot.NodeCount, entitySnapshot.EdgeCount, nil -} - -func countGraphEntitySnapshot(ctx context.Context, db graph.Database, targetGraph graph.Graph) (graphEntitySnapshot, error) { - var entitySnapshot graphEntitySnapshot - if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { - tx = tx.WithGraph(targetGraph) - - var err error - if entitySnapshot.NodeCount, err = tx.Nodes().Count(); err != nil { - return fmt.Errorf("count nodes: %w", err) - } - - if entitySnapshot.EdgeCount, err = tx.Relationships().Count(); err != nil { - return fmt.Errorf("count relationships: %w", err) - } - - return nil - }); err != nil { - return graphEntitySnapshot{}, err - } - - return entitySnapshot, nil +type dumpPhaseConfig struct { + Graph string + Phase Phase + Planned int64 + Options DumpOptions } -func dumpNodePhase(ctx context.Context, db graph.Database, targetGraph graph.Graph, options DumpOptions, activeScrubber *scrubber, nodeKinds map[string]struct{}, graphActionCounts map[string]int, entitySnapshot graphEntitySnapshot, metricsBuilder *metricsBuilder) ([]FileManifest, error) { - if entitySnapshot.NodeCount == 0 { +func dumpEntityPhase[S, T any]( + ctx context.Context, + source faucet[S], + config dumpPhaseConfig, + transform func([]S) transformedBatch[T], + observe func(T, map[string]int) error, + output shardOutput[T], + parquetPublisher parquetPublisher, +) ([]FileManifest, error) { + if config.Planned == 0 { return nil, nil } - var ( - files []FileManifest - fragmentWriter *compressedJSONLinesWriter - fragmentRelativePath string - shardActionCounts = map[string]int{} - shardNumber = 1 - ) - - flush := func() error { - if fragmentWriter == nil { - return nil - } - - fileEntry, err := closeFragmentWriter(fragmentWriter, fragmentRelativePath, PhaseNodes, shardActionCounts) - fragmentWriter = nil - if err != nil { - return err - } - - files = append(files, fileEntry) - shardActionCounts = map[string]int{} - shardNumber++ - - return nil - } - - if _, err := scanDatabaseNodesWithProgressInterval(ctx, db, targetGraph, entitySnapshot.NodeCount, options.BatchSize, options.ProgressInterval, func(nodes []*graph.Node) error { - for _, node := range nodes { - if fragmentWriter == nil { - nextWriter, nextRelativePath, err := openFragmentWriter(options.OutputDir, targetGraph.Name, PhaseNodes, shardNumber, options) - if err != nil { - return err - } - - fragmentWriter = nextWriter - fragmentRelativePath = nextRelativePath - } - - kinds := node.Kinds.Strings() - sort.Strings(kinds) - addKindsToSet(nodeKinds, kinds) - - properties := node.Properties.MapOrEmpty() - if activeScrubber != nil { - var actionCounts map[string]int - properties, actionCounts = activeScrubber.scrubProperties(properties) - addActionCounts(shardActionCounts, actionCounts) - addActionCounts(graphActionCounts, actionCounts) - } - - item := FragmentNode{ - ID: node.ID.String(), - Kinds: kinds, - Properties: properties, - } - - if err := metricsBuilder.observeFragmentNode(item); err != nil { - return err - } - - if err := fragmentWriter.Write(item); err != nil { - return err - } - - if fragmentWriter.Count() >= options.ShardSize { - if err := flush(); err != nil { - return err - } - } - } - - return nil - }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { - return logRetrieverEntityProgressInterval("retriever dump node phase progress", targetGraph.Name, PhaseNodes, processed, entitySnapshot.NodeCount, startedAt, nextProgressAt, options.Progress, options.ProgressInterval) - }); err != nil { - if fragmentWriter != nil { - fragmentWriter.Abort() - } - - return nil, err - } - - if err := flush(); err != nil { + sharder, err := newLogicalSharder[T](config.Graph, config.Phase, config.Options.ShardSize) + if err != nil { return nil, err } - return files, nil -} - -func dumpEdgePhase(ctx context.Context, db graph.Database, targetGraph graph.Graph, options DumpOptions, activeScrubber *scrubber, edgeKinds map[string]struct{}, graphActionCounts map[string]int, entitySnapshot graphEntitySnapshot, metricsBuilder *metricsBuilder) ([]FileManifest, error) { - if entitySnapshot.EdgeCount == 0 { - return nil, nil - } - - var ( - files []FileManifest - fragmentWriter *compressedJSONLinesWriter - fragmentRelativePath string - shardActionCounts = map[string]int{} - shardNumber = 1 - ) - - flush := func() error { - if fragmentWriter == nil { - return nil - } - - fileEntry, err := closeFragmentWriter(fragmentWriter, fragmentRelativePath, PhaseEdges, shardActionCounts) - fragmentWriter = nil - if err != nil { - return err + var files []FileManifest + receiver := newShardOutputReceiver(ctx, output, func(summary shardSummary, committed committedShard) error { + files = append(files, newJSONLFileManifest(summary, committed.JSONL)) + if committed.Parquet != nil { + if parquetPublisher == nil { + return fmt.Errorf("received Parquet %s fragment without a publisher", config.Phase) + } + parquetPublisher.AddFragment(summary, *committed.Parquet) } - - files = append(files, fileEntry) - shardActionCounts = map[string]int{} - shardNumber++ - return nil - } - - if _, err := scanDatabaseRelationshipsWithProgressInterval(ctx, db, targetGraph, entitySnapshot.EdgeCount, options.BatchSize, options.ProgressInterval, func(relationships []*graph.Relationship) error { - for _, relationship := range relationships { - if fragmentWriter == nil { - nextWriter, nextRelativePath, err := openFragmentWriter(options.OutputDir, targetGraph.Name, PhaseEdges, shardNumber, options) - if err != nil { - return err - } - - fragmentWriter = nextWriter - fragmentRelativePath = nextRelativePath - } - - kind := "" - if relationship.Kind != nil { - kind = relationship.Kind.String() - edgeKinds[kind] = struct{}{} - } - - properties := relationship.Properties.MapOrEmpty() - if activeScrubber != nil { - var actionCounts map[string]int - properties, actionCounts = activeScrubber.scrubProperties(properties) - addActionCounts(shardActionCounts, actionCounts) - addActionCounts(graphActionCounts, actionCounts) - } - - item := FragmentEdge{ - StartID: relationship.StartID.String(), - EndID: relationship.EndID.String(), - Kind: kind, - Properties: properties, - } - - if err := metricsBuilder.observeFragmentEdge(item); err != nil { - return err - } + }) - if err := fragmentWriter.Write(item); err != nil { + if _, err := runFaucetWithProgress(ctx, source, config.Planned, config.Options.ProgressInterval, func(sourceBatch []S) error { + batch := transform(sourceBatch) + for index, record := range batch.Records { + if err := observe(record, batch.ActionCounts[index]); err != nil { return err } - - if fragmentWriter.Count() >= options.ShardSize { - if err := flush(); err != nil { - return err - } - } } - return nil + return sharder.Add(batch, receiver) }, func(processed int64, startedAt time.Time, nextProgressAt int64) int64 { - return logRetrieverEntityProgressInterval("retriever dump edge phase progress", targetGraph.Name, PhaseEdges, processed, entitySnapshot.EdgeCount, startedAt, nextProgressAt, options.Progress, options.ProgressInterval) + message := fmt.Sprintf("retriever dump %s phase progress", dumpPhaseLabel(config.Phase)) + return logRetrieverEntityProgressInterval(message, config.Graph, config.Phase, processed, config.Planned, startedAt, nextProgressAt, config.Options.Progress, config.Options.ProgressInterval) }); err != nil { - if fragmentWriter != nil { - fragmentWriter.Abort() - } - - return nil, err + return nil, cleanupOnError(err, receiver.Abort) } - if err := flush(); err != nil { - return nil, err + if err := sharder.Flush(receiver); err != nil { + return nil, cleanupOnError(err, receiver.Abort) } return files, nil } -func writeNodeFragment(outputDir, graphName string, shardNumber int, options DumpOptions, items []FragmentNode, actionCounts map[string]int) (FileManifest, error) { - relativePath, err := fragmentPath(graphName, PhaseNodes, shardNumber, options.Compression) - if err != nil { - return FileManifest{}, err - } - - absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) - fileEntry, err := writeCompressedJSONLines(absolutePath, options.Compression, options.ZstdLevel, items) - if err != nil { - return FileManifest{}, err - } - - fileEntry.Phase = PhaseNodes - fileEntry.Path = relativePath - fileEntry.Count = len(items) - fileEntry.ActionCounts = cloneActionCounts(actionCounts) - - return fileEntry, nil -} - -func writeEdgeFragment(outputDir, graphName string, shardNumber int, options DumpOptions, items []FragmentEdge, actionCounts map[string]int) (FileManifest, error) { - relativePath, err := fragmentPath(graphName, PhaseEdges, shardNumber, options.Compression) - if err != nil { - return FileManifest{}, err - } - - absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) - fileEntry, err := writeCompressedJSONLines(absolutePath, options.Compression, options.ZstdLevel, items) - if err != nil { - return FileManifest{}, err - } - - fileEntry.Phase = PhaseEdges - fileEntry.Path = relativePath - fileEntry.Count = len(items) - fileEntry.ActionCounts = cloneActionCounts(actionCounts) - - return fileEntry, nil -} - -func fragmentPath(graphName string, fragmentPhase Phase, shardNumber int, codec CompressionCodec) (string, error) { - if shardNumber <= 0 { - return "", fmt.Errorf("shard number must be > 0") +func dumpPhaseLabel(phase Phase) string { + if phase == PhaseEdges { + return "edge" } - - extension, err := compressionExtension(codec) - if err != nil { - return "", err - } - - var prefix string - switch fragmentPhase { - case PhaseNodes: - prefix = "nodes" - case PhaseEdges: - prefix = "edges" - default: - return "", fmt.Errorf("unsupported fragment phase %q", fragmentPhase) - } - - return path.Join("graphs", graphDirectoryName(graphName), fmt.Sprintf("%s-%06d.jsonl%s", prefix, shardNumber, extension)), nil -} - -func openFragmentWriter(outputDir, graphName string, fragmentPhase Phase, shardNumber int, options DumpOptions) (*compressedJSONLinesWriter, string, error) { - relativePath, err := fragmentPath(graphName, fragmentPhase, shardNumber, options.Compression) - if err != nil { - return nil, "", err - } - - absolutePath := filepath.Join(outputDir, filepath.FromSlash(relativePath)) - writer, err := newCompressedJSONLinesWriter(absolutePath, options.Compression, options.ZstdLevel) - if err != nil { - return nil, "", err - } - - return writer, relativePath, nil -} - -func closeFragmentWriter(writer *compressedJSONLinesWriter, relativePath string, fragmentPhase Phase, actionCounts map[string]int) (FileManifest, error) { - fileEntry, err := writer.Close() - if err != nil { - return FileManifest{}, err - } - - fileEntry.Phase = fragmentPhase - fileEntry.Path = relativePath - fileEntry.ActionCounts = cloneActionCounts(actionCounts) - - return fileEntry, nil + return "node" } func fileTotal(files []FileManifest) int64 { @@ -760,63 +518,3 @@ func fileTotal(files []FileManifest) int64 { return total } - -func readDatabaseNodes(ctx context.Context, db graph.Database, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Node, error) { - var nodes []*graph.Node - if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { - tx = tx.WithGraph(targetGraph) - nodeQuery := tx.Nodes(). - OrderBy(query.NodeID()). - Limit(batchSize) - - if hasAfterID { - nodeQuery = nodeQuery.Filter(query.GreaterThan(query.NodeID(), afterID)) - } - - return nodeQuery.Fetch(func(cursor graph.Cursor[*graph.Node]) error { - for node := range cursor.Chan() { - nodes = append(nodes, node) - } - - return cursor.Error() - }) - }); err != nil { - if hasAfterID { - return nil, fmt.Errorf("read node batch after ID %d: %w", afterID.Uint64(), err) - } - - return nil, fmt.Errorf("read initial node batch: %w", err) - } - - return nodes, nil -} - -func readDatabaseRelationships(ctx context.Context, db graph.Database, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Relationship, error) { - var relationships []*graph.Relationship - if err := db.ReadTransaction(ctx, func(tx graph.Transaction) error { - tx = tx.WithGraph(targetGraph) - relationshipQuery := tx.Relationships(). - OrderBy(query.RelationshipID()). - Limit(batchSize) - - if hasAfterID { - relationshipQuery = relationshipQuery.Filter(query.GreaterThan(query.RelationshipID(), afterID)) - } - - return relationshipQuery.Fetch(func(cursor graph.Cursor[*graph.Relationship]) error { - for relationship := range cursor.Chan() { - relationships = append(relationships, relationship) - } - - return cursor.Error() - }) - }); err != nil { - if hasAfterID { - return nil, fmt.Errorf("read relationship batch after ID %d: %w", afterID.Uint64(), err) - } - - return nil, fmt.Errorf("read initial relationship batch: %w", err) - } - - return relationships, nil -} diff --git a/retriever/dump_characterization_test.go b/retriever/dump_characterization_test.go new file mode 100644 index 00000000..d3c9a994 --- /dev/null +++ b/retriever/dump_characterization_test.go @@ -0,0 +1,670 @@ +package retriever + +import ( + "context" + "errors" + "fmt" + "os" + "path/filepath" + "reflect" + "sort" + "strings" + "testing" + + cypherModel "github.com/specterops/dawgs/cypher/models/cypher" + "github.com/specterops/dawgs/graph" +) + +func TestDumpCharacterizationJSONLContract(t *testing.T) { + database := newScriptedDumpDatabase(map[string]*scriptedDumpGraph{ + "graph/name": { + nodes: []*graph.Node{ + graph.NewNode(3, graph.AsProperties(map[string]any{"note": "third"}), graph.StringKind("User")), + graph.NewNode(1, graph.AsProperties(map[string]any{"enabled": true, "name": "alice"}), graph.StringKind("User"), graph.StringKind("Admin")), + graph.NewNode(2, nil, graph.StringKind("Computer")), + }, + relationships: []*graph.Relationship{ + graph.NewRelationship(11, 2, 3, nil, nil), + graph.NewRelationship(10, 1, 2, graph.AsProperties(map[string]any{"route": "north"}), graph.StringKind("AdminTo")), + }, + }, + "empty": {}, + }) + + outputDir := t.TempDir() + var progressMessages []string + result, err := Dump(context.Background(), database, "scripted", []GraphTarget{ + {Name: "graph/name"}, + {Name: "empty"}, + }, DumpOptions{ + OutputDir: outputDir, + Scrub: ScrubNone, + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + ShardSize: 2, + BatchSize: 3, + Progress: func(event ProgressEvent) { + progressMessages = append(progressMessages, event.Message) + }, + }) + if err != nil { + t.Fatalf("dump: %v", err) + } + + if result.NodeCount != 3 || result.EdgeCount != 2 { + t.Fatalf("dump counts: nodes=%d edges=%d", result.NodeCount, result.EdgeCount) + } + + manifest := result.Manifest + if _, offset := manifest.GeneratedAt.Zone(); offset != 0 { + t.Fatalf("manifest generated_at is not UTC: %s", manifest.GeneratedAt) + } + if manifest.Format != manifestFormat || manifest.Driver != "scripted" || manifest.Source.GraphCount != 2 || manifest.IDStrategy != idStrategy { + t.Fatalf("manifest identity changed: %+v", manifest) + } + if manifest.Compression != CompressionGzip || manifest.CompressionLevel != DefaultZstdLevel { + t.Fatalf("manifest compression changed: %+v", manifest) + } + if !reflect.DeepEqual(manifest.Scrub, ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + }) { + t.Fatalf("manifest scrub metadata changed: %+v", manifest.Scrub) + } + + expectedGraphs := []GraphManifest{ + { + Name: "graph/name", + NodeCount: 3, + EdgeCount: 2, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + Files: []FileManifest{ + {Phase: PhaseNodes, Path: "graphs/graph%2Fname/nodes-000001.jsonl.gz", Count: 2, CompressedBytes: 115, UncompressedBytes: 113, SHA256: "6d123f3d464e1737cff426cbe890972d1ec8abaecbd15a2739c203ffa43b97a4", ActionCounts: map[string]int{}}, + {Phase: PhaseNodes, Path: "graphs/graph%2Fname/nodes-000002.jsonl.gz", Count: 1, CompressedBytes: 81, UncompressedBytes: 58, SHA256: "a5976778349131421837eb8a69aab6e1efd75d607d50bc5f233ac4107809f739", ActionCounts: map[string]int{}}, + {Phase: PhaseEdges, Path: "graphs/graph%2Fname/edges-000001.jsonl.gz", Count: 2, CompressedBytes: 107, UncompressedBytes: 118, SHA256: "748baddea892c36540aadf75f6412b57bd1e7e4dfe961953f8dcdacf2d0f685d", ActionCounts: map[string]int{}}, + }, + }, + { + Name: "empty", + NodeCount: 0, + EdgeCount: 0, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + }, + } + if !reflect.DeepEqual(manifest.Graphs, expectedGraphs) { + t.Fatalf("manifest graph contract changed:\nactual: %#v\nexpected: %#v", manifest.Graphs, expectedGraphs) + } + + expectedSchema := SchemaMetadata{Graphs: []GraphSchemaMetadata{ + {Name: "graph/name", NodeKinds: []string{"Admin", "Computer", "User"}, EdgeKinds: []string{"AdminTo"}}, + {Name: "empty", NodeKinds: []string{}, EdgeKinds: []string{}}, + }} + if !reflect.DeepEqual(manifest.Schema, expectedSchema) { + t.Fatalf("manifest schema contract changed: actual=%+v expected=%+v", manifest.Schema, expectedSchema) + } + + if manifest.Metrics == nil || len(manifest.Metrics.Graphs) != 2 { + t.Fatalf("manifest metrics contract changed: %+v", manifest.Metrics) + } + if manifest.Metrics.Version != metricsVersion { + t.Fatalf("metrics version = %q", manifest.Metrics.Version) + } + expectedMetrics := []GraphMetrics{ + { + Name: "graph/name", + NodeCount: 3, + EdgeCount: 2, + NodeKindHistogram: map[string]int64{"4:User": 1, "5:Admin+4:User": 1, "8:Computer": 1}, + EdgeKindHistogram: map[string]int64{"7:AdminTo": 1, metricsNoneKind: 1}, + InDegreeHistogram: map[string]int64{"0": 1, "1": 2}, + OutDegreeHistogram: map[string]int64{"0": 1, "1": 2}, + TotalDegreeHistogram: map[string]int64{"1": 2, "2": 1}, + EndpointKindHistogram: map[string]int64{ + "10:8:Computer|2:0:|6:4:User": 1, + "14:5:Admin+4:User|9:7:AdminTo|10:8:Computer": 1, + }, + Fingerprint: "sha256:4a4a992c9c605b220356b38e042929dd69a671e7b5cd6da17f91ee355e6a1dbc", + }, + { + Name: "empty", + NodeKindHistogram: map[string]int64{}, + EdgeKindHistogram: map[string]int64{}, + InDegreeHistogram: map[string]int64{}, + OutDegreeHistogram: map[string]int64{}, + TotalDegreeHistogram: map[string]int64{}, + EndpointKindHistogram: map[string]int64{}, + Fingerprint: "sha256:3fffdda8af847c800a59c98edd6655fc37a8dd70cb1b4b74731cbd283c1af3eb", + }, + } + if !reflect.DeepEqual(manifest.Metrics.Graphs, expectedMetrics) { + t.Fatalf("metrics contract changed:\nactual: %#v\nexpected: %#v", manifest.Metrics.Graphs, expectedMetrics) + } + + assertFragmentNodes(t, outputDir, manifest.Graphs[0].Files[0], []FragmentNode{ + {ID: "1", Kinds: []string{"Admin", "User"}, Properties: map[string]any{"enabled": true, "name": "alice"}}, + {ID: "2", Kinds: []string{"Computer"}}, + }) + assertFragmentNodes(t, outputDir, manifest.Graphs[0].Files[1], []FragmentNode{ + {ID: "3", Kinds: []string{"User"}, Properties: map[string]any{"note": "third"}}, + }) + assertFragmentEdges(t, outputDir, manifest.Graphs[0].Files[2], []FragmentEdge{ + {StartID: "1", EndID: "2", Kind: "AdminTo", Properties: map[string]any{"route": "north"}}, + {StartID: "2", EndID: "3", Kind: ""}, + }) + + writtenManifest, err := readManifest(outputDir) + if err != nil { + t.Fatalf("read written manifest: %v", err) + } + if !reflect.DeepEqual(writtenManifest, result.Manifest) { + t.Fatalf("returned and written manifests differ") + } + + expectedProgressMessages := []string{ + "retriever dump started", + "retriever dump output directory ready", + "retriever dump graph started", + "retriever dump graph counts ready", + "retriever dump node phase started", + "retriever dump node phase completed", + "retriever dump edge phase started", + "retriever dump edge phase completed", + "retriever dump graph completed", + "retriever dump graph started", + "retriever dump graph counts ready", + "retriever dump node phase started", + "retriever dump node phase completed", + "retriever dump edge phase started", + "retriever dump edge phase completed", + "retriever dump graph completed", + "retriever dump completed", + } + if !reflect.DeepEqual(progressMessages, expectedProgressMessages) { + t.Fatalf("progress event order changed: actual=%v expected=%v", progressMessages, expectedProgressMessages) + } +} + +func TestDumpCharacterizationScrubRegistryGraphOrdering(t *testing.T) { + const sourceSID = "S-1-5-21-111111111-222222222-333333333-1001" + database := newScriptedDumpDatabase(map[string]*scriptedDumpGraph{ + "first": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{ + "description": "private description", + "owner_sid": sourceSID, + }), graph.StringKind("User")), + }, + }, + "second": { + nodes: []*graph.Node{ + graph.NewNode(2, graph.AsProperties(map[string]any{ + "objectid": sourceSID, + }), graph.StringKind("User")), + }, + }, + }) + + outputDir := t.TempDir() + result, err := Dump(context.Background(), database, "scripted", []GraphTarget{ + {Name: "first"}, + {Name: "second"}, + }, DumpOptions{ + OutputDir: outputDir, + Scrub: ScrubFull, + Salt: "characterization-salt", + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + ShardSize: 1, + BatchSize: 1, + }) + if err != nil { + t.Fatalf("dump: %v", err) + } + + expectedOperations := []string{ + "first:nodes:count", + "first:edges:count", + "first:nodes:fetch", + "first:nodes:fetch", + "second:nodes:count", + "second:edges:count", + "second:nodes:fetch", + "second:nodes:fetch", + } + if !reflect.DeepEqual(database.operations, expectedOperations) { + t.Fatalf("scrub graph operation order changed: actual=%v expected=%v", database.operations, expectedOperations) + } + + if !reflect.DeepEqual(result.Manifest.Scrub.NodeActionCounts, map[string]int{ + string(actionPseudonymize): 2, + string(actionRedact): 1, + }) { + t.Fatalf("collection scrub counts changed: %+v", result.Manifest.Scrub.NodeActionCounts) + } + if !reflect.DeepEqual(result.Manifest.Graphs[0].NodeActionCounts, map[string]int{ + string(actionPseudonymize): 1, + string(actionRedact): 1, + }) { + t.Fatalf("first graph scrub counts changed: %+v", result.Manifest.Graphs[0].NodeActionCounts) + } + if !reflect.DeepEqual(result.Manifest.Graphs[0].Files[0].ActionCounts, result.Manifest.Graphs[0].NodeActionCounts) { + t.Fatalf("first shard scrub counts changed: %+v", result.Manifest.Graphs[0].Files[0].ActionCounts) + } + + var firstRows []FragmentNode + readFragmentNodes(t, outputDir, result.Manifest.Graphs[0].Files[0], &firstRows) + if len(firstRows) != 1 { + t.Fatalf("first graph row count = %d", len(firstRows)) + } + if firstRows[0].Properties["description"] != "[REDACTED]" { + t.Fatalf("description scrub output changed: %#v", firstRows[0].Properties["description"]) + } + if firstRows[0].Properties["owner_sid"] == sourceSID { + t.Fatalf("owner_sid was not scrubbed: %#v", firstRows[0].Properties["owner_sid"]) + } + + var secondRows []FragmentNode + readFragmentNodes(t, outputDir, result.Manifest.Graphs[1].Files[0], &secondRows) + if len(secondRows) != 1 { + t.Fatalf("second graph row count = %d", len(secondRows)) + } + if secondRows[0].Properties["objectid"] != firstRows[0].Properties["owner_sid"] { + t.Fatalf("shared scrub registry identity changed: first=%#v second=%#v", firstRows[0].Properties["owner_sid"], secondRows[0].Properties["objectid"]) + } +} + +func TestDumpCharacterizationCountDrift(t *testing.T) { + t.Run("short scan commits fragment then reports mismatch", func(t *testing.T) { + database := newScriptedDumpDatabase(map[string]*scriptedDumpGraph{ + "default": { + nodeCount: 3, + nodes: []*graph.Node{ + graph.NewNode(1, nil, graph.StringKind("User")), + graph.NewNode(2, nil, graph.StringKind("User")), + }, + }, + }) + outputDir := t.TempDir() + + _, err := Dump(context.Background(), database, "scripted", []GraphTarget{{Name: "default"}}, characterizationDumpOptions(outputDir)) + var mismatch EntityCountMismatchError + if !errors.As(err, &mismatch) { + t.Fatalf("expected entity count mismatch, got %v", err) + } + if mismatch.Phase != PhaseNodes || mismatch.Expected != 3 || mismatch.Actual != 2 { + t.Fatalf("count mismatch contract changed: %+v", mismatch) + } + assertPathExists(t, filepath.Join(outputDir, "graphs", "default", "nodes-000001.jsonl.gz")) + assertPathDoesNotExist(t, filepath.Join(outputDir, manifestFileName)) + assertNoTemporaryArtifacts(t, outputDir) + }) + + t.Run("growth beyond inventory is capped and succeeds", func(t *testing.T) { + database := newScriptedDumpDatabase(map[string]*scriptedDumpGraph{ + "default": { + nodeCount: 2, + ignoreNodeLimit: true, + nodes: []*graph.Node{ + graph.NewNode(1, nil, graph.StringKind("User")), + graph.NewNode(2, nil, graph.StringKind("User")), + graph.NewNode(3, nil, graph.StringKind("User")), + }, + }, + }) + outputDir := t.TempDir() + + result, err := Dump(context.Background(), database, "scripted", []GraphTarget{{Name: "default"}}, characterizationDumpOptions(outputDir)) + if err != nil { + t.Fatalf("dump: %v", err) + } + if result.NodeCount != 2 || result.Manifest.Graphs[0].NodeCount != 2 { + t.Fatalf("growth cap contract changed: result=%+v", result) + } + assertFragmentNodes(t, outputDir, result.Manifest.Graphs[0].Files[0], []FragmentNode{ + {ID: "1", Kinds: []string{"User"}}, + {ID: "2", Kinds: []string{"User"}}, + }) + }) +} + +func TestDumpCharacterizationFailurePublication(t *testing.T) { + t.Run("source failure retains completed fragments", func(t *testing.T) { + sourceErr := errors.New("scripted relationship read failure") + database := newScriptedDumpDatabase(map[string]*scriptedDumpGraph{ + "default": { + nodes: []*graph.Node{ + graph.NewNode(1, nil, graph.StringKind("User")), + graph.NewNode(2, nil, graph.StringKind("User")), + }, + relationships: []*graph.Relationship{ + graph.NewRelationship(10, 1, 2, nil, graph.StringKind("AdminTo")), + }, + relationshipFetchErr: sourceErr, + }, + }) + outputDir := t.TempDir() + + _, err := Dump(context.Background(), database, "scripted", []GraphTarget{{Name: "default"}}, characterizationDumpOptions(outputDir)) + if !errors.Is(err, sourceErr) { + t.Fatalf("expected relationship read error, got %v", err) + } + + assertPathExists(t, filepath.Join(outputDir, "graphs", "default", "nodes-000001.jsonl.gz")) + assertPathDoesNotExist(t, filepath.Join(outputDir, "graphs", "default", "edges-000001.jsonl.gz")) + assertPathDoesNotExist(t, filepath.Join(outputDir, manifestFileName)) + assertNoTemporaryArtifacts(t, outputDir) + }) + + t.Run("encoding failure aborts active fragment", func(t *testing.T) { + database := newScriptedDumpDatabase(map[string]*scriptedDumpGraph{ + "default": { + nodes: []*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{ + "unsupported": func() {}, + }), graph.StringKind("User")), + }, + }, + }) + outputDir := t.TempDir() + + _, err := Dump(context.Background(), database, "scripted", []GraphTarget{{Name: "default"}}, characterizationDumpOptions(outputDir)) + if err == nil || !strings.Contains(err.Error(), "encode JSONL record 1") { + t.Fatalf("expected JSONL encoding error, got %v", err) + } + + assertPathDoesNotExist(t, filepath.Join(outputDir, "graphs", "default", "nodes-000001.jsonl.gz")) + assertPathDoesNotExist(t, filepath.Join(outputDir, manifestFileName)) + assertNoTemporaryArtifacts(t, outputDir) + }) +} + +func characterizationDumpOptions(outputDir string) DumpOptions { + return DumpOptions{ + OutputDir: outputDir, + Scrub: ScrubNone, + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + ShardSize: 10, + BatchSize: 10, + } +} + +func assertFragmentNodes(t *testing.T, outputDir string, fileEntry FileManifest, expected []FragmentNode) { + t.Helper() + var actual []FragmentNode + readFragmentNodes(t, outputDir, fileEntry, &actual) + if !reflect.DeepEqual(actual, expected) { + t.Fatalf("node fragment %q changed:\nactual: %#v\nexpected: %#v", fileEntry.Path, actual, expected) + } +} + +func readFragmentNodes(t *testing.T, outputDir string, fileEntry FileManifest, target *[]FragmentNode) { + t.Helper() + count, err := readCompressedJSONLines(filepath.Join(outputDir, filepath.FromSlash(fileEntry.Path)), CompressionGzip, func(row FragmentNode) error { + *target = append(*target, row) + return nil + }) + if err != nil { + t.Fatalf("read node fragment %q: %v", fileEntry.Path, err) + } + if count != fileEntry.Count { + t.Fatalf("node fragment %q count=%d manifest=%d", fileEntry.Path, count, fileEntry.Count) + } +} + +func assertFragmentEdges(t *testing.T, outputDir string, fileEntry FileManifest, expected []FragmentEdge) { + t.Helper() + var actual []FragmentEdge + count, err := readCompressedJSONLines(filepath.Join(outputDir, filepath.FromSlash(fileEntry.Path)), CompressionGzip, func(row FragmentEdge) error { + actual = append(actual, row) + return nil + }) + if err != nil { + t.Fatalf("read edge fragment %q: %v", fileEntry.Path, err) + } + if count != fileEntry.Count { + t.Fatalf("edge fragment %q count=%d manifest=%d", fileEntry.Path, count, fileEntry.Count) + } + if !reflect.DeepEqual(actual, expected) { + t.Fatalf("edge fragment %q changed:\nactual: %#v\nexpected: %#v", fileEntry.Path, actual, expected) + } +} + +func assertPathExists(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path); err != nil { + t.Fatalf("expected path %q: %v", path, err) + } +} + +func assertPathDoesNotExist(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("expected path %q to be absent, got %v", path, err) + } +} + +func assertNoTemporaryArtifacts(t *testing.T, root string) { + t.Helper() + if err := filepath.WalkDir(root, func(path string, entry os.DirEntry, err error) error { + if err != nil { + return err + } + if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".tmp") { + return fmt.Errorf("temporary artifact remains at %s", path) + } + return nil + }); err != nil { + t.Fatal(err) + } +} + +type scriptedDumpGraph struct { + nodeCount int64 + relationshipCount int64 + nodes []*graph.Node + relationships []*graph.Relationship + ignoreNodeLimit bool + ignoreRelationshipLimit bool + nodeFetchErr error + relationshipFetchErr error +} + +type scriptedDumpDatabase struct { + graph.Database + graphs map[string]*scriptedDumpGraph + operations []string +} + +func newScriptedDumpDatabase(graphs map[string]*scriptedDumpGraph) *scriptedDumpDatabase { + for _, graphFixture := range graphs { + if graphFixture.nodeCount == 0 { + graphFixture.nodeCount = int64(len(graphFixture.nodes)) + } + if graphFixture.relationshipCount == 0 { + graphFixture.relationshipCount = int64(len(graphFixture.relationships)) + } + } + return &scriptedDumpDatabase{graphs: graphs} +} + +func (s *scriptedDumpDatabase) ReadTransaction(_ context.Context, delegate graph.TransactionDelegate, _ ...graph.TransactionOption) error { + return delegate(&scriptedDumpTransaction{database: s}) +} + +type scriptedDumpTransaction struct { + graph.Transaction + database *scriptedDumpDatabase + graphName string +} + +func (s *scriptedDumpTransaction) WithGraph(target graph.Graph) graph.Transaction { + return &scriptedDumpTransaction{database: s.database, graphName: target.Name} +} + +func (s *scriptedDumpTransaction) Nodes() graph.NodeQuery { + return &scriptedDumpNodeQuery{database: s.database, graphName: s.graphName} +} + +func (s *scriptedDumpTransaction) Relationships() graph.RelationshipQuery { + return &scriptedDumpRelationshipQuery{database: s.database, graphName: s.graphName} +} + +type scriptedDumpNodeQuery struct { + graph.NodeQuery + database *scriptedDumpDatabase + graphName string + limit int + afterID graph.ID + hasAfter bool +} + +func (s *scriptedDumpNodeQuery) OrderBy(...graph.Criteria) graph.NodeQuery { + return s +} + +func (s *scriptedDumpNodeQuery) Limit(limit int) graph.NodeQuery { + s.limit = limit + return s +} + +func (s *scriptedDumpNodeQuery) Filter(criteria graph.Criteria) graph.NodeQuery { + s.afterID = scriptedDumpAfterID(criteria) + s.hasAfter = true + return s +} + +func (s *scriptedDumpNodeQuery) Count() (int64, error) { + fixture := s.database.graphFixture(s.graphName) + s.database.recordOperation(s.graphName, PhaseNodes, "count") + return fixture.nodeCount, nil +} + +func (s *scriptedDumpNodeQuery) Fetch(delegate func(graph.Cursor[*graph.Node]) error, _ ...graph.Criteria) error { + fixture := s.database.graphFixture(s.graphName) + s.database.recordOperation(s.graphName, PhaseNodes, "fetch") + if fixture.nodeFetchErr != nil { + return fixture.nodeFetchErr + } + + rows := append([]*graph.Node(nil), fixture.nodes...) + sort.Slice(rows, func(left, right int) bool { return rows[left].ID < rows[right].ID }) + rows = filterScriptedDumpRows(rows, s.afterID, s.hasAfter, func(node *graph.Node) graph.ID { return node.ID }) + if !fixture.ignoreNodeLimit && s.limit >= 0 && len(rows) > s.limit { + rows = rows[:s.limit] + } + return delegate(newScriptedDumpCursor(rows, nil)) +} + +type scriptedDumpRelationshipQuery struct { + graph.RelationshipQuery + database *scriptedDumpDatabase + graphName string + limit int + afterID graph.ID + hasAfter bool +} + +func (s *scriptedDumpRelationshipQuery) OrderBy(...graph.Criteria) graph.RelationshipQuery { + return s +} + +func (s *scriptedDumpRelationshipQuery) Limit(limit int) graph.RelationshipQuery { + s.limit = limit + return s +} + +func (s *scriptedDumpRelationshipQuery) Filter(criteria graph.Criteria) graph.RelationshipQuery { + s.afterID = scriptedDumpAfterID(criteria) + s.hasAfter = true + return s +} + +func (s *scriptedDumpRelationshipQuery) Count() (int64, error) { + fixture := s.database.graphFixture(s.graphName) + s.database.recordOperation(s.graphName, PhaseEdges, "count") + return fixture.relationshipCount, nil +} + +func (s *scriptedDumpRelationshipQuery) Fetch(delegate func(graph.Cursor[*graph.Relationship]) error) error { + fixture := s.database.graphFixture(s.graphName) + s.database.recordOperation(s.graphName, PhaseEdges, "fetch") + if fixture.relationshipFetchErr != nil { + return fixture.relationshipFetchErr + } + + rows := append([]*graph.Relationship(nil), fixture.relationships...) + sort.Slice(rows, func(left, right int) bool { return rows[left].ID < rows[right].ID }) + rows = filterScriptedDumpRows(rows, s.afterID, s.hasAfter, func(relationship *graph.Relationship) graph.ID { return relationship.ID }) + if !fixture.ignoreRelationshipLimit && s.limit >= 0 && len(rows) > s.limit { + rows = rows[:s.limit] + } + return delegate(newScriptedDumpCursor(rows, nil)) +} + +func (s *scriptedDumpDatabase) graphFixture(name string) *scriptedDumpGraph { + fixture, found := s.graphs[name] + if !found { + panic(fmt.Sprintf("missing scripted graph %q", name)) + } + return fixture +} + +func (s *scriptedDumpDatabase) recordOperation(graphName string, phase Phase, operation string) { + s.operations = append(s.operations, fmt.Sprintf("%s:%s:%s", graphName, phase, operation)) +} + +func scriptedDumpAfterID(criteria graph.Criteria) graph.ID { + comparison, ok := criteria.(*cypherModel.Comparison) + if !ok || len(comparison.Partials) != 1 { + panic(fmt.Sprintf("unexpected scripted keyset criteria %T", criteria)) + } + parameter, ok := comparison.Partials[0].Right.(*cypherModel.Parameter) + if !ok { + panic(fmt.Sprintf("unexpected scripted keyset value %T", comparison.Partials[0].Right)) + } + value, ok := parameter.Value.(graph.ID) + if !ok { + panic(fmt.Sprintf("unexpected scripted keyset ID %T", parameter.Value)) + } + return value +} + +func filterScriptedDumpRows[T any](rows []T, afterID graph.ID, hasAfter bool, id func(T) graph.ID) []T { + if !hasAfter { + return rows + } + for index, row := range rows { + if id(row) > afterID { + return rows[index:] + } + } + return nil +} + +type scriptedDumpCursor[T any] struct { + values chan T + err error +} + +func newScriptedDumpCursor[T any](values []T, err error) *scriptedDumpCursor[T] { + valueChannel := make(chan T, len(values)) + for _, value := range values { + valueChannel <- value + } + close(valueChannel) + return &scriptedDumpCursor[T]{values: valueChannel, err: err} +} + +func (s *scriptedDumpCursor[T]) Error() error { + return s.err +} + +func (s *scriptedDumpCursor[T]) Close() {} + +func (s *scriptedDumpCursor[T]) Chan() chan T { + return s.values +} diff --git a/retriever/dump_failure_test.go b/retriever/dump_failure_test.go new file mode 100644 index 00000000..8df6a086 --- /dev/null +++ b/retriever/dump_failure_test.go @@ -0,0 +1,256 @@ +package retriever + +import ( + "context" + "errors" + "os" + "path/filepath" + "testing" + + "github.com/specterops/dawgs/graph" +) + +type fragmentFailurePoint string + +const ( + fragmentFailureOpen fragmentFailurePoint = "open" + fragmentFailurePrepare fragmentFailurePoint = "prepare" + fragmentFailureCommit fragmentFailurePoint = "commit" +) + +type fragmentLifecycle struct { + opened bool + wrote bool + prepared bool + writerAborted bool + preparedAborted bool +} + +type failingNodeSink struct { + point fragmentFailurePoint + failure error + lifecycle *fragmentLifecycle +} + +func (s failingNodeSink) Open(context.Context, shardID) (fragmentWriter[normalizedNode, jsonlFragmentMetadata], error) { + if s.point == fragmentFailureOpen { + return nil, s.failure + } + s.lifecycle.opened = true + return &failingNodeWriter{ + point: s.point, + failure: s.failure, + lifecycle: s.lifecycle, + }, nil +} + +type failingNodeWriter struct { + point fragmentFailurePoint + failure error + lifecycle *fragmentLifecycle +} + +func (s *failingNodeWriter) WriteBatch(context.Context, []normalizedNode) error { + s.lifecycle.wrote = true + return nil +} + +func (s *failingNodeWriter) Prepare(context.Context) (preparedFragment[jsonlFragmentMetadata], error) { + if s.point == fragmentFailurePrepare { + return nil, s.failure + } + s.lifecycle.prepared = true + return &failingPreparedNodeFragment{ + failure: s.failure, + lifecycle: s.lifecycle, + }, nil +} + +func (s *failingNodeWriter) Abort() error { + s.lifecycle.writerAborted = true + return nil +} + +type failingPreparedNodeFragment struct { + failure error + lifecycle *fragmentLifecycle +} + +type failingParquetWriteSink struct { + failure error +} + +func (s failingParquetWriteSink) Open(context.Context, shardID) (fragmentWriter[normalizedNode, parquetFragmentMetadata], error) { + return &failingParquetWriteWriter{failure: s.failure}, nil +} + +type failingParquetWriteWriter struct { + failure error + aborted bool +} + +func (s *failingParquetWriteWriter) WriteBatch(context.Context, []normalizedNode) error { + return s.failure +} + +func (s *failingParquetWriteWriter) Prepare(context.Context) (preparedFragment[parquetFragmentMetadata], error) { + return nil, errors.New("unexpected prepare") +} + +func (s *failingParquetWriteWriter) Abort() error { + s.aborted = true + return nil +} + +func (s *failingPreparedNodeFragment) Metadata() jsonlFragmentMetadata { + return jsonlFragmentMetadata{Path: "unused", Rows: 1, SHA256: "unused"} +} + +func (s *failingPreparedNodeFragment) Commit(context.Context) error { + return s.failure +} + +func (s *failingPreparedNodeFragment) Abort() error { + s.lifecycle.preparedAborted = true + return nil +} + +type publishFailureWorkspace struct { + collectionWorkspace + failure error +} + +func (s publishFailureWorkspace) Publish(context.Context, string, []byte) (string, error) { + return "", s.failure +} + +func TestDumpFragmentFailuresDoNotPublishManifest(t *testing.T) { + for _, point := range []fragmentFailurePoint{ + fragmentFailureOpen, + fragmentFailurePrepare, + fragmentFailureCommit, + } { + t.Run(string(point), func(t *testing.T) { + outputDir := t.TempDir() + failure := errors.New("injected " + string(point) + " failure") + lifecycle := &fragmentLifecycle{} + source := oneNodeGraphSource() + + _, err := runDump( + context.Background(), + source, + "failure-test", + []GraphTarget{{Name: "source"}}, + DefaultDumpOptions(outputDir), + dumpOverrides{ + nodeOutput: newShardSinkSet(newJSONLShardSink(failingNodeSink{ + point: point, + failure: failure, + lifecycle: lifecycle, + })), + }, + ) + if !errors.Is(err, failure) { + t.Fatalf("dump error = %v, want injected failure", err) + } + assertShardSinkOperation(t, err, jsonlFragmentFormat, shardID{Graph: "source", Phase: PhaseNodes, Number: 1}, string(point), failure) + assertNoPublishedManifest(t, outputDir) + + switch point { + case fragmentFailureOpen: + if lifecycle.opened || lifecycle.wrote || lifecycle.prepared { + t.Fatalf("open failure lifecycle = %+v", lifecycle) + } + case fragmentFailurePrepare: + if !lifecycle.opened || !lifecycle.wrote || !lifecycle.writerAborted || lifecycle.prepared { + t.Fatalf("prepare failure lifecycle = %+v", lifecycle) + } + case fragmentFailureCommit: + if !lifecycle.opened || !lifecycle.wrote || !lifecycle.prepared || !lifecycle.preparedAborted { + t.Fatalf("commit failure lifecycle = %+v", lifecycle) + } + } + }) + } +} + +func TestDumpPublishFailureLeavesFragmentsWithoutManifest(t *testing.T) { + outputDir := t.TempDir() + failure := errors.New("injected publish failure") + workspace := publishFailureWorkspace{ + collectionWorkspace: newLocalCollectionWorkspace(outputDir, false), + failure: failure, + } + + _, err := runDump( + context.Background(), + oneNodeGraphSource(), + "failure-test", + []GraphTarget{{Name: "source"}}, + DefaultDumpOptions(outputDir), + dumpOverrides{workspace: workspace}, + ) + if !errors.Is(err, failure) { + t.Fatalf("dump error = %v, want injected failure", err) + } + assertNoPublishedManifest(t, outputDir) + + fragments, err := filepath.Glob(filepath.Join(outputDir, "graphs", "source", "nodes-*.jsonl.zst")) + if err != nil { + t.Fatalf("find committed fragments: %v", err) + } + if len(fragments) != 1 { + t.Fatalf("committed fragments = %v, want one", fragments) + } +} + +func TestRequestedParquetFailurePublishesNoManifestOrSuccessMarker(t *testing.T) { + outputDir := t.TempDir() + workspace := newLocalCollectionWorkspace(outputDir, false) + options := DefaultDumpOptions(outputDir) + options.Parquet = true + failure := errors.New("injected Parquet failure") + nodeOutput := newShardSinkSet( + newJSONLShardSink(newJSONLNodeSinkInWorkspace(options, workspace)), + newParquetShardSink[normalizedNode](failingParquetWriteSink{failure: failure}), + ) + + _, err := runDump( + context.Background(), + oneNodeGraphSource(), + "failure-test", + []GraphTarget{{Name: "source"}}, + options, + dumpOverrides{workspace: workspace, nodeOutput: nodeOutput}, + ) + assertShardSinkOperation(t, err, parquetFragmentFormat, shardID{Graph: "source", Phase: PhaseNodes, Number: 1}, "write", failure) + assertNoPublishedManifest(t, outputDir) + for _, relativePath := range []string{parquetManifestFileName, parquetSuccessFileName} { + if _, statErr := os.Stat(filepath.Join(outputDir, filepath.FromSlash(relativePath))); !os.IsNotExist(statErr) { + t.Fatalf("%s should not be published, stat error = %v", relativePath, statErr) + } + } + jsonlPath, pathErr := jsonlFragmentPath("source", PhaseNodes, 1, options.Compression) + if pathErr != nil { + t.Fatalf("JSONL path: %v", pathErr) + } + if _, statErr := os.Stat(filepath.Join(outputDir, filepath.FromSlash(jsonlPath))); !os.IsNotExist(statErr) { + t.Fatalf("JSONL sibling should be aborted, stat error = %v", statErr) + } +} + +func oneNodeGraphSource() *scriptedGraphSource { + return &scriptedGraphSource{ + snapshot: graphEntitySnapshot{NodeCount: 1}, + nodeBatches: [][]*graph.Node{{ + graph.NewNode(1, nil, graph.StringKind("User")), + }}, + } +} + +func assertNoPublishedManifest(t *testing.T, outputDir string) { + t.Helper() + if _, err := os.Stat(filepath.Join(outputDir, manifestFileName)); !os.IsNotExist(err) { + t.Fatalf("manifest should not be published, stat error = %v", err) + } +} diff --git a/retriever/dump_test.go b/retriever/dump_test.go index 011cae77..5e810391 100644 --- a/retriever/dump_test.go +++ b/retriever/dump_test.go @@ -1,12 +1,9 @@ package retriever -import ( - "path/filepath" - "testing" -) +import "testing" -func TestFragmentPath(t *testing.T) { - nodePath, err := fragmentPath("graph/name", PhaseNodes, 7, CompressionZstd) +func TestJSONLFragmentPath(t *testing.T) { + nodePath, err := jsonlFragmentPath("graph/name", PhaseNodes, 7, CompressionZstd) if err != nil { t.Fatalf("node fragment path: %v", err) } @@ -14,7 +11,7 @@ func TestFragmentPath(t *testing.T) { t.Fatalf("unexpected node fragment path %q", nodePath) } - edgePath, err := fragmentPath("default", PhaseEdges, 3, CompressionGzip) + edgePath, err := jsonlFragmentPath("default", PhaseEdges, 3, CompressionGzip) if err != nil { t.Fatalf("edge fragment path: %v", err) } @@ -22,52 +19,14 @@ func TestFragmentPath(t *testing.T) { t.Fatalf("unexpected edge fragment path %q", edgePath) } - if _, err := fragmentPath("default", Phase("bad"), 1, CompressionGzip); err == nil { + if _, err := jsonlFragmentPath("default", Phase("bad"), 1, CompressionGzip); err == nil { t.Fatalf("expected unsupported Phase error") } - if _, err := fragmentPath("default", PhaseNodes, 0, CompressionGzip); err == nil { + if _, err := jsonlFragmentPath("default", PhaseNodes, 0, CompressionGzip); err == nil { t.Fatalf("expected invalid shard number error") } } -func TestWriteFragmentMetadata(t *testing.T) { - options := DumpOptions{ - OutputDir: t.TempDir(), - Compression: CompressionGzip, - ZstdLevel: DefaultZstdLevel, - } - - fileEntry, err := writeNodeFragment(options.OutputDir, "default", 1, options, []FragmentNode{{ - ID: "1", - Kinds: []string{"User"}, - Properties: map[string]any{"name": "alice"}, - }}, map[string]int{"pseudonymize": 1}) - if err != nil { - t.Fatalf("write node fragment: %v", err) - } - if fileEntry.Phase != PhaseNodes || fileEntry.Path != "graphs/default/nodes-000001.jsonl.gz" || fileEntry.Count != 1 { - t.Fatalf("unexpected node file Manifest: %+v", fileEntry) - } - if fileEntry.ActionCounts["pseudonymize"] != 1 { - t.Fatalf("missing action count: %+v", fileEntry.ActionCounts) - } - if _, err := readManifest(filepath.Join(options.OutputDir, "graphs")); err == nil { - t.Fatalf("fragment write should not create Manifest") - } - - edgeEntry, err := writeEdgeFragment(options.OutputDir, "default", 2, options, []FragmentEdge{{ - StartID: "1", - EndID: "2", - Kind: "AdminTo", - }}, nil) - if err != nil { - t.Fatalf("write edge fragment: %v", err) - } - if edgeEntry.Phase != PhaseEdges || edgeEntry.Path != "graphs/default/edges-000002.jsonl.gz" || edgeEntry.Count != 1 { - t.Fatalf("unexpected edge file Manifest: %+v", edgeEntry) - } -} - func TestKindAndActionHelpers(t *testing.T) { kinds := map[string]struct{}{} addKindsToSet(kinds, []string{"User", "", "Computer", "User"}) diff --git a/retriever/load_benchmark_test.go b/retriever/load_benchmark_test.go index 5f97372b..8c24cbd4 100644 --- a/retriever/load_benchmark_test.go +++ b/retriever/load_benchmark_test.go @@ -48,9 +48,14 @@ func BenchmarkLoadFragmentPath(b *testing.B) { if err != nil { b.Fatalf("write node benchmark fragment: %v", err) } - entry.Phase = PhaseNodes - entry.Path = path - files = append(files, entry) + files = append(files, FileManifest{ + Phase: PhaseNodes, + Path: path, + Count: entry.Rows, + CompressedBytes: entry.CompressedBytes, + UncompressedBytes: entry.UncompressedBytes, + SHA256: entry.SHA256, + }) } for start, shard := 0, 1; start < len(edges); start, shard = start+shardSize, shard+1 { end := min(start+shardSize, len(edges)) @@ -59,9 +64,14 @@ func BenchmarkLoadFragmentPath(b *testing.B) { if err != nil { b.Fatalf("write edge benchmark fragment: %v", err) } - entry.Phase = PhaseEdges - entry.Path = path - files = append(files, entry) + files = append(files, FileManifest{ + Phase: PhaseEdges, + Path: path, + Count: entry.Rows, + CompressedBytes: entry.CompressedBytes, + UncompressedBytes: entry.UncompressedBytes, + SHA256: entry.SHA256, + }) } value := newValidTestManifest(1) diff --git a/retriever/manifest.go b/retriever/manifest.go index 11c02ce1..5a872c18 100644 --- a/retriever/manifest.go +++ b/retriever/manifest.go @@ -1,6 +1,7 @@ package retriever import ( + "context" "encoding/json" "fmt" "os" @@ -35,29 +36,27 @@ func WriteManifest(outputDir string, value Manifest) error { } func writeManifest(outputDir string, value Manifest) error { - if err := value.validate(); err != nil { - return err - } - - tempPath := filepath.Join(outputDir, manifestFileName+".tmp") - finalPath := filepath.Join(outputDir, manifestFileName) - payload, err := json.MarshalIndent(value, "", " ") + payload, err := encodeManifest(value) if err != nil { - return fmt.Errorf("encode manifest: %w", err) + return err } - payload = append(payload, '\n') + _, err = newLocalCollectionWorkspace(outputDir, false).Publish(context.Background(), manifestFileName, payload) + return err +} - if err := os.WriteFile(tempPath, payload, 0o600); err != nil { - return fmt.Errorf("write manifest temp file: %w", err) +func encodeManifest(value Manifest) ([]byte, error) { + if err := value.validate(); err != nil { + return nil, err } - if err := os.Rename(tempPath, finalPath); err != nil { - os.Remove(tempPath) - return fmt.Errorf("rename manifest: %w", err) + payload, err := json.MarshalIndent(value, "", " ") + if err != nil { + return nil, fmt.Errorf("encode manifest: %w", err) } - return nil + payload = append(payload, '\n') + return payload, nil } func VerifyManifestFiles(inputDir string, value Manifest) error { diff --git a/retriever/observer.go b/retriever/observer.go new file mode 100644 index 00000000..a1d11f39 --- /dev/null +++ b/retriever/observer.go @@ -0,0 +1,59 @@ +package retriever + +type graphObserver struct { + graphName string + nodeCount int64 + edgeCount int64 + nodeKinds map[string]struct{} + edgeKinds map[string]struct{} + nodeActionCounts map[string]int + edgeActionCounts map[string]int + metrics *metricsBuilder +} + +func newGraphObserver(graphName string, expectedNodeCount int64) *graphObserver { + return &graphObserver{ + graphName: graphName, + nodeKinds: map[string]struct{}{}, + edgeKinds: map[string]struct{}{}, + nodeActionCounts: map[string]int{}, + edgeActionCounts: map[string]int{}, + metrics: newMetricsBuilder(graphName, expectedNodeCount), + } +} + +func (s *graphObserver) ObserveNode(record normalizedNode, actionCounts map[string]int) error { + if err := s.metrics.observeFragmentNode(jsonlV1NodeFromNormalized(record)); err != nil { + return err + } + + addKindsToSet(s.nodeKinds, record.Kinds) + addActionCounts(s.nodeActionCounts, actionCounts) + s.nodeCount++ + return nil +} + +func (s *graphObserver) ObserveEdge(record normalizedEdge, actionCounts map[string]int) error { + if err := s.metrics.observeFragmentEdge(jsonlV1EdgeFromNormalized(record)); err != nil { + return err + } + + if record.Kind != "" { + s.edgeKinds[record.Kind] = struct{}{} + } + addActionCounts(s.edgeActionCounts, actionCounts) + s.edgeCount++ + return nil +} + +func (s *graphObserver) Schema() GraphSchemaMetadata { + return GraphSchemaMetadata{ + Name: s.graphName, + NodeKinds: stringsFromKindSet(s.nodeKinds), + EdgeKinds: stringsFromKindSet(s.edgeKinds), + } +} + +func (s *graphObserver) Metrics() GraphMetrics { + return s.metrics.finalize() +} diff --git a/retriever/observer_test.go b/retriever/observer_test.go new file mode 100644 index 00000000..42596ec5 --- /dev/null +++ b/retriever/observer_test.go @@ -0,0 +1,38 @@ +package retriever + +import ( + "reflect" + "testing" +) + +func TestGraphObserverCollectsSharedPipelineMetadata(t *testing.T) { + observer := newGraphObserver("example", 2) + if err := observer.ObserveNode(normalizedNode{ID: "1", Kinds: []string{"Admin", "User"}, Properties: map[string]any{}}, map[string]int{"preserve": 1}); err != nil { + t.Fatalf("observe first node: %v", err) + } + if err := observer.ObserveNode(normalizedNode{ID: "2", Kinds: []string{"Group"}, Properties: map[string]any{}}, map[string]int{"redact": 1}); err != nil { + t.Fatalf("observe second node: %v", err) + } + if err := observer.ObserveEdge(normalizedEdge{ID: "source-edge-id", StartID: "1", EndID: "2", Kind: "MemberOf", Properties: map[string]any{}}, map[string]int{"preserve": 1}); err != nil { + t.Fatalf("observe edge: %v", err) + } + + if observer.nodeCount != 2 || observer.edgeCount != 1 { + t.Fatalf("observer totals: nodes=%d edges=%d", observer.nodeCount, observer.edgeCount) + } + if !reflect.DeepEqual(observer.nodeActionCounts, map[string]int{"preserve": 1, "redact": 1}) || !reflect.DeepEqual(observer.edgeActionCounts, map[string]int{"preserve": 1}) { + t.Fatalf("observer actions: nodes=%v edges=%v", observer.nodeActionCounts, observer.edgeActionCounts) + } + if !reflect.DeepEqual(observer.Schema(), GraphSchemaMetadata{ + Name: "example", + NodeKinds: []string{"Admin", "Group", "User"}, + EdgeKinds: []string{"MemberOf"}, + }) { + t.Fatalf("observer schema = %+v", observer.Schema()) + } + + metrics := observer.Metrics() + if metrics.NodeCount != 2 || metrics.EdgeCount != 1 { + t.Fatalf("observer metrics = %+v", metrics) + } +} diff --git a/retriever/options.go b/retriever/options.go index 4eaba173..2ab8c352 100644 --- a/retriever/options.go +++ b/retriever/options.go @@ -48,6 +48,7 @@ type ProgressEvent struct { NodeCount int64 EdgeCount int64 Compression CompressionCodec + Parquet bool Scrub ScrubMode Elapsed time.Duration EntitiesPerSecond float64 @@ -64,6 +65,7 @@ func (s ProgressFunc) emit(event ProgressEvent) { type DumpOptions struct { OutputDir string Force bool + Parquet bool Scrub ScrubMode Salt string ScrubConfig io.Reader diff --git a/retriever/options_test.go b/retriever/options_test.go index 89e096a4..f82ad5cc 100644 --- a/retriever/options_test.go +++ b/retriever/options_test.go @@ -124,7 +124,7 @@ func TestOptionsValidate(t *testing.T) { func TestDefaultOptions(t *testing.T) { dump := DefaultDumpOptions(t.TempDir()) - if dump.Scrub != ScrubNone || dump.Compression != CompressionZstd || dump.ZstdLevel != DefaultZstdLevel { + if dump.Parquet || dump.Scrub != ScrubNone || dump.Compression != CompressionZstd || dump.ZstdLevel != DefaultZstdLevel { t.Fatalf("unexpected dump defaults: %+v", dump) } if dump.ShardSize != DefaultShardSize || dump.BatchSize != DefaultBatchSize || dump.ProgressInterval != DefaultProgressInterval { diff --git a/retriever/output.go b/retriever/output.go new file mode 100644 index 00000000..37c47aae --- /dev/null +++ b/retriever/output.go @@ -0,0 +1,361 @@ +package retriever + +import ( + "context" + "fmt" +) + +type committedShard struct { + JSONL jsonlFragmentMetadata + Parquet *parquetFragmentMetadata +} + +type shardOutput[T any] interface { + OpenShard(context.Context, shardID) (shardOutputWriter[T], error) +} + +type shardOutputWriter[T any] interface { + WriteBatch(context.Context, []T) error + Finish(context.Context, shardSummary) (committedShard, error) + Abort() error +} + +type shardSink[T any] interface { + format() string + open(context.Context, shardID) (shardSinkWriter[T], error) +} + +type shardSinkWriter[T any] interface { + writeBatch(context.Context, []T) error + prepare(context.Context) (preparedShardSink, error) + abort() error +} + +type preparedShardSink interface { + rowCount() int + commit(context.Context) error + abort() error + addTo(*committedShard) +} + +type typedShardSink[T any, M fragmentMetadata] struct { + name string + sink fragmentSink[T, M] + collect func(*committedShard, M) +} + +func newShardSink[T any, M fragmentMetadata](name string, sink fragmentSink[T, M], collect func(*committedShard, M)) shardSink[T] { + return typedShardSink[T, M]{name: name, sink: sink, collect: collect} +} + +func newJSONLShardSink[T any](sink fragmentSink[T, jsonlFragmentMetadata]) shardSink[T] { + return newShardSink(jsonlFragmentFormat, sink, func(committed *committedShard, metadata jsonlFragmentMetadata) { + committed.JSONL = metadata + }) +} + +func newParquetShardSink[T any](sink fragmentSink[T, parquetFragmentMetadata]) shardSink[T] { + return newShardSink(parquetFragmentFormat, sink, func(committed *committedShard, metadata parquetFragmentMetadata) { + committed.Parquet = &metadata + }) +} + +func (s typedShardSink[T, M]) format() string { + return s.name +} + +func (s typedShardSink[T, M]) open(ctx context.Context, id shardID) (shardSinkWriter[T], error) { + writer, err := s.sink.Open(ctx, id) + if err != nil { + return nil, err + } + return typedShardSinkWriter[T, M]{writer: writer, collect: s.collect}, nil +} + +type typedShardSinkWriter[T any, M fragmentMetadata] struct { + writer fragmentWriter[T, M] + collect func(*committedShard, M) +} + +func (s typedShardSinkWriter[T, M]) writeBatch(ctx context.Context, records []T) error { + return s.writer.WriteBatch(ctx, records) +} + +func (s typedShardSinkWriter[T, M]) prepare(ctx context.Context) (preparedShardSink, error) { + prepared, err := s.writer.Prepare(ctx) + if err != nil { + return nil, err + } + return typedPreparedShardSink[M]{prepared: prepared, collect: s.collect}, nil +} + +func (s typedShardSinkWriter[T, M]) abort() error { + return s.writer.Abort() +} + +type typedPreparedShardSink[M fragmentMetadata] struct { + prepared preparedFragment[M] + collect func(*committedShard, M) +} + +func (s typedPreparedShardSink[M]) rowCount() int { + return s.prepared.Metadata().rowCount() +} + +func (s typedPreparedShardSink[M]) commit(ctx context.Context) error { + return s.prepared.Commit(ctx) +} + +func (s typedPreparedShardSink[M]) abort() error { + return s.prepared.Abort() +} + +func (s typedPreparedShardSink[M]) addTo(committed *committedShard) { + s.collect(committed, s.prepared.Metadata()) +} + +type shardSinkSet[T any] struct { + sinks []shardSink[T] +} + +func newShardSinkSet[T any](sinks ...shardSink[T]) shardSinkSet[T] { + return shardSinkSet[T]{sinks: sinks} +} + +func (s shardSinkSet[T]) OpenShard(ctx context.Context, id shardID) (shardOutputWriter[T], error) { + if len(s.sinks) == 0 { + return nil, fmt.Errorf("open shard %d without a sink", id.Number) + } + + writer := &shardSinkSetWriter[T]{id: id, sinks: make([]activeShardSink[T], len(s.sinks))} + for index, sink := range s.sinks { + writer.sinks[index].format = sink.format() + } + errs := runSinkOperations(ctx, len(s.sinks), func(index int, operationCtx context.Context) error { + opened, err := s.sinks[index].open(operationCtx, id) + writer.sinks[index].writer = opened + return writer.failure(index, "open", err) + }) + if err := collectErrors(errs...); err != nil { + return nil, cleanupOnError(err, writer.Abort) + } + return writer, nil +} + +type activeShardSink[T any] struct { + format string + writer shardSinkWriter[T] + prepared preparedShardSink + committed bool +} + +type shardSinkOperationError struct { + Format string + ID shardID + Operation string + Err error +} + +func (s *shardSinkOperationError) Error() string { + return fmt.Sprintf("%s sink graph %q phase %q shard %d operation %q: %v", s.Format, s.ID.Graph, s.ID.Phase, s.ID.Number, s.Operation, s.Err) +} + +func (s *shardSinkOperationError) Unwrap() error { + return s.Err +} + +type shardSinkSetWriter[T any] struct { + id shardID + sinks []activeShardSink[T] + closed bool +} + +func (s *shardSinkSetWriter[T]) WriteBatch(ctx context.Context, records []T) error { + if s.closed { + return fmt.Errorf("write shard %d after output closed", s.id.Number) + } + errs := runSinkOperations(ctx, len(s.sinks), func(index int, operationCtx context.Context) error { + return s.failure(index, "write", s.sinks[index].writer.writeBatch(operationCtx, records)) + }) + if err := collectErrors(errs...); err != nil { + return cleanupOnError(err, s.Abort) + } + return nil +} + +func (s *shardSinkSetWriter[T]) Finish(ctx context.Context, summary shardSummary) (committedShard, error) { + if s.closed { + return committedShard{}, fmt.Errorf("finish shard %d after output closed", s.id.Number) + } + if summary.ID != s.id { + err := fmt.Errorf("finish shard %d while shard %d is active", summary.ID.Number, s.id.Number) + return committedShard{}, cleanupOnError(s.allFailures("validate", err), s.Abort) + } + + errs := runSinkOperations(ctx, len(s.sinks), func(index int, operationCtx context.Context) error { + prepared, err := s.sinks[index].writer.prepare(operationCtx) + s.sinks[index].prepared = prepared + return s.failure(index, "prepare", err) + }) + if err := collectErrors(errs...); err != nil { + return committedShard{}, cleanupOnError(err, s.Abort) + } + + if err := s.validate(summary); err != nil { + return committedShard{}, cleanupOnError(err, s.Abort) + } + + for index := range s.sinks { + if err := s.sinks[index].prepared.commit(ctx); err != nil { + return committedShard{}, cleanupOnError(s.failure(index, "commit", err), s.Abort) + } + s.sinks[index].committed = true + } + + var committed committedShard + for _, sink := range s.sinks { + sink.prepared.addTo(&committed) + } + s.closed = true + return committed, nil +} + +func (s *shardSinkSetWriter[T]) validate(summary shardSummary) error { + errs := make([]error, 0, len(s.sinks)*2) + peerRows := s.sinks[0].prepared.rowCount() + for index, sink := range s.sinks { + rows := sink.prepared.rowCount() + if rows != summary.Rows { + err := fmt.Errorf("prepared shard %d has %d rows, expected %d", summary.ID.Number, rows, summary.Rows) + errs = append(errs, s.failure(index, "validate", err)) + } + if index > 0 && rows != peerRows { + err := fmt.Errorf("prepared shard %d has %d rows, first sink has %d", summary.ID.Number, rows, peerRows) + errs = append(errs, s.failure(index, "validate", err)) + } + } + return collectErrors(errs...) +} + +func (s *shardSinkSetWriter[T]) Abort() error { + if s.closed { + return nil + } + s.closed = true + + errs := make([]error, 0, len(s.sinks)) + for index, sink := range s.sinks { + if sink.committed || sink.writer == nil { + continue + } + if sink.prepared != nil { + errs = append(errs, s.failure(index, "abort", sink.prepared.abort())) + } else { + errs = append(errs, s.failure(index, "abort", sink.writer.abort())) + } + } + return collectErrors(errs...) +} + +func (s *shardSinkSetWriter[T]) failure(index int, operation string, err error) error { + if err == nil { + return nil + } + return &shardSinkOperationError{Format: s.sinks[index].format, ID: s.id, Operation: operation, Err: err} +} + +func (s *shardSinkSetWriter[T]) allFailures(operation string, err error) error { + errs := make([]error, len(s.sinks)) + for index := range s.sinks { + errs[index] = s.failure(index, operation, err) + } + return collectErrors(errs...) +} + +type sinkOperationResult struct { + index int + err error +} + +func runSinkOperations(ctx context.Context, count int, operation func(int, context.Context) error) []error { + if count == 1 { + return []error{operation(0, ctx)} + } + + operationCtx, cancel := context.WithCancel(ctx) + defer cancel() + results := make(chan sinkOperationResult, count) + for index := range count { + go func() { + results <- sinkOperationResult{index: index, err: operation(index, operationCtx)} + }() + } + + errs := make([]error, count) + for range count { + result := <-results + errs[result.index] = result.err + if result.err != nil { + cancel() + } + } + return errs +} + +type shardOutputReceiver[T any] struct { + ctx context.Context + output shardOutput[T] + accept func(shardSummary, committedShard) error + writer shardOutputWriter[T] +} + +func newShardOutputReceiver[T any](ctx context.Context, output shardOutput[T], accept func(shardSummary, committedShard) error) *shardOutputReceiver[T] { + return &shardOutputReceiver[T]{ctx: ctx, output: output, accept: accept} +} + +func (s *shardOutputReceiver[T]) BeginShard(id shardID) error { + if s.writer != nil { + return fmt.Errorf("begin shard %d while another shard is active", id.Number) + } + + writer, err := s.output.OpenShard(s.ctx, id) + if err != nil { + return err + } + s.writer = writer + return nil +} + +func (s *shardOutputReceiver[T]) WriteBatch(records []T) error { + if s.writer == nil { + return fmt.Errorf("write shard batch without an active shard") + } + if err := s.writer.WriteBatch(s.ctx, records); err != nil { + writer := s.writer + s.writer = nil + return cleanupOnError(err, writer.Abort) + } + return nil +} + +func (s *shardOutputReceiver[T]) FinishShard(summary shardSummary) error { + if s.writer == nil { + return fmt.Errorf("finish shard %d without an active shard", summary.ID.Number) + } + + committed, err := s.writer.Finish(s.ctx, summary) + s.writer = nil + if err != nil { + return err + } + return s.accept(summary, committed) +} + +func (s *shardOutputReceiver[T]) Abort() error { + if s.writer == nil { + return nil + } + err := s.writer.Abort() + s.writer = nil + return err +} diff --git a/retriever/output_fanout_test.go b/retriever/output_fanout_test.go new file mode 100644 index 00000000..fc714709 --- /dev/null +++ b/retriever/output_fanout_test.go @@ -0,0 +1,572 @@ +package retriever + +import ( + "context" + "errors" + "reflect" + "sync" + "testing" + "time" +) + +const testFragmentFormat = "TEST" + +type testFragmentMetadata struct { + Rows int +} + +func (s testFragmentMetadata) rowCount() int { + return s.Rows +} + +type fanoutEventLog struct { + mu sync.Mutex + events []string +} + +func (s *fanoutEventLog) add(event string) { + s.mu.Lock() + defer s.mu.Unlock() + s.events = append(s.events, event) +} + +func (s *fanoutEventLog) snapshot() []string { + s.mu.Lock() + defer s.mu.Unlock() + return append([]string(nil), s.events...) +} + +type fanoutSinkState[T any] struct { + mu sync.Mutex + + name string + events *fanoutEventLog + rowOffset int + + openHook func(context.Context) error + writeHook func(context.Context) error + prepareHook func(context.Context) error + commitHook func(context.Context) error + abortErr error + + openedOnce sync.Once + preparedOnce sync.Once + opened chan struct{} + prepared chan struct{} + + ids []shardID + batches [][]T + batchStarts []*T + rows int + writerAborts int + preparedAborts int + commitAttempts int + committed bool +} + +func newFanoutSinkState[T any](name string, events *fanoutEventLog) *fanoutSinkState[T] { + return &fanoutSinkState[T]{ + name: name, + events: events, + opened: make(chan struct{}), + prepared: make(chan struct{}), + } +} + +func (s *fanoutSinkState[T]) addEvent(event string) { + if s.events != nil { + s.events.add(s.name + "." + event) + } +} + +type fanoutSinkSnapshot[T any] struct { + ids []shardID + batches [][]T + batchStarts []*T + writerAborts int + preparedAborts int + commitAttempts int + committed bool +} + +func (s *fanoutSinkState[T]) snapshot() fanoutSinkSnapshot[T] { + s.mu.Lock() + defer s.mu.Unlock() + + batches := make([][]T, len(s.batches)) + for index, batch := range s.batches { + batches[index] = append([]T(nil), batch...) + } + return fanoutSinkSnapshot[T]{ + ids: append([]shardID(nil), s.ids...), + batches: batches, + batchStarts: append([]*T(nil), s.batchStarts...), + writerAborts: s.writerAborts, + preparedAborts: s.preparedAborts, + commitAttempts: s.commitAttempts, + committed: s.committed, + } +} + +type fanoutTestSink[T any, M fragmentMetadata] struct { + state *fanoutSinkState[T] + metadata func(int) M +} + +func (s fanoutTestSink[T, M]) Open(ctx context.Context, id shardID) (fragmentWriter[T, M], error) { + s.state.mu.Lock() + s.state.ids = append(s.state.ids, id) + s.state.mu.Unlock() + if s.state.openHook != nil { + if err := s.state.openHook(ctx); err != nil { + return nil, err + } + } + s.state.addEvent("open") + s.state.openedOnce.Do(func() { close(s.state.opened) }) + return &fanoutTestWriter[T, M]{sink: s}, nil +} + +type fanoutTestWriter[T any, M fragmentMetadata] struct { + sink fanoutTestSink[T, M] +} + +func (s *fanoutTestWriter[T, M]) WriteBatch(ctx context.Context, records []T) error { + s.sink.state.mu.Lock() + s.sink.state.batches = append(s.sink.state.batches, append([]T(nil), records...)) + if len(records) > 0 { + s.sink.state.batchStarts = append(s.sink.state.batchStarts, &records[0]) + } else { + s.sink.state.batchStarts = append(s.sink.state.batchStarts, nil) + } + s.sink.state.rows += len(records) + s.sink.state.mu.Unlock() + s.sink.state.addEvent("write") + if s.sink.state.writeHook != nil { + return s.sink.state.writeHook(ctx) + } + return nil +} + +func (s *fanoutTestWriter[T, M]) Prepare(ctx context.Context) (preparedFragment[M], error) { + if s.sink.state.prepareHook != nil { + if err := s.sink.state.prepareHook(ctx); err != nil { + return nil, err + } + } + s.sink.state.mu.Lock() + rows := s.sink.state.rows + s.sink.state.rowOffset + s.sink.state.mu.Unlock() + prepared := &fanoutTestPrepared[T, M]{state: s.sink.state, metadata: s.sink.metadata(rows)} + s.sink.state.addEvent("prepare") + s.sink.state.preparedOnce.Do(func() { close(s.sink.state.prepared) }) + return prepared, nil +} + +func (s *fanoutTestWriter[T, M]) Abort() error { + s.sink.state.mu.Lock() + s.sink.state.writerAborts++ + s.sink.state.mu.Unlock() + s.sink.state.addEvent("writer-abort") + return s.sink.state.abortErr +} + +type fanoutTestPrepared[T any, M fragmentMetadata] struct { + state *fanoutSinkState[T] + metadata M +} + +func (s *fanoutTestPrepared[T, M]) Metadata() M { + return s.metadata +} + +func (s *fanoutTestPrepared[T, M]) Commit(ctx context.Context) error { + s.state.mu.Lock() + s.state.commitAttempts++ + s.state.mu.Unlock() + s.state.addEvent("commit-attempt") + if s.state.commitHook != nil { + if err := s.state.commitHook(ctx); err != nil { + return err + } + } + s.state.mu.Lock() + s.state.committed = true + s.state.mu.Unlock() + s.state.addEvent("commit") + return nil +} + +func (s *fanoutTestPrepared[T, M]) Abort() error { + s.state.mu.Lock() + s.state.preparedAborts++ + s.state.mu.Unlock() + s.state.addEvent("prepared-abort") + return s.state.abortErr +} + +func newFanoutJSONLSink[T any](state *fanoutSinkState[T]) fragmentSink[T, jsonlFragmentMetadata] { + return fanoutTestSink[T, jsonlFragmentMetadata]{ + state: state, + metadata: func(rows int) jsonlFragmentMetadata { + return jsonlFragmentMetadata{Path: "fragment.jsonl.gz", Rows: rows} + }, + } +} + +func newFanoutSecondarySink[T any](state *fanoutSinkState[T], collected *testFragmentMetadata) shardSink[T] { + leaf := fanoutTestSink[T, testFragmentMetadata]{ + state: state, + metadata: func(rows int) testFragmentMetadata { + return testFragmentMetadata{Rows: rows} + }, + } + return newShardSink(testFragmentFormat, leaf, func(_ *committedShard, metadata testFragmentMetadata) { + *collected = metadata + }) +} + +func newTestSinkSet[T any](jsonl, secondary *fanoutSinkState[T], collected *testFragmentMetadata) shardSinkSet[T] { + return newShardSinkSet( + newJSONLShardSink(newFanoutJSONLSink(jsonl)), + newFanoutSecondarySink(secondary, collected), + ) +} + +func TestShardSinkSetFansOutSameSlicesAndSummary(t *testing.T) { + events := &fanoutEventLog{} + jsonlState := newFanoutSinkState[int]("jsonl", events) + secondaryState := newFanoutSinkState[int]("secondary", events) + var collected testFragmentMetadata + output := newTestSinkSet(jsonlState, secondaryState, &collected) + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 4} + summary := shardSummary{ID: id, Rows: 4, ActionCounts: map[string]int{"drop": 2}} + var accepted shardSummary + receiver := newShardOutputReceiver(context.Background(), output, func(got shardSummary, committed committedShard) error { + accepted = got + if committed.JSONL.Rows != summary.Rows { + t.Fatalf("committed JSONL rows = %d", committed.JSONL.Rows) + } + return nil + }) + first := []int{1, 2} + second := []int{3, 4} + if err := receiver.BeginShard(id); err != nil { + t.Fatalf("begin shard: %v", err) + } + if err := receiver.WriteBatch(first); err != nil { + t.Fatalf("write first batch: %v", err) + } + if err := receiver.WriteBatch(second); err != nil { + t.Fatalf("write second batch: %v", err) + } + if err := receiver.FinishShard(summary); err != nil { + t.Fatalf("finish shard: %v", err) + } + + jsonl := jsonlState.snapshot() + secondary := secondaryState.snapshot() + if !reflect.DeepEqual(jsonl.ids, []shardID{id}) || !reflect.DeepEqual(secondary.ids, jsonl.ids) { + t.Fatalf("JSONL ids = %+v, secondary ids = %+v", jsonl.ids, secondary.ids) + } + if !reflect.DeepEqual(jsonl.batches, [][]int{{1, 2}, {3, 4}}) || !reflect.DeepEqual(secondary.batches, jsonl.batches) { + t.Fatalf("JSONL batches = %+v, secondary batches = %+v", jsonl.batches, secondary.batches) + } + if !reflect.DeepEqual(jsonl.batchStarts, []*int{&first[0], &second[0]}) || !reflect.DeepEqual(secondary.batchStarts, jsonl.batchStarts) { + t.Fatalf("sink set did not offer the same slices to both sinks") + } + if !reflect.DeepEqual(accepted, summary) || collected.Rows != summary.Rows || !jsonl.committed || !secondary.committed { + t.Fatalf("summary/result mismatch: accepted=%+v collected=%+v JSONL=%+v secondary=%+v", accepted, collected, jsonl, secondary) + } + + lifecycle := events.snapshot() + jsonlPrepare := eventPosition(lifecycle, "jsonl.prepare") + secondaryPrepare := eventPosition(lifecycle, "secondary.prepare") + jsonlCommit := eventPosition(lifecycle, "jsonl.commit-attempt") + secondaryCommit := eventPosition(lifecycle, "secondary.commit-attempt") + if jsonlPrepare < 0 || secondaryPrepare < 0 || jsonlCommit < jsonlPrepare || jsonlCommit < secondaryPrepare || secondaryCommit < jsonlCommit { + t.Fatalf("prepare and deterministic commit order = %v", lifecycle) + } +} + +func TestShardSinkSetBackpressuresCurrentBatch(t *testing.T) { + jsonlState := newFanoutSinkState[int]("jsonl", nil) + secondaryState := newFanoutSinkState[int]("secondary", nil) + jsonlStarted := make(chan struct{}) + secondaryStarted := make(chan struct{}) + release := make(chan struct{}) + jsonlState.writeHook = func(context.Context) error { + close(jsonlStarted) + return nil + } + secondaryState.writeHook = func(context.Context) error { + close(secondaryStarted) + <-release + return nil + } + var collected testFragmentMetadata + writer, err := newTestSinkSet(jsonlState, secondaryState, &collected).OpenShard(context.Background(), shardID{Graph: "example", Phase: PhaseNodes, Number: 1}) + if err != nil { + t.Fatalf("open shard: %v", err) + } + batch := []int{1, 2, 3} + returned := make(chan error, 1) + go func() { returned <- writer.WriteBatch(context.Background(), batch) }() + waitForFanoutSignal(t, jsonlStarted, "JSONL write") + waitForFanoutSignal(t, secondaryStarted, "secondary write") + select { + case err := <-returned: + t.Fatalf("write returned before the slow sink accepted the batch: %v", err) + default: + } + close(release) + if err := <-returned; err != nil { + t.Fatalf("write batch: %v", err) + } + jsonl := jsonlState.snapshot() + secondary := secondaryState.snapshot() + if len(jsonl.batches) != 1 || len(secondary.batches) != 1 || jsonl.batchStarts[0] != &batch[0] || secondary.batchStarts[0] != &batch[0] { + t.Fatalf("active batch was copied or queued: JSONL=%+v secondary=%+v", jsonl, secondary) + } + if err := writer.Abort(); err != nil { + t.Fatalf("abort: %v", err) + } +} + +func TestShardSinkSetOpenFailureAbortsOpenedSibling(t *testing.T) { + for _, failingFormat := range []string{jsonlFragmentFormat, testFragmentFormat} { + t.Run(failingFormat, func(t *testing.T) { + cause := errors.New("open failed") + jsonlState := newFanoutSinkState[int]("jsonl", nil) + secondaryState := newFanoutSinkState[int]("secondary", nil) + if failingFormat == jsonlFragmentFormat { + jsonlState.openHook = func(context.Context) error { + <-secondaryState.opened + return cause + } + } else { + secondaryState.openHook = func(context.Context) error { + <-jsonlState.opened + return cause + } + } + var collected testFragmentMetadata + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 1} + _, err := newTestSinkSet(jsonlState, secondaryState, &collected).OpenShard(context.Background(), id) + assertShardSinkOperation(t, err, failingFormat, id, "open", cause) + if failingFormat == jsonlFragmentFormat && secondaryState.snapshot().writerAborts != 1 { + t.Fatalf("secondary sibling was not aborted: %+v", secondaryState.snapshot()) + } + if failingFormat == testFragmentFormat && jsonlState.snapshot().writerAborts != 1 { + t.Fatalf("JSONL sibling was not aborted: %+v", jsonlState.snapshot()) + } + }) + } +} + +func TestShardSinkSetWriteFailureCancelsAndAbortsSibling(t *testing.T) { + for _, failingFormat := range []string{jsonlFragmentFormat, testFragmentFormat} { + t.Run(failingFormat, func(t *testing.T) { + cause := errors.New("write failed") + cleanupCause := errors.New("sibling abort failed") + jsonlState := newFanoutSinkState[int]("jsonl", nil) + secondaryState := newFanoutSinkState[int]("secondary", nil) + siblingStarted := make(chan struct{}) + siblingCanceled := make(chan struct{}) + failure := func(context.Context) error { + <-siblingStarted + return cause + } + blocking := func(ctx context.Context) error { + close(siblingStarted) + <-ctx.Done() + close(siblingCanceled) + return ctx.Err() + } + if failingFormat == jsonlFragmentFormat { + jsonlState.writeHook = failure + secondaryState.writeHook = blocking + secondaryState.abortErr = cleanupCause + } else { + jsonlState.writeHook = blocking + jsonlState.abortErr = cleanupCause + secondaryState.writeHook = failure + } + var collected testFragmentMetadata + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 1} + writer, err := newTestSinkSet(jsonlState, secondaryState, &collected).OpenShard(context.Background(), id) + if err != nil { + t.Fatalf("open shard: %v", err) + } + err = writer.WriteBatch(context.Background(), []int{1}) + assertShardSinkOperation(t, err, failingFormat, id, "write", cause) + if !errors.Is(err, context.Canceled) { + t.Fatalf("write error does not retain sibling cancellation: %v", err) + } + if !errors.Is(err, cleanupCause) { + t.Fatalf("write error does not retain sibling cleanup failure: %v", err) + } + waitForFanoutSignal(t, siblingCanceled, "sibling cancellation") + if jsonlState.snapshot().writerAborts != 1 || secondaryState.snapshot().writerAborts != 1 { + t.Fatalf("write cleanup: JSONL=%+v secondary=%+v", jsonlState.snapshot(), secondaryState.snapshot()) + } + }) + } +} + +func TestShardSinkSetPrepareFailureAbortsPreparedSibling(t *testing.T) { + for _, failingFormat := range []string{jsonlFragmentFormat, testFragmentFormat} { + t.Run(failingFormat, func(t *testing.T) { + cause := errors.New("prepare failed") + jsonlState := newFanoutSinkState[int]("jsonl", nil) + secondaryState := newFanoutSinkState[int]("secondary", nil) + if failingFormat == jsonlFragmentFormat { + jsonlState.prepareHook = func(context.Context) error { + <-secondaryState.prepared + return cause + } + } else { + secondaryState.prepareHook = func(context.Context) error { + <-jsonlState.prepared + return cause + } + } + var collected testFragmentMetadata + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 1} + writer, err := newTestSinkSet(jsonlState, secondaryState, &collected).OpenShard(context.Background(), id) + if err != nil { + t.Fatalf("open shard: %v", err) + } + if err := writer.WriteBatch(context.Background(), []int{1}); err != nil { + t.Fatalf("write batch: %v", err) + } + _, err = writer.Finish(context.Background(), shardSummary{ID: id, Rows: 1}) + assertShardSinkOperation(t, err, failingFormat, id, "prepare", cause) + jsonl := jsonlState.snapshot() + secondary := secondaryState.snapshot() + if jsonl.commitAttempts != 0 || secondary.commitAttempts != 0 { + t.Fatalf("commit attempted after prepare failure: JSONL=%+v secondary=%+v", jsonl, secondary) + } + if failingFormat == jsonlFragmentFormat && (jsonl.writerAborts != 1 || secondary.preparedAborts != 1) { + t.Fatalf("JSONL prepare cleanup: JSONL=%+v secondary=%+v", jsonl, secondary) + } + if failingFormat == testFragmentFormat && (secondary.writerAborts != 1 || jsonl.preparedAborts != 1) { + t.Fatalf("secondary prepare cleanup: JSONL=%+v secondary=%+v", jsonl, secondary) + } + }) + } +} + +func TestShardSinkSetValidatesAllCountsBeforeCommit(t *testing.T) { + jsonlState := newFanoutSinkState[int]("jsonl", nil) + secondaryState := newFanoutSinkState[int]("secondary", nil) + secondaryState.rowOffset = 1 + var collected testFragmentMetadata + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 1} + writer, err := newTestSinkSet(jsonlState, secondaryState, &collected).OpenShard(context.Background(), id) + if err != nil { + t.Fatalf("open shard: %v", err) + } + if err := writer.WriteBatch(context.Background(), []int{1}); err != nil { + t.Fatalf("write batch: %v", err) + } + _, err = writer.Finish(context.Background(), shardSummary{ID: id, Rows: 1}) + assertShardSinkOperation(t, err, testFragmentFormat, id, "validate", nil) + jsonl := jsonlState.snapshot() + secondary := secondaryState.snapshot() + if jsonl.commitAttempts != 0 || secondary.commitAttempts != 0 || jsonl.preparedAborts != 1 || secondary.preparedAborts != 1 { + t.Fatalf("validation lifecycle: JSONL=%+v secondary=%+v", jsonl, secondary) + } +} + +func TestShardSinkSetCommitFailureAbortsOnlyUncommittedSinks(t *testing.T) { + for _, failingFormat := range []string{jsonlFragmentFormat, testFragmentFormat} { + t.Run(failingFormat, func(t *testing.T) { + cause := errors.New("commit failed") + jsonlState := newFanoutSinkState[int]("jsonl", nil) + secondaryState := newFanoutSinkState[int]("secondary", nil) + if failingFormat == jsonlFragmentFormat { + jsonlState.commitHook = func(context.Context) error { return cause } + } else { + secondaryState.commitHook = func(context.Context) error { return cause } + } + var collected testFragmentMetadata + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 1} + writer, err := newTestSinkSet(jsonlState, secondaryState, &collected).OpenShard(context.Background(), id) + if err != nil { + t.Fatalf("open shard: %v", err) + } + if err := writer.WriteBatch(context.Background(), []int{1}); err != nil { + t.Fatalf("write batch: %v", err) + } + _, err = writer.Finish(context.Background(), shardSummary{ID: id, Rows: 1}) + assertShardSinkOperation(t, err, failingFormat, id, "commit", cause) + jsonl := jsonlState.snapshot() + secondary := secondaryState.snapshot() + if failingFormat == jsonlFragmentFormat { + if jsonl.committed || secondary.commitAttempts != 0 || jsonl.preparedAborts != 1 || secondary.preparedAborts != 1 { + t.Fatalf("JSONL commit failure: JSONL=%+v secondary=%+v", jsonl, secondary) + } + } else if !jsonl.committed || secondary.committed || jsonl.preparedAborts != 0 || secondary.preparedAborts != 1 { + t.Fatalf("secondary commit failure: JSONL=%+v secondary=%+v", jsonl, secondary) + } + }) + } +} + +func TestShardSinkSetSingleSinkUsesDirectPath(t *testing.T) { + state := newFanoutSinkState[int]("jsonl", nil) + started := make(chan struct{}) + release := make(chan struct{}) + type contextKey struct{} + callerCtx := context.WithValue(context.Background(), contextKey{}, "caller") + state.writeHook = func(ctx context.Context) error { + if ctx != callerCtx { + return errors.New("single-sink path replaced the caller context") + } + close(started) + <-release + return nil + } + output := newShardSinkSet(newJSONLShardSink(newFanoutJSONLSink(state))) + writer, err := output.OpenShard(callerCtx, shardID{Graph: "example", Phase: PhaseNodes, Number: 1}) + if err != nil { + t.Fatalf("open shard: %v", err) + } + returned := make(chan error, 1) + go func() { returned <- writer.WriteBatch(callerCtx, []int{1}) }() + waitForFanoutSignal(t, started, "direct write") + select { + case err := <-returned: + t.Fatalf("single sink did not backpressure its caller: %v", err) + default: + } + close(release) + if err := <-returned; err != nil { + t.Fatalf("write batch: %v", err) + } + if state.snapshot().writerAborts != 0 { + t.Fatalf("single-sink write unexpectedly aborted: %+v", state.snapshot()) + } + if err := writer.Abort(); err != nil { + t.Fatalf("abort: %v", err) + } +} + +func eventPosition(events []string, target string) int { + for index, event := range events { + if event == target { + return index + } + } + return -1 +} + +func waitForFanoutSignal(t *testing.T, signal <-chan struct{}, description string) { + t.Helper() + select { + case <-signal: + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for %s", description) + } +} diff --git a/retriever/output_test.go b/retriever/output_test.go new file mode 100644 index 00000000..70215dce --- /dev/null +++ b/retriever/output_test.go @@ -0,0 +1,216 @@ +package retriever + +import ( + "context" + "errors" + "reflect" + "testing" +) + +type recordingJSONLSinkState[T any] struct { + fragments []*recordingJSONLFragmentState[T] + rowOffset int + writeErr error +} + +type recordingJSONLFragmentState[T any] struct { + id shardID + batches [][]T + prepared bool + committed bool + writerAborted bool + preparedAborted bool +} + +type recordingJSONLFragmentSink[T any] struct { + state *recordingJSONLSinkState[T] +} + +func (s recordingJSONLFragmentSink[T]) Open(_ context.Context, id shardID) (fragmentWriter[T, jsonlFragmentMetadata], error) { + fragment := &recordingJSONLFragmentState[T]{id: id} + s.state.fragments = append(s.state.fragments, fragment) + return &recordingJSONLFragmentWriter[T]{sink: s.state, fragment: fragment}, nil +} + +type recordingJSONLFragmentWriter[T any] struct { + sink *recordingJSONLSinkState[T] + fragment *recordingJSONLFragmentState[T] +} + +func (s *recordingJSONLFragmentWriter[T]) WriteBatch(_ context.Context, records []T) error { + if s.sink.writeErr != nil { + return s.sink.writeErr + } + s.fragment.batches = append(s.fragment.batches, append([]T(nil), records...)) + return nil +} + +func (s *recordingJSONLFragmentWriter[T]) Prepare(context.Context) (preparedFragment[jsonlFragmentMetadata], error) { + s.fragment.prepared = true + return &recordingPreparedJSONLFragment[T]{ + fragment: s.fragment, + metadata: jsonlFragmentMetadata{ + Path: "fragment.jsonl.gz", + Rows: countBatchRows(s.fragment.batches) + s.sink.rowOffset, + }, + }, nil +} + +func (s *recordingJSONLFragmentWriter[T]) Abort() error { + s.fragment.writerAborted = true + return nil +} + +type recordingPreparedJSONLFragment[T any] struct { + fragment *recordingJSONLFragmentState[T] + metadata jsonlFragmentMetadata +} + +func (s *recordingPreparedJSONLFragment[T]) Metadata() jsonlFragmentMetadata { + return s.metadata +} + +func (s *recordingPreparedJSONLFragment[T]) Commit(context.Context) error { + s.fragment.committed = true + return nil +} + +func (s *recordingPreparedJSONLFragment[T]) Abort() error { + s.fragment.preparedAborted = true + return nil +} + +func TestJSONLShardOutputStreamsLogicalShardSlices(t *testing.T) { + state := &recordingJSONLSinkState[int]{} + output := newShardSinkSet(newJSONLShardSink(recordingJSONLFragmentSink[int]{state: state})) + var committed []shardSummary + receiver := newShardOutputReceiver(context.Background(), output, func(summary shardSummary, result committedShard) error { + if result.JSONL.Rows != summary.Rows { + t.Fatalf("committed metadata = %+v, summary = %+v", result.JSONL, summary) + } + committed = append(committed, summary) + return nil + }) + sharder, err := newLogicalSharder[int]("example", PhaseNodes, 3) + if err != nil { + t.Fatalf("new sharder: %v", err) + } + + first := transformedBatch[int]{Records: []int{1, 2}, ActionCounts: make([]map[string]int, 2)} + if err := sharder.Add(first, receiver); err != nil { + t.Fatalf("add first batch: %v", err) + } + if len(state.fragments) != 1 || !reflect.DeepEqual(state.fragments[0].batches, [][]int{{1, 2}}) || state.fragments[0].prepared { + t.Fatalf("first batch was not streamed into the active writer: %+v", state.fragments) + } + + second := transformedBatch[int]{Records: []int{3, 4}, ActionCounts: make([]map[string]int, 2)} + if err := sharder.Add(second, receiver); err != nil { + t.Fatalf("add second batch: %v", err) + } + if err := sharder.Flush(receiver); err != nil { + t.Fatalf("flush: %v", err) + } + + if len(state.fragments) != 2 || len(committed) != 2 { + t.Fatalf("fragments=%d committed=%d", len(state.fragments), len(committed)) + } + if !reflect.DeepEqual(state.fragments[0].batches, [][]int{{1, 2}, {3}}) || !reflect.DeepEqual(state.fragments[1].batches, [][]int{{4}}) { + t.Fatalf("fragment batches = %+v", state.fragments) + } + for index, fragment := range state.fragments { + if fragment.id.Number != index+1 || !fragment.prepared || !fragment.committed || fragment.writerAborted || fragment.preparedAborted { + t.Fatalf("fragment %d lifecycle = %+v", index, fragment) + } + } +} + +func TestJSONLShardOutputAbortsPreparedRowCountMismatch(t *testing.T) { + state := &recordingJSONLSinkState[int]{rowOffset: 1} + output := newShardSinkSet(newJSONLShardSink(recordingJSONLFragmentSink[int]{state: state})) + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 1} + writer, err := output.OpenShard(context.Background(), id) + if err != nil { + t.Fatalf("open shard: %v", err) + } + if err := writer.WriteBatch(context.Background(), []int{1}); err != nil { + t.Fatalf("write batch: %v", err) + } + _, err = writer.Finish(context.Background(), shardSummary{ID: id, Rows: 1}) + assertShardSinkOperation(t, err, jsonlFragmentFormat, id, "validate", nil) + + fragment := state.fragments[0] + if !fragment.prepared || !fragment.preparedAborted || fragment.committed || fragment.writerAborted { + t.Fatalf("mismatch lifecycle = %+v", fragment) + } +} + +func assertShardSinkOperation(t *testing.T, err error, format string, id shardID, operation string, cause error) { + t.Helper() + if err == nil { + t.Fatalf("expected %s %s error", format, operation) + } + if cause != nil && !errors.Is(err, cause) { + t.Fatalf("%s %s error = %v, want cause %v", format, operation, err, cause) + } + + if !containsShardSinkOperation(err, format, id, operation) { + t.Fatalf("error lacks %s %s context for %+v: %v", format, operation, id, err) + } +} + +func containsShardSinkOperation(err error, format string, id shardID, operation string) bool { + if err == nil { + return false + } + if operationErr, ok := err.(*shardSinkOperationError); ok && operationErr.Format == format && operationErr.ID == id && operationErr.Operation == operation { + return true + } + if joined, ok := err.(interface{ Unwrap() []error }); ok { + for _, child := range joined.Unwrap() { + if containsShardSinkOperation(child, format, id, operation) { + return true + } + } + return false + } + return containsShardSinkOperation(errors.Unwrap(err), format, id, operation) +} + +func TestShardOutputReceiverAbortsAfterWriteAndUpstreamFailures(t *testing.T) { + t.Run("write failure", func(t *testing.T) { + writeErr := errors.New("write failed") + state := &recordingJSONLSinkState[int]{writeErr: writeErr} + receiver := newShardOutputReceiver(context.Background(), newShardSinkSet(newJSONLShardSink(recordingJSONLFragmentSink[int]{state: state})), func(shardSummary, committedShard) error { + return nil + }) + if err := receiver.BeginShard(shardID{Graph: "example", Phase: PhaseNodes, Number: 1}); err != nil { + t.Fatalf("begin shard: %v", err) + } + if err := receiver.WriteBatch([]int{1}); !errors.Is(err, writeErr) { + t.Fatalf("write error = %v", err) + } + if !state.fragments[0].writerAborted { + t.Fatalf("writer was not aborted after write failure") + } + }) + + t.Run("upstream failure", func(t *testing.T) { + state := &recordingJSONLSinkState[int]{} + receiver := newShardOutputReceiver(context.Background(), newShardSinkSet(newJSONLShardSink(recordingJSONLFragmentSink[int]{state: state})), func(shardSummary, committedShard) error { + return nil + }) + if err := receiver.BeginShard(shardID{Graph: "example", Phase: PhaseNodes, Number: 1}); err != nil { + t.Fatalf("begin shard: %v", err) + } + if err := receiver.WriteBatch([]int{1}); err != nil { + t.Fatalf("write batch: %v", err) + } + if err := receiver.Abort(); err != nil { + t.Fatalf("abort receiver: %v", err) + } + if !state.fragments[0].writerAborted { + t.Fatalf("writer was not aborted after upstream failure") + } + }) +} diff --git a/retriever/publisher_jsonl.go b/retriever/publisher_jsonl.go new file mode 100644 index 00000000..8db044b1 --- /dev/null +++ b/retriever/publisher_jsonl.go @@ -0,0 +1,68 @@ +package retriever + +import "context" + +type collectionPublication struct { + Manifest Manifest + Path string +} + +type collectionPublisher interface { + AddGraph(GraphManifest, GraphSchemaMetadata, GraphMetrics) + Publish(context.Context) (collectionPublication, error) +} + +type jsonlCollectionPublisher struct { + workspace collectionWorkspace + manifest Manifest + metrics MetricsManifest +} + +func newJSONLFileManifest(summary shardSummary, metadata jsonlFragmentMetadata) FileManifest { + return FileManifest{ + Phase: summary.ID.Phase, + Path: metadata.Path, + Count: metadata.Rows, + CompressedBytes: metadata.CompressedBytes, + UncompressedBytes: metadata.UncompressedBytes, + SHA256: metadata.SHA256, + ActionCounts: cloneActionCounts(summary.ActionCounts), + } +} + +func newJSONLCollectionPublisher(workspace collectionWorkspace, driverName string, options DumpOptions, scrub ScrubMetadata, graphCount int) *jsonlCollectionPublisher { + return &jsonlCollectionPublisher{ + workspace: workspace, + manifest: newManifest(driverName, options.Compression, options.ZstdLevel, scrub, graphCount), + metrics: newMetricsManifest(graphCount), + } +} + +func (s *jsonlCollectionPublisher) AddGraph(graphEntry GraphManifest, schemaEntry GraphSchemaMetadata, metricsEntry GraphMetrics) { + s.manifest.Graphs = append(s.manifest.Graphs, graphEntry) + s.manifest.Schema.Graphs = append(s.manifest.Schema.Graphs, schemaEntry) + s.metrics.Graphs = append(s.metrics.Graphs, metricsEntry) + addActionCounts(s.manifest.Scrub.NodeActionCounts, graphEntry.NodeActionCounts) + addActionCounts(s.manifest.Scrub.EdgeActionCounts, graphEntry.EdgeActionCounts) +} + +func (s *jsonlCollectionPublisher) Publish(ctx context.Context) (collectionPublication, error) { + nextManifest := s.manifest + nextMetrics := s.metrics + nextManifest.Metrics = &nextMetrics + + payload, err := encodeManifest(nextManifest) + if err != nil { + return collectionPublication{}, err + } + + manifestPath, err := s.workspace.Publish(ctx, manifestFileName, payload) + if err != nil { + return collectionPublication{}, err + } + + return collectionPublication{ + Manifest: nextManifest, + Path: manifestPath, + }, nil +} diff --git a/retriever/publisher_jsonl_test.go b/retriever/publisher_jsonl_test.go new file mode 100644 index 00000000..b51cb294 --- /dev/null +++ b/retriever/publisher_jsonl_test.go @@ -0,0 +1,85 @@ +package retriever + +import ( + "context" + "reflect" + "testing" +) + +func TestNewJSONLFileManifestCombinesLogicalAndPhysicalMetadata(t *testing.T) { + actionCounts := map[string]int{"redact": 2} + summary := shardSummary{ + ID: shardID{Graph: "example", Phase: PhaseNodes, Number: 3}, + Rows: 4, + ActionCounts: actionCounts, + } + metadata := jsonlFragmentMetadata{ + Path: "graphs/example/nodes-000003.jsonl.gz", + Rows: 4, + CompressedBytes: 100, + UncompressedBytes: 200, + SHA256: "checksum", + } + + actual := newJSONLFileManifest(summary, metadata) + expected := FileManifest{ + Phase: PhaseNodes, + Path: metadata.Path, + Count: 4, + CompressedBytes: 100, + UncompressedBytes: 200, + SHA256: "checksum", + ActionCounts: map[string]int{"redact": 2}, + } + if !reflect.DeepEqual(actual, expected) { + t.Fatalf("file manifest = %+v, want %+v", actual, expected) + } + + actionCounts["redact"] = 9 + if actual.ActionCounts["redact"] != 2 { + t.Fatalf("file manifest retained logical action-count map") + } +} + +func TestJSONLCollectionPublisherAggregatesGraphMetadata(t *testing.T) { + workspace := newLocalCollectionWorkspace(t.TempDir(), false) + options := DefaultDumpOptions(workspace.Root()) + publisher := newJSONLCollectionPublisher(workspace, "scripted", options, ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + }, 1) + graphEntry := GraphManifest{ + Name: "source", + NodeActionCounts: map[string]int{"drop": 2}, + EdgeActionCounts: map[string]int{"redact": 1}, + } + schemaEntry := GraphSchemaMetadata{Name: "source", NodeKinds: []string{}, EdgeKinds: []string{}} + metricsEntry := newMetricsBuilder("source", 0).finalize() + publisher.AddGraph(graphEntry, schemaEntry, metricsEntry) + + publication, err := publisher.Publish(context.Background()) + if err != nil { + t.Fatalf("publish collection: %v", err) + } + stored, err := readManifest(workspace.Root()) + if err != nil { + t.Fatalf("read published manifest: %v", err) + } + + if !reflect.DeepEqual(publication.Manifest, stored) { + t.Fatalf("published manifest differs from stored manifest") + } + if !reflect.DeepEqual(stored.Graphs, []GraphManifest{graphEntry}) { + t.Fatalf("manifest graphs = %+v", stored.Graphs) + } + if !reflect.DeepEqual(stored.Schema.Graphs, []GraphSchemaMetadata{schemaEntry}) { + t.Fatalf("manifest schema = %+v", stored.Schema.Graphs) + } + if stored.Metrics == nil || !reflect.DeepEqual(stored.Metrics.Graphs, []GraphMetrics{metricsEntry}) { + t.Fatalf("manifest metrics = %+v", stored.Metrics) + } + if stored.Scrub.NodeActionCounts["drop"] != 2 || stored.Scrub.EdgeActionCounts["redact"] != 1 { + t.Fatalf("manifest action counts = %+v", stored.Scrub) + } +} diff --git a/retriever/publisher_parquet.go b/retriever/publisher_parquet.go new file mode 100644 index 00000000..62d3cc9e --- /dev/null +++ b/retriever/publisher_parquet.go @@ -0,0 +1,207 @@ +package retriever + +import ( + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "time" +) + +const ( + parquetManifestFormat = "retriever-parquet-export-v1" + parquetManifestFileName = "parquet/manifest.json" + parquetSuccessFileName = "parquet/_SUCCESS" +) + +type parquetManifest struct { + Format string `json:"format"` + GeneratedAt time.Time `json:"generated_at"` + Compression string `json:"compression"` + PropertiesEncoding string `json:"properties_encoding"` + Graphs []parquetGraphManifest `json:"graphs"` +} + +type parquetGraphManifest struct { + Name string `json:"name"` + NodeCount int64 `json:"node_count"` + EdgeCount int64 `json:"edge_count"` + Files []parquetFileManifest `json:"files"` +} + +type parquetFileManifest struct { + Phase Phase `json:"phase"` + Path string `json:"path"` + Count int `json:"count"` + Bytes int64 `json:"bytes"` + SHA256 string `json:"sha256"` + ActionCounts map[string]int `json:"action_counts"` +} + +type parquetPublication struct { + ManifestPath string + SuccessPath string +} + +type parquetPublisher interface { + AddFragment(shardSummary, parquetFragmentMetadata) + AddGraph(string, int64, int64) + PublishManifest(context.Context) (string, error) + PublishSuccess(context.Context) (string, error) +} + +type parquetCollectionPublisher struct { + workspace collectionWorkspace + manifest parquetManifest + files map[string][]parquetFileManifest +} + +func newParquetCollectionPublisher(workspace collectionWorkspace, graphCount int) *parquetCollectionPublisher { + return &parquetCollectionPublisher{ + workspace: workspace, + manifest: parquetManifest{ + Format: parquetManifestFormat, + GeneratedAt: time.Now().UTC(), + Compression: "zstd", + PropertiesEncoding: "json", + Graphs: make([]parquetGraphManifest, 0, graphCount), + }, + files: make(map[string][]parquetFileManifest, graphCount), + } +} + +func (s *parquetCollectionPublisher) AddFragment(summary shardSummary, metadata parquetFragmentMetadata) { + s.files[summary.ID.Graph] = append(s.files[summary.ID.Graph], parquetFileManifest{ + Phase: summary.ID.Phase, + Path: metadata.Path, + Count: metadata.Rows, + Bytes: metadata.Bytes, + SHA256: metadata.SHA256, + ActionCounts: cloneActionCounts(summary.ActionCounts), + }) +} + +func (s *parquetCollectionPublisher) AddGraph(name string, nodeCount, edgeCount int64) { + s.manifest.Graphs = append(s.manifest.Graphs, parquetGraphManifest{ + Name: name, + NodeCount: nodeCount, + EdgeCount: edgeCount, + Files: append([]parquetFileManifest(nil), s.files[name]...), + }) +} + +func (s *parquetCollectionPublisher) PublishManifest(ctx context.Context) (string, error) { + payload, err := encodeParquetManifest(s.manifest) + if err != nil { + return "", err + } + return s.workspace.Publish(ctx, parquetManifestFileName, payload) +} + +func (s *parquetCollectionPublisher) PublishSuccess(ctx context.Context) (string, error) { + return s.workspace.Publish(ctx, parquetSuccessFileName, []byte(parquetManifestFormat+"\n")) +} + +func publishDumpOutputs(ctx context.Context, jsonl collectionPublisher, parquet parquetPublisher) (collectionPublication, parquetPublication, error) { + var parquetResult parquetPublication + if parquet != nil { + path, err := parquet.PublishManifest(ctx) + if err != nil { + return collectionPublication{}, parquetResult, err + } + parquetResult.ManifestPath = path + } + + jsonlResult, err := jsonl.Publish(ctx) + if err != nil { + return collectionPublication{}, parquetResult, err + } + + if parquet != nil { + path, err := parquet.PublishSuccess(ctx) + if err != nil { + return jsonlResult, parquetResult, err + } + parquetResult.SuccessPath = path + } + return jsonlResult, parquetResult, nil +} + +func readParquetManifest(outputDir string) (parquetManifest, error) { + var manifest parquetManifest + payload, err := os.ReadFile(filepath.Join(outputDir, filepath.FromSlash(parquetManifestFileName))) + if err != nil { + return manifest, fmt.Errorf("read Parquet manifest: %w", err) + } + if err := json.Unmarshal(payload, &manifest); err != nil { + return manifest, fmt.Errorf("decode Parquet manifest: %w", err) + } + if err := manifest.validate(); err != nil { + return manifest, err + } + return manifest, nil +} + +func encodeParquetManifest(manifest parquetManifest) ([]byte, error) { + if err := manifest.validate(); err != nil { + return nil, err + } + payload, err := json.MarshalIndent(manifest, "", " ") + if err != nil { + return nil, fmt.Errorf("encode Parquet manifest: %w", err) + } + return append(payload, '\n'), nil +} + +func (s parquetManifest) validate() error { + if s.Format != parquetManifestFormat { + return fmt.Errorf("unsupported Parquet manifest format %q", s.Format) + } + if s.Compression != "zstd" { + return fmt.Errorf("unsupported Parquet compression %q", s.Compression) + } + if s.PropertiesEncoding != "json" { + return fmt.Errorf("unsupported Parquet properties encoding %q", s.PropertiesEncoding) + } + + seenGraphs := make(map[string]struct{}, len(s.Graphs)) + seenPaths := map[string]struct{}{} + for _, graph := range s.Graphs { + if graph.Name == "" { + return fmt.Errorf("Parquet manifest graph entry has empty name") + } + if _, seen := seenGraphs[graph.Name]; seen { + return fmt.Errorf("Parquet manifest contains duplicate graph %q", graph.Name) + } + seenGraphs[graph.Name] = struct{}{} + + var nodes, edges int64 + seenEdgePhase := false + for _, file := range graph.Files { + if file.Path == "" || file.Count < 0 || file.Bytes < 0 || file.SHA256 == "" { + return fmt.Errorf("Parquet manifest contains invalid file entry %+v", file) + } + if _, seen := seenPaths[file.Path]; seen { + return fmt.Errorf("Parquet manifest contains duplicate path %q", file.Path) + } + seenPaths[file.Path] = struct{}{} + switch file.Phase { + case PhaseNodes: + if seenEdgePhase { + return fmt.Errorf("Parquet manifest graph %q lists node file after edge file", graph.Name) + } + nodes += int64(file.Count) + case PhaseEdges: + seenEdgePhase = true + edges += int64(file.Count) + default: + return fmt.Errorf("Parquet manifest graph %q contains unsupported phase %q", graph.Name, file.Phase) + } + } + if graph.NodeCount != nodes || graph.EdgeCount != edges { + return fmt.Errorf("Parquet manifest graph %q counts (%d nodes, %d edges) do not match file totals (%d nodes, %d edges)", graph.Name, graph.NodeCount, graph.EdgeCount, nodes, edges) + } + } + return nil +} diff --git a/retriever/publisher_parquet_test.go b/retriever/publisher_parquet_test.go new file mode 100644 index 00000000..96c3fa5c --- /dev/null +++ b/retriever/publisher_parquet_test.go @@ -0,0 +1,129 @@ +package retriever + +import ( + "context" + "errors" + "reflect" + "testing" +) + +func TestParquetCollectionPublisherCombinesLogicalAndPhysicalMetadata(t *testing.T) { + workspace := newLocalCollectionWorkspace(t.TempDir(), false) + publisher := newParquetCollectionPublisher(workspace, 1) + summary := shardSummary{ + ID: shardID{Graph: "source", Phase: PhaseNodes, Number: 1}, + Rows: 2, + ActionCounts: map[string]int{"redact": 1}, + } + publisher.AddFragment(summary, parquetFragmentMetadata{ + Path: "parquet/graphs/source/nodes-000001.parquet", + Rows: 2, + Bytes: 100, + SHA256: "checksum", + }) + publisher.AddGraph("source", 2, 0) + + manifestPath, err := publisher.PublishManifest(context.Background()) + if err != nil { + t.Fatalf("publish manifest: %v", err) + } + successPath, err := publisher.PublishSuccess(context.Background()) + if err != nil { + t.Fatalf("publish success: %v", err) + } + stored, err := readParquetManifest(workspace.Root()) + if err != nil { + t.Fatalf("read manifest: %v", err) + } + if manifestPath == "" || successPath == "" || len(stored.Graphs) != 1 || len(stored.Graphs[0].Files) != 1 { + t.Fatalf("publication paths or manifest invalid: manifest=%q success=%q stored=%+v", manifestPath, successPath, stored) + } + file := stored.Graphs[0].Files[0] + if file.Phase != PhaseNodes || file.Count != 2 || file.Bytes != 100 || file.ActionCounts["redact"] != 1 { + t.Fatalf("Parquet file manifest = %+v", file) + } + summary.ActionCounts["redact"] = 9 + if file.ActionCounts["redact"] != 1 { + t.Fatalf("Parquet manifest retained logical action-count map") + } +} + +type recordingCollectionPublisher struct { + events *[]string + err error +} + +func (s recordingCollectionPublisher) AddGraph(GraphManifest, GraphSchemaMetadata, GraphMetrics) {} + +func (s recordingCollectionPublisher) Publish(context.Context) (collectionPublication, error) { + *s.events = append(*s.events, "jsonl-manifest") + return collectionPublication{Path: "manifest.json"}, s.err +} + +type recordingParquetPublisher struct { + events *[]string + manifestErr error + successErr error +} + +func (s recordingParquetPublisher) AddFragment(shardSummary, parquetFragmentMetadata) {} + +func (s recordingParquetPublisher) AddGraph(string, int64, int64) {} + +func (s recordingParquetPublisher) PublishManifest(context.Context) (string, error) { + *s.events = append(*s.events, "parquet-manifest") + return parquetManifestFileName, s.manifestErr +} + +func (s recordingParquetPublisher) PublishSuccess(context.Context) (string, error) { + *s.events = append(*s.events, "parquet-success") + return parquetSuccessFileName, s.successErr +} + +func TestDumpOutputPublicationOrder(t *testing.T) { + manifestFailure := errors.New("Parquet manifest failed") + jsonlFailure := errors.New("JSONL manifest failed") + successFailure := errors.New("Parquet success failed") + tests := []struct { + name string + jsonlErr error + manifestErr error + successErr error + wantEvents []string + wantErr error + }{ + {name: "success", wantEvents: []string{"parquet-manifest", "jsonl-manifest", "parquet-success"}}, + {name: "Parquet manifest failure", manifestErr: manifestFailure, wantEvents: []string{"parquet-manifest"}, wantErr: manifestFailure}, + {name: "JSONL manifest failure", jsonlErr: jsonlFailure, wantEvents: []string{"parquet-manifest", "jsonl-manifest"}, wantErr: jsonlFailure}, + {name: "Parquet success failure", successErr: successFailure, wantEvents: []string{"parquet-manifest", "jsonl-manifest", "parquet-success"}, wantErr: successFailure}, + } + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + var events []string + jsonl := recordingCollectionPublisher{events: &events, err: test.jsonlErr} + parquet := recordingParquetPublisher{events: &events, manifestErr: test.manifestErr, successErr: test.successErr} + jsonlResult, parquetResult, err := publishDumpOutputs(context.Background(), jsonl, parquet) + if !errors.Is(err, test.wantErr) { + t.Fatalf("publish error = %v, want %v", err, test.wantErr) + } + if !reflect.DeepEqual(events, test.wantEvents) { + t.Fatalf("publication events = %v, want %v", events, test.wantEvents) + } + if test.wantErr == nil && (jsonlResult.Path == "" || parquetResult.ManifestPath == "" || parquetResult.SuccessPath == "") { + t.Fatalf("successful publication results = JSONL %+v Parquet %+v", jsonlResult, parquetResult) + } + }) + } +} + +func TestDumpOutputPublicationWithoutParquetUsesJSONLOnly(t *testing.T) { + var events []string + jsonl := recordingCollectionPublisher{events: &events} + jsonlResult, parquetResult, err := publishDumpOutputs(context.Background(), jsonl, nil) + if err != nil { + t.Fatalf("publish: %v", err) + } + if !reflect.DeepEqual(events, []string{"jsonl-manifest"}) || jsonlResult.Path == "" || parquetResult != (parquetPublication{}) { + t.Fatalf("JSONL-only publication = events %v, JSONL %+v, Parquet %+v", events, jsonlResult, parquetResult) + } +} diff --git a/retriever/records.go b/retriever/records.go new file mode 100644 index 00000000..eaa4a712 --- /dev/null +++ b/retriever/records.go @@ -0,0 +1,71 @@ +package retriever + +import ( + "maps" + "sort" + + "github.com/specterops/dawgs/graph" +) + +// normalizedNode is the output-neutral node representation owned by the dump +// pipeline. Its kinds and top-level property map do not alias the source graph +// entity; property values are treated as immutable. +type normalizedNode struct { + ID string + Kinds []string + Properties map[string]any +} + +// normalizedEdge is the output-neutral relationship representation owned by +// the dump pipeline. ID is retained for sinks that need source relationship +// identity even though the JSONL v1 adapter intentionally omits it. +type normalizedEdge struct { + ID string + StartID string + EndID string + Kind string + Properties map[string]any +} + +func normalizeNode(node *graph.Node) normalizedNode { + kinds := node.Kinds.Strings() + sort.Strings(kinds) + + return normalizedNode{ + ID: node.ID.String(), + Kinds: kinds, + Properties: maps.Clone(node.Properties.MapOrEmpty()), + } +} + +func normalizeEdge(relationship *graph.Relationship) normalizedEdge { + kind := "" + if relationship.Kind != nil { + kind = relationship.Kind.String() + } + + return normalizedEdge{ + ID: relationship.ID.String(), + StartID: relationship.StartID.String(), + EndID: relationship.EndID.String(), + Kind: kind, + Properties: maps.Clone(relationship.Properties.MapOrEmpty()), + } +} + +func jsonlV1NodeFromNormalized(node normalizedNode) FragmentNode { + return FragmentNode{ + ID: node.ID, + Kinds: node.Kinds, + Properties: node.Properties, + } +} + +func jsonlV1EdgeFromNormalized(edge normalizedEdge) FragmentEdge { + return FragmentEdge{ + StartID: edge.StartID, + EndID: edge.EndID, + Kind: edge.Kind, + Properties: edge.Properties, + } +} diff --git a/retriever/records_test.go b/retriever/records_test.go new file mode 100644 index 00000000..63262fc8 --- /dev/null +++ b/retriever/records_test.go @@ -0,0 +1,134 @@ +package retriever + +import ( + "reflect" + "testing" + + "github.com/specterops/dawgs/graph" +) + +func TestNormalizeNode(t *testing.T) { + properties := map[string]any{ + "nested": map[string]any{"names": []string{"alice", "bob"}}, + "bytes": []byte{1, 2, 3}, + "nil": nil, + } + node := graph.NewNode( + 42, + graph.AsProperties(properties), + graph.StringKind("User"), + graph.StringKind("Admin"), + ) + + normalized := normalizeNode(node) + if normalized.ID != "42" { + t.Fatalf("normalized ID = %q", normalized.ID) + } + if !reflect.DeepEqual(normalized.Kinds, []string{"Admin", "User"}) { + t.Fatalf("normalized kinds = %v", normalized.Kinds) + } + if !reflect.DeepEqual(node.Kinds.Strings(), []string{"User", "Admin"}) { + t.Fatalf("normalization modified source kinds: %v", node.Kinds.Strings()) + } + if !reflect.DeepEqual(normalized.Properties, properties) { + t.Fatalf("normalized properties = %#v, want %#v", normalized.Properties, properties) + } + + properties["new"] = "source-only" + if _, found := normalized.Properties["new"]; found { + t.Fatalf("normalized properties alias source map: %#v", normalized.Properties) + } + + normalized.Kinds[0] = "Changed" + normalized.Properties["other"] = "normalized-only" + if !reflect.DeepEqual(node.Kinds.Strings(), []string{"User", "Admin"}) { + t.Fatalf("normalized kinds alias source kinds: %v", node.Kinds.Strings()) + } + if _, found := properties["other"]; found { + t.Fatalf("source properties changed through normalized map: %#v", properties) + } +} + +func TestNormalizeNodeWithoutPropertiesOwnsEmptyMap(t *testing.T) { + normalized := normalizeNode(graph.NewNode(1, nil)) + if normalized.Properties == nil || len(normalized.Properties) != 0 { + t.Fatalf("normalized properties = %#v", normalized.Properties) + } +} + +func TestNormalizeEdge(t *testing.T) { + properties := map[string]any{ + "routes": []any{"north", map[string]any{"weight": int64(2)}}, + } + relationship := graph.NewRelationship( + 99, + 10, + 20, + graph.AsProperties(properties), + graph.StringKind("AdminTo"), + ) + + normalized := normalizeEdge(relationship) + expected := normalizedEdge{ + ID: "99", + StartID: "10", + EndID: "20", + Kind: "AdminTo", + Properties: map[string]any{"routes": []any{"north", map[string]any{"weight": int64(2)}}}, + } + if !reflect.DeepEqual(normalized, expected) { + t.Fatalf("normalized edge = %#v, want %#v", normalized, expected) + } + + properties["routes"] = []any{"changed"} + if !reflect.DeepEqual(normalized, expected) { + t.Fatalf("normalized edge properties alias source map: %#v", normalized.Properties) + } +} + +func TestNormalizeEdgeWithoutKind(t *testing.T) { + normalized := normalizeEdge(graph.NewRelationship(3, 1, 2, nil, nil)) + if normalized.Kind != "" { + t.Fatalf("normalized missing kind = %q", normalized.Kind) + } + if normalized.Properties == nil || len(normalized.Properties) != 0 { + t.Fatalf("normalized properties = %#v", normalized.Properties) + } +} + +func TestJSONLV1NormalizedRecordAdapters(t *testing.T) { + node := normalizedNode{ + ID: "1", + Kinds: []string{"Admin", "User"}, + Properties: map[string]any{"name": "alice"}, + } + if actual := jsonlV1NodeFromNormalized(node); !reflect.DeepEqual(actual, FragmentNode{ + ID: "1", + Kinds: []string{"Admin", "User"}, + Properties: map[string]any{"name": "alice"}, + }) { + t.Fatalf("JSONL node adapter = %#v", actual) + } + + edge := normalizedEdge{ + ID: "source-edge-id", + StartID: "1", + EndID: "2", + Kind: "AdminTo", + Properties: map[string]any{"route": "north"}, + } + expectedEdge := FragmentEdge{ + StartID: "1", + EndID: "2", + Kind: "AdminTo", + Properties: map[string]any{"route": "north"}, + } + if actual := jsonlV1EdgeFromNormalized(edge); !reflect.DeepEqual(actual, expectedEdge) { + t.Fatalf("JSONL edge adapter = %#v, want %#v", actual, expectedEdge) + } + + edge.ID = "different-source-edge-id" + if actual := jsonlV1EdgeFromNormalized(edge); !reflect.DeepEqual(actual, expectedEdge) { + t.Fatalf("source edge ID affected JSONL v1 adapter: %#v", actual) + } +} diff --git a/retriever/scan.go b/retriever/scan.go index 938eda05..e1f7b10b 100644 --- a/retriever/scan.go +++ b/retriever/scan.go @@ -115,20 +115,14 @@ func scanDatabaseNodesWithProgressInterval( handle entityBatchHandler[*graph.Node], logProgress entityProgressLogger, ) (int64, error) { - return scanEntityBatches(entityScanOptions[*graph.Node]{ - Total: total, - BatchSize: batchSize, - ProgressInterval: progressInterval, - EntityName: "node", - Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Node, error) { - return readDatabaseNodes(ctx, db, targetGraph, afterID, hasAfterID, limit) - }, - ID: func(node *graph.Node) graph.ID { - return node.ID - }, - Handle: handle, - LogProgress: logProgress, - }) + return runFaucetWithProgress( + ctx, + newDatabaseGraphSource(db).Nodes(targetGraph, total, batchSize), + total, + progressInterval, + handle, + logProgress, + ) } func scanDatabaseRelationships( @@ -153,18 +147,48 @@ func scanDatabaseRelationshipsWithProgressInterval( handle entityBatchHandler[*graph.Relationship], logProgress entityProgressLogger, ) (int64, error) { - return scanEntityBatches(entityScanOptions[*graph.Relationship]{ - Total: total, - BatchSize: batchSize, - ProgressInterval: progressInterval, - EntityName: "relationship", - Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Relationship, error) { - return readDatabaseRelationships(ctx, db, targetGraph, afterID, hasAfterID, limit) - }, - ID: func(relationship *graph.Relationship) graph.ID { - return relationship.ID - }, - Handle: handle, - LogProgress: logProgress, + return runFaucetWithProgress( + ctx, + newDatabaseGraphSource(db).Edges(targetGraph, total, batchSize), + total, + progressInterval, + handle, + logProgress, + ) +} + +type batchProgressObserver struct { + processed int64 + startedAt time.Time + nextProgressAt int64 + logProgress entityProgressLogger +} + +func newBatchProgressObserver(total, progressInterval int64, logProgress entityProgressLogger) *batchProgressObserver { + return &batchProgressObserver{ + startedAt: time.Now(), + nextProgressAt: retrieverInitialProgressAtInterval(total, progressInterval), + logProgress: logProgress, + } +} + +func (s *batchProgressObserver) Observe(count int) { + s.processed += int64(count) + if s.logProgress != nil { + s.nextProgressAt = s.logProgress(s.processed, s.startedAt, s.nextProgressAt) + } +} + +func runFaucetWithProgress[T any](ctx context.Context, source faucet[T], total, progressInterval int64, handle entityBatchHandler[T], logProgress entityProgressLogger) (int64, error) { + progress := newBatchProgressObserver(total, progressInterval, logProgress) + return source.Run(ctx, func(batch []T) error { + if handle != nil { + if err := handle(batch); err != nil { + return err + } + } + + progress.Observe(len(batch)) + return nil }) } diff --git a/retriever/scan_test.go b/retriever/scan_test.go index d502c334..d0f0a6c0 100644 --- a/retriever/scan_test.go +++ b/retriever/scan_test.go @@ -1,6 +1,7 @@ package retriever import ( + "context" "errors" "reflect" "strings" @@ -219,3 +220,30 @@ func TestScanEntityBatchesValidationAndReaderErrors(t *testing.T) { t.Fatalf("processed=%d err=%v", processed, err) } } + +func TestRunFaucetWithProgressObservesSuccessfulBatchesOnce(t *testing.T) { + source := scriptedFaucet[scanTestEntity]{ + batches: [][]scanTestEntity{{{id: 1}, {id: 2}}, {{id: 3}}}, + total: 3, + } + var ( + handled []graph.ID + progress []int64 + ) + + processed, err := runFaucetWithProgress(context.Background(), source, 3, 1, func(batch []scanTestEntity) error { + for _, entity := range batch { + handled = append(handled, entity.id) + } + return nil + }, func(processed int64, _ time.Time, nextProgressAt int64) int64 { + progress = append(progress, processed) + return nextProgressAt + }) + if err != nil { + t.Fatalf("run faucet: %v", err) + } + if processed != 3 || !reflect.DeepEqual(handled, []graph.ID{1, 2, 3}) || !reflect.DeepEqual(progress, []int64{2, 3}) { + t.Fatalf("processed=%d handled=%v progress=%v", processed, handled, progress) + } +} diff --git a/retriever/sharder.go b/retriever/sharder.go new file mode 100644 index 00000000..6b6c9aa4 --- /dev/null +++ b/retriever/sharder.go @@ -0,0 +1,110 @@ +package retriever + +import "fmt" + +type shardID struct { + Graph string + Phase Phase + Number int +} + +type shardSummary struct { + ID shardID + Rows int + ActionCounts map[string]int +} + +type logicalShardReceiver[T any] interface { + BeginShard(shardID) error + WriteBatch([]T) error + FinishShard(shardSummary) error +} + +type logicalSharder[T any] struct { + graph string + phase Phase + shardSize int + nextNumber int + active bool + rows int + actionCounts map[string]int +} + +func newLogicalSharder[T any](graphName string, phase Phase, shardSize int) (*logicalSharder[T], error) { + if shardSize <= 0 { + return nil, fmt.Errorf("shard size must be > 0") + } + + return &logicalSharder[T]{ + graph: graphName, + phase: phase, + shardSize: shardSize, + nextNumber: 1, + actionCounts: map[string]int{}, + }, nil +} + +func (s *logicalSharder[T]) Add(batch transformedBatch[T], receiver logicalShardReceiver[T]) error { + if len(batch.Records) != len(batch.ActionCounts) { + return fmt.Errorf("transformed batch has %d records and %d action-count entries", len(batch.Records), len(batch.ActionCounts)) + } + + for offset := 0; offset < len(batch.Records); { + if !s.active { + if err := receiver.BeginShard(s.id()); err != nil { + return err + } + s.active = true + } + + batchEnd := min(offset+s.shardSize-s.rows, len(batch.Records)) + if err := receiver.WriteBatch(batch.Records[offset:batchEnd]); err != nil { + return err + } + for _, actionCounts := range batch.ActionCounts[offset:batchEnd] { + addActionCounts(s.actionCounts, actionCounts) + } + s.rows += batchEnd - offset + offset = batchEnd + + if s.rows == s.shardSize { + if err := s.finish(receiver); err != nil { + return err + } + } + } + + return nil +} + +func (s *logicalSharder[T]) Flush(receiver logicalShardReceiver[T]) error { + if !s.active || s.rows == 0 { + return nil + } + return s.finish(receiver) +} + +func (s *logicalSharder[T]) id() shardID { + return shardID{ + Graph: s.graph, + Phase: s.phase, + Number: s.nextNumber, + } +} + +func (s *logicalSharder[T]) finish(receiver logicalShardReceiver[T]) error { + summary := shardSummary{ + ID: s.id(), + Rows: s.rows, + ActionCounts: cloneActionCounts(s.actionCounts), + } + if err := receiver.FinishShard(summary); err != nil { + return err + } + + s.nextNumber++ + s.active = false + s.rows = 0 + s.actionCounts = map[string]int{} + return nil +} diff --git a/retriever/sharder_test.go b/retriever/sharder_test.go new file mode 100644 index 00000000..fdbcdae6 --- /dev/null +++ b/retriever/sharder_test.go @@ -0,0 +1,281 @@ +package retriever + +import ( + "errors" + "reflect" + "testing" +) + +type recordedLogicalShard[T any] struct { + id shardID + batches [][]T + summary shardSummary + finished bool +} + +type recordingLogicalShardReceiver[T any] struct { + shards []recordedLogicalShard[T] +} + +func (s *recordingLogicalShardReceiver[T]) BeginShard(id shardID) error { + s.shards = append(s.shards, recordedLogicalShard[T]{id: id}) + return nil +} + +func (s *recordingLogicalShardReceiver[T]) WriteBatch(records []T) error { + index := len(s.shards) - 1 + batch := append([]T(nil), records...) + s.shards[index].batches = append(s.shards[index].batches, batch) + return nil +} + +func (s *recordingLogicalShardReceiver[T]) FinishShard(summary shardSummary) error { + index := len(s.shards) - 1 + s.shards[index].summary = summary + s.shards[index].finished = true + return nil +} + +func TestLogicalSharderStreamsExactBoundaries(t *testing.T) { + testCases := []struct { + name string + shardSize int + batches [][]int + expectedBatches [][][]int + }{ + {name: "empty", shardSize: 2}, + { + name: "exact boundaries", + shardSize: 2, + batches: [][]int{{1, 2}, {3, 4}}, + expectedBatches: [][][]int{{{1, 2}}, {{3, 4}}}, + }, + { + name: "partial final shard", + shardSize: 3, + batches: [][]int{{1, 2}, {3, 4}}, + expectedBatches: [][][]int{{{1, 2}, {3}}, {{4}}}, + }, + { + name: "batch larger than shard", + shardSize: 2, + batches: [][]int{{1, 2, 3, 4, 5}}, + expectedBatches: [][][]int{{{1, 2}}, {{3, 4}}, {{5}}}, + }, + } + + for _, testCase := range testCases { + t.Run(testCase.name, func(t *testing.T) { + sharder, err := newLogicalSharder[int]("example", PhaseNodes, testCase.shardSize) + if err != nil { + t.Fatalf("new sharder: %v", err) + } + + receiver := &recordingLogicalShardReceiver[int]{} + for _, records := range testCase.batches { + batch := transformedBatch[int]{ + Records: records, + ActionCounts: make([]map[string]int, len(records)), + } + if err := sharder.Add(batch, receiver); err != nil { + t.Fatalf("add batch: %v", err) + } + } + if err := sharder.Flush(receiver); err != nil { + t.Fatalf("flush: %v", err) + } + + if len(receiver.shards) != len(testCase.expectedBatches) { + t.Fatalf("shard count = %d, want %d", len(receiver.shards), len(testCase.expectedBatches)) + } + for index, shard := range receiver.shards { + if shard.id != (shardID{Graph: "example", Phase: PhaseNodes, Number: index + 1}) { + t.Fatalf("shard %d ID = %+v", index, shard.id) + } + if !shard.finished || shard.summary.ID != shard.id { + t.Fatalf("shard %d lifecycle = %+v", index, shard) + } + if !reflect.DeepEqual(shard.batches, testCase.expectedBatches[index]) { + t.Fatalf("shard %d batches = %v, want %v", index, shard.batches, testCase.expectedBatches[index]) + } + if shard.summary.Rows != countBatchRows(shard.batches) { + t.Fatalf("shard %d summary = %+v", index, shard.summary) + } + } + }) + } +} + +func TestLogicalSharderStreamsBeforeFlushWithoutRetainingRecords(t *testing.T) { + sharder, err := newLogicalSharder[int]("example", PhaseNodes, 1_000_000) + if err != nil { + t.Fatalf("new sharder: %v", err) + } + + receiver := &recordingLogicalShardReceiver[int]{} + batch := transformedBatch[int]{Records: []int{1, 2}, ActionCounts: make([]map[string]int, 2)} + if err := sharder.Add(batch, receiver); err != nil { + t.Fatalf("add batch: %v", err) + } + if len(receiver.shards) != 1 || !reflect.DeepEqual(receiver.shards[0].batches, [][]int{{1, 2}}) || receiver.shards[0].finished { + t.Fatalf("streamed state before flush = %+v", receiver.shards) + } + + batch.Records[0] = 99 + if err := sharder.Flush(receiver); err != nil { + t.Fatalf("flush: %v", err) + } + if !reflect.DeepEqual(receiver.shards[0].batches, [][]int{{1, 2}}) || receiver.shards[0].summary.Rows != 2 { + t.Fatalf("finished shard = %+v", receiver.shards[0]) + } +} + +type addressRecordingLogicalShardReceiver struct { + starts []*int + lengths []int +} + +func (*addressRecordingLogicalShardReceiver) BeginShard(shardID) error { + return nil +} + +func (s *addressRecordingLogicalShardReceiver) WriteBatch(records []int) error { + s.starts = append(s.starts, &records[0]) + s.lengths = append(s.lengths, len(records)) + return nil +} + +func (*addressRecordingLogicalShardReceiver) FinishShard(shardSummary) error { + return nil +} + +func TestLogicalSharderStreamsSourceBatchSlicesWithoutCopying(t *testing.T) { + sharder, err := newLogicalSharder[int]("example", PhaseNodes, 2) + if err != nil { + t.Fatalf("new sharder: %v", err) + } + + batch := transformedBatch[int]{Records: []int{1, 2, 3, 4, 5}, ActionCounts: make([]map[string]int, 5)} + receiver := &addressRecordingLogicalShardReceiver{} + if err := sharder.Add(batch, receiver); err != nil { + t.Fatalf("add batch: %v", err) + } + if err := sharder.Flush(receiver); err != nil { + t.Fatalf("flush: %v", err) + } + + if !reflect.DeepEqual(receiver.lengths, []int{2, 2, 1}) { + t.Fatalf("batch lengths = %v", receiver.lengths) + } + for index, offset := range []int{0, 2, 4} { + if receiver.starts[index] != &batch.Records[offset] { + t.Fatalf("batch %d does not share the transformed batch backing slice", index) + } + } +} + +func TestLogicalSharderSplitsActionCounts(t *testing.T) { + sharder, err := newLogicalSharder[string]("example", PhaseEdges, 2) + if err != nil { + t.Fatalf("new sharder: %v", err) + } + + batch := transformedBatch[string]{ + Records: []string{"a", "b", "c"}, + ActionCounts: []map[string]int{ + {"preserve": 1}, + {"redact": 1}, + {"preserve": 2}, + }, + } + receiver := &recordingLogicalShardReceiver[string]{} + if err := sharder.Add(batch, receiver); err != nil { + t.Fatalf("add batch: %v", err) + } + if err := sharder.Flush(receiver); err != nil { + t.Fatalf("flush: %v", err) + } + + if len(receiver.shards) != 2 || !reflect.DeepEqual(receiver.shards[0].summary.ActionCounts, map[string]int{"preserve": 1, "redact": 1}) || !reflect.DeepEqual(receiver.shards[1].summary.ActionCounts, map[string]int{"preserve": 2}) { + t.Fatalf("shard action counts = %+v", receiver.shards) + } +} + +type failingLogicalShardReceiver[T any] struct { + point string + failure error +} + +func (s failingLogicalShardReceiver[T]) BeginShard(shardID) error { + if s.point == "begin" { + return s.failure + } + return nil +} + +func (s failingLogicalShardReceiver[T]) WriteBatch([]T) error { + if s.point == "write" { + return s.failure + } + return nil +} + +func (s failingLogicalShardReceiver[T]) FinishShard(shardSummary) error { + if s.point == "finish" { + return s.failure + } + return nil +} + +func TestLogicalSharderValidationAndDownstreamErrors(t *testing.T) { + if _, err := newLogicalSharder[int]("example", PhaseNodes, 0); err == nil { + t.Fatalf("expected invalid shard size error") + } + + validBatch := transformedBatch[int]{Records: []int{1}, ActionCounts: make([]map[string]int, 1)} + for _, testCase := range []struct { + name string + point string + shardSize int + flush bool + }{ + {name: "begin", point: "begin", shardSize: 1}, + {name: "write", point: "write", shardSize: 1}, + {name: "finish at boundary", point: "finish", shardSize: 1}, + {name: "finish on flush", point: "finish", shardSize: 2, flush: true}, + } { + t.Run(testCase.name, func(t *testing.T) { + failure := errors.New("downstream failed") + receiver := failingLogicalShardReceiver[int]{point: testCase.point, failure: failure} + sharder, err := newLogicalSharder[int]("example", PhaseNodes, testCase.shardSize) + if err != nil { + t.Fatalf("new sharder: %v", err) + } + + err = sharder.Add(validBatch, receiver) + if testCase.flush && err == nil { + err = sharder.Flush(receiver) + } + if !errors.Is(err, failure) { + t.Fatalf("downstream error = %v", err) + } + }) + } + + sharder, err := newLogicalSharder[int]("example", PhaseNodes, 1) + if err != nil { + t.Fatalf("new sharder: %v", err) + } + invalidBatch := transformedBatch[int]{Records: []int{1}} + if err := sharder.Add(invalidBatch, &recordingLogicalShardReceiver[int]{}); err == nil { + t.Fatalf("expected mismatched transformed batch error") + } +} + +func countBatchRows[T any](batches [][]T) int { + var rows int + for _, batch := range batches { + rows += len(batch) + } + return rows +} diff --git a/retriever/sink.go b/retriever/sink.go new file mode 100644 index 00000000..72690a90 --- /dev/null +++ b/retriever/sink.go @@ -0,0 +1,39 @@ +package retriever + +import ( + "context" + "errors" +) + +func collectErrors(errs ...error) error { + return errors.Join(errs...) +} + +func cleanupOnError(primary error, cleanups ...func() error) error { + errs := make([]error, 1, len(cleanups)+1) + errs[0] = primary + for _, cleanup := range cleanups { + errs = append(errs, cleanup()) + } + return collectErrors(errs...) +} + +type fragmentMetadata interface { + rowCount() int +} + +type fragmentSink[T any, M fragmentMetadata] interface { + Open(context.Context, shardID) (fragmentWriter[T, M], error) +} + +type fragmentWriter[T any, M fragmentMetadata] interface { + WriteBatch(context.Context, []T) error + Prepare(context.Context) (preparedFragment[M], error) + Abort() error +} + +type preparedFragment[M fragmentMetadata] interface { + Metadata() M + Commit(context.Context) error + Abort() error +} diff --git a/retriever/sink_contract_test.go b/retriever/sink_contract_test.go new file mode 100644 index 00000000..3f6e3e89 --- /dev/null +++ b/retriever/sink_contract_test.go @@ -0,0 +1,383 @@ +package retriever + +import ( + "bytes" + "context" + "errors" + "os" + "path/filepath" + "testing" +) + +type leafSinkContract[T any, M fragmentMetadata] struct { + New func(*testing.T) (fragmentSink[T, M], shardID, string) + Record T +} + +func TestJSONLFragmentSinkConformance(t *testing.T) { + t.Run("nodes", func(t *testing.T) { + runLeafSinkContract(t, leafSinkContract[normalizedNode, jsonlFragmentMetadata]{ + New: func(t *testing.T) (fragmentSink[normalizedNode, jsonlFragmentMetadata], shardID, string) { + return newJSONLNodeSinkContract(t) + }, + Record: normalizedNode{ID: "1", Kinds: []string{"User"}}, + }) + }) + + t.Run("edges", func(t *testing.T) { + runLeafSinkContract(t, leafSinkContract[normalizedEdge, jsonlFragmentMetadata]{ + New: func(t *testing.T) (fragmentSink[normalizedEdge, jsonlFragmentMetadata], shardID, string) { + options := DumpOptions{ + OutputDir: t.TempDir(), + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + } + id := shardID{Graph: "graph/name", Phase: PhaseEdges, Number: 2} + relativePath, err := jsonlFragmentPath(id.Graph, id.Phase, id.Number, options.Compression) + if err != nil { + t.Fatalf("fragment path: %v", err) + } + return newJSONLEdgeSink(options), id, filepath.Join(options.OutputDir, filepath.FromSlash(relativePath)) + }, + Record: normalizedEdge{StartID: "1", EndID: "2", Kind: "MemberOf"}, + }) + }) +} + +func TestParquetFragmentSinkConformance(t *testing.T) { + t.Run("nodes", func(t *testing.T) { + runLeafSinkContract(t, leafSinkContract[normalizedNode, parquetFragmentMetadata]{ + New: func(t *testing.T) (fragmentSink[normalizedNode, parquetFragmentMetadata], shardID, string) { + outputDir := t.TempDir() + id := shardID{Graph: "graph/name", Phase: PhaseNodes, Number: 2} + relativePath, err := parquetFragmentPath(id.Graph, id.Phase, id.Number) + if err != nil { + t.Fatalf("fragment path: %v", err) + } + workspace := newLocalCollectionWorkspace(outputDir, false) + return newParquetNodeSinkInWorkspace(workspace), id, filepath.Join(outputDir, filepath.FromSlash(relativePath)) + }, + Record: normalizedNode{ID: "1", Kinds: []string{"User"}}, + }) + }) + + t.Run("edges", func(t *testing.T) { + runLeafSinkContract(t, leafSinkContract[normalizedEdge, parquetFragmentMetadata]{ + New: func(t *testing.T) (fragmentSink[normalizedEdge, parquetFragmentMetadata], shardID, string) { + outputDir := t.TempDir() + id := shardID{Graph: "graph/name", Phase: PhaseEdges, Number: 2} + relativePath, err := parquetFragmentPath(id.Graph, id.Phase, id.Number) + if err != nil { + t.Fatalf("fragment path: %v", err) + } + workspace := newLocalCollectionWorkspace(outputDir, false) + return newParquetEdgeSinkInWorkspace(workspace), id, filepath.Join(outputDir, filepath.FromSlash(relativePath)) + }, + Record: normalizedEdge{ID: "3", StartID: "1", EndID: "2", Kind: "MemberOf"}, + }) + }) +} + +func newJSONLNodeSinkContract(t *testing.T) (fragmentSink[normalizedNode, jsonlFragmentMetadata], shardID, string) { + t.Helper() + options := DumpOptions{ + OutputDir: t.TempDir(), + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + } + id := shardID{Graph: "graph/name", Phase: PhaseNodes, Number: 2} + relativePath, err := jsonlFragmentPath(id.Graph, id.Phase, id.Number, options.Compression) + if err != nil { + t.Fatalf("fragment path: %v", err) + } + return newJSONLNodeSink(options), id, filepath.Join(options.OutputDir, filepath.FromSlash(relativePath)) +} + +func runLeafSinkContract[T any, M fragmentMetadata](t *testing.T, contract leafSinkContract[T, M]) { + t.Helper() + + t.Run("prepare and commit transfer ownership", func(t *testing.T) { + sink, id, finalPath := contract.New(t) + writer := openContractWriter(t, sink, id) + if err := writer.WriteBatch(context.Background(), []T{contract.Record}); err != nil { + t.Fatalf("write batch: %v", err) + } + prepared, err := writer.Prepare(context.Background()) + if err != nil { + t.Fatalf("prepare: %v", err) + } + if prepared.Metadata().rowCount() != 1 { + t.Fatalf("prepared rows = %d", prepared.Metadata().rowCount()) + } + assertPathAbsent(t, finalPath) + assertPathPresent(t, finalPath+".tmp") + + assertLeafSinkError(t, writer.WriteBatch(context.Background(), []T{contract.Record}), nil) + _, err = writer.Prepare(context.Background()) + assertLeafSinkError(t, err, nil) + assertLeafSinkError(t, writer.Abort(), nil) + + if err := prepared.Commit(context.Background()); err != nil { + t.Fatalf("commit: %v", err) + } + assertPathPresent(t, finalPath) + assertPathAbsent(t, finalPath+".tmp") + assertLeafSinkError(t, prepared.Commit(context.Background()), nil) + assertLeafSinkError(t, prepared.Abort(), nil) + }) + + t.Run("writer abort is repeatable", func(t *testing.T) { + sink, id, finalPath := contract.New(t) + writer := openContractWriter(t, sink, id) + if err := writer.WriteBatch(context.Background(), []T{contract.Record}); err != nil { + t.Fatalf("write batch: %v", err) + } + if err := writer.Abort(); err != nil { + t.Fatalf("abort writer: %v", err) + } + if err := writer.Abort(); err != nil { + t.Fatalf("repeat writer abort: %v", err) + } + assertPathAbsent(t, finalPath) + assertPathAbsent(t, finalPath+".tmp") + assertLeafSinkError(t, writer.WriteBatch(context.Background(), []T{contract.Record}), nil) + _, err := writer.Prepare(context.Background()) + assertLeafSinkError(t, err, nil) + }) + + t.Run("prepared abort is repeatable", func(t *testing.T) { + sink, id, finalPath := contract.New(t) + writer := openContractWriter(t, sink, id) + if err := writer.WriteBatch(context.Background(), []T{contract.Record}); err != nil { + t.Fatalf("write batch: %v", err) + } + prepared, err := writer.Prepare(context.Background()) + if err != nil { + t.Fatalf("prepare: %v", err) + } + if err := prepared.Abort(); err != nil { + t.Fatalf("abort prepared fragment: %v", err) + } + if err := prepared.Abort(); err != nil { + t.Fatalf("repeat prepared abort: %v", err) + } + assertPathAbsent(t, finalPath) + assertPathAbsent(t, finalPath+".tmp") + assertLeafSinkError(t, prepared.Commit(context.Background()), nil) + }) + + t.Run("context cancellation", func(t *testing.T) { + t.Run("open", func(t *testing.T) { + sink, id, _ := contract.New(t) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := sink.Open(ctx, id) + assertLeafSinkError(t, err, context.Canceled) + }) + + t.Run("write", func(t *testing.T) { + sink, id, finalPath := contract.New(t) + writer := openContractWriter(t, sink, id) + ctx, cancel := context.WithCancel(context.Background()) + cancel() + assertLeafSinkError(t, writer.WriteBatch(ctx, []T{contract.Record}), context.Canceled) + if err := writer.Abort(); err != nil { + t.Fatalf("abort canceled writer: %v", err) + } + assertPathAbsent(t, finalPath+".tmp") + }) + + t.Run("prepare", func(t *testing.T) { + sink, id, finalPath := contract.New(t) + writer := openContractWriter(t, sink, id) + if err := writer.WriteBatch(context.Background(), []T{contract.Record}); err != nil { + t.Fatalf("write batch: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := writer.Prepare(ctx) + assertLeafSinkError(t, err, context.Canceled) + if err := writer.Abort(); err != nil { + t.Fatalf("abort canceled writer: %v", err) + } + assertPathAbsent(t, finalPath+".tmp") + }) + + t.Run("commit", func(t *testing.T) { + sink, id, finalPath := contract.New(t) + writer := openContractWriter(t, sink, id) + if err := writer.WriteBatch(context.Background(), []T{contract.Record}); err != nil { + t.Fatalf("write batch: %v", err) + } + prepared, err := writer.Prepare(context.Background()) + if err != nil { + t.Fatalf("prepare: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + assertLeafSinkError(t, prepared.Commit(ctx), context.Canceled) + if err := prepared.Abort(); err != nil { + t.Fatalf("abort canceled commit: %v", err) + } + assertPathAbsent(t, finalPath) + assertPathAbsent(t, finalPath+".tmp") + }) + }) +} + +func openContractWriter[T any, M fragmentMetadata](t *testing.T, sink fragmentSink[T, M], id shardID) fragmentWriter[T, M] { + t.Helper() + writer, err := sink.Open(context.Background(), id) + if err != nil { + t.Fatalf("open sink: %v", err) + } + return writer +} + +func assertLeafSinkError(t *testing.T, err error, cause error) { + t.Helper() + if err == nil { + t.Fatalf("expected leaf sink error") + } + if cause != nil && !errors.Is(err, cause) { + t.Fatalf("leaf sink error = %v, want cause %v", err, cause) + } +} + +func assertPathPresent(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path); err != nil { + t.Fatalf("path %q is absent: %v", path, err) + } +} + +func assertPathAbsent(t *testing.T, path string) { + t.Helper() + if _, err := os.Stat(path); !os.IsNotExist(err) { + t.Fatalf("path %q exists or cannot be inspected: %v", path, err) + } +} + +type fragmentFailureWorkspace struct { + collectionWorkspace + stageErr error + artifact stagedWorkspaceFile +} + +func (s fragmentFailureWorkspace) Stage(context.Context, string) (stagedWorkspaceFile, error) { + if s.stageErr != nil { + return nil, s.stageErr + } + return s.artifact, nil +} + +type fragmentFailureArtifact struct { + bytes.Buffer + closeErr error + commitErr error + abortErr error + aborted bool +} + +func (s *fragmentFailureArtifact) Close() error { + return s.closeErr +} + +func (s *fragmentFailureArtifact) Commit(context.Context) error { + return s.commitErr +} + +func (s *fragmentFailureArtifact) Abort() error { + s.aborted = true + return s.abortErr +} + +type failingJSONValue struct { + err error +} + +func (s failingJSONValue) MarshalJSON() ([]byte, error) { + return nil, s.err +} + +func TestJSONLFragmentSinkFailureConformance(t *testing.T) { + id := shardID{Graph: "graph/name", Phase: PhaseNodes, Number: 7} + options := DumpOptions{Compression: CompressionGzip, ZstdLevel: DefaultZstdLevel} + newSink := func(workspace collectionWorkspace, adapt func(int) any) jsonlFragmentSink[int] { + return newJSONLFragmentSink(options, workspace, PhaseNodes, adapt) + } + + t.Run("open", func(t *testing.T) { + cause := errors.New("stage failed") + sink := newSink(fragmentFailureWorkspace{stageErr: cause}, func(value int) any { return value }) + _, err := sink.Open(context.Background(), id) + assertLeafSinkError(t, err, cause) + }) + + t.Run("write", func(t *testing.T) { + cause := errors.New("encode failed") + artifact := &fragmentFailureArtifact{} + sink := newSink(fragmentFailureWorkspace{artifact: artifact}, func(int) any { return failingJSONValue{err: cause} }) + writer := openContractWriter(t, sink, id) + assertLeafSinkError(t, writer.WriteBatch(context.Background(), []int{1}), cause) + if err := writer.Abort(); err != nil { + t.Fatalf("abort failed write: %v", err) + } + }) + + t.Run("prepare", func(t *testing.T) { + cause := errors.New("close failed") + cleanupCause := errors.New("prepare cleanup failed") + artifact := &fragmentFailureArtifact{closeErr: cause, abortErr: cleanupCause} + sink := newSink(fragmentFailureWorkspace{artifact: artifact}, func(value int) any { return value }) + writer := openContractWriter(t, sink, id) + if err := writer.WriteBatch(context.Background(), []int{1}); err != nil { + t.Fatalf("write batch: %v", err) + } + _, err := writer.Prepare(context.Background()) + assertLeafSinkError(t, err, cause) + if !errors.Is(err, cleanupCause) { + t.Fatalf("prepare error does not retain cleanup failure: %v", err) + } + if !artifact.aborted { + t.Fatalf("prepare failure did not abort staged artifact") + } + if err := writer.Abort(); err != nil { + t.Fatalf("repeat cleanup after prepare failure: %v", err) + } + }) + + t.Run("commit", func(t *testing.T) { + cause := errors.New("commit failed") + cleanupCause := errors.New("commit cleanup failed") + artifact := &fragmentFailureArtifact{commitErr: cause, abortErr: cleanupCause} + sink := newSink(fragmentFailureWorkspace{artifact: artifact}, func(value int) any { return value }) + writer := openContractWriter(t, sink, id) + if err := writer.WriteBatch(context.Background(), []int{1}); err != nil { + t.Fatalf("write batch: %v", err) + } + prepared, err := writer.Prepare(context.Background()) + if err != nil { + t.Fatalf("prepare: %v", err) + } + err = prepared.Commit(context.Background()) + assertLeafSinkError(t, err, cause) + if !errors.Is(err, cleanupCause) { + t.Fatalf("commit error does not retain cleanup failure: %v", err) + } + if !artifact.aborted { + t.Fatalf("commit failure did not abort staged artifact") + } + if err := prepared.Abort(); err != nil { + t.Fatalf("repeat cleanup after commit failure: %v", err) + } + }) + + t.Run("abort", func(t *testing.T) { + cause := errors.New("abort failed") + artifact := &fragmentFailureArtifact{abortErr: cause} + sink := newSink(fragmentFailureWorkspace{artifact: artifact}, func(value int) any { return value }) + writer := openContractWriter(t, sink, id) + assertLeafSinkError(t, writer.Abort(), cause) + }) +} diff --git a/retriever/sink_jsonl.go b/retriever/sink_jsonl.go new file mode 100644 index 00000000..3dc354d9 --- /dev/null +++ b/retriever/sink_jsonl.go @@ -0,0 +1,174 @@ +package retriever + +import ( + "context" + "fmt" + "path" +) + +const jsonlFragmentFormat = "JSONL" + +type jsonlFragmentSink[T any] struct { + workspace collectionWorkspace + codec CompressionCodec + zstdLevel int + phase Phase + adapt func(T) any +} + +type jsonlFragmentMetadata struct { + Path string + Rows int + CompressedBytes int64 + UncompressedBytes int64 + SHA256 string +} + +func (s jsonlFragmentMetadata) rowCount() int { + return s.Rows +} + +func newJSONLNodeSink(options DumpOptions) jsonlFragmentSink[normalizedNode] { + return newJSONLNodeSinkInWorkspace(options, newLocalCollectionWorkspace(options.OutputDir, options.Force)) +} + +func newJSONLNodeSinkInWorkspace(options DumpOptions, workspace collectionWorkspace) jsonlFragmentSink[normalizedNode] { + return newJSONLFragmentSink(options, workspace, PhaseNodes, func(record normalizedNode) any { + return jsonlV1NodeFromNormalized(record) + }) +} + +func newJSONLEdgeSink(options DumpOptions) jsonlFragmentSink[normalizedEdge] { + return newJSONLEdgeSinkInWorkspace(options, newLocalCollectionWorkspace(options.OutputDir, options.Force)) +} + +func newJSONLEdgeSinkInWorkspace(options DumpOptions, workspace collectionWorkspace) jsonlFragmentSink[normalizedEdge] { + return newJSONLFragmentSink(options, workspace, PhaseEdges, func(record normalizedEdge) any { + return jsonlV1EdgeFromNormalized(record) + }) +} + +func newJSONLFragmentSink[T any](options DumpOptions, workspace collectionWorkspace, phase Phase, adapt func(T) any) jsonlFragmentSink[T] { + return jsonlFragmentSink[T]{ + workspace: workspace, + codec: options.Compression, + zstdLevel: options.ZstdLevel, + phase: phase, + adapt: adapt, + } +} + +func (s jsonlFragmentSink[T]) Open(ctx context.Context, id shardID) (fragmentWriter[T, jsonlFragmentMetadata], error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if id.Phase != s.phase { + return nil, fmt.Errorf("JSONL sink for phase %q cannot open shard for phase %q", s.phase, id.Phase) + } + + relativePath, err := jsonlFragmentPath(id.Graph, id.Phase, id.Number, s.codec) + if err != nil { + return nil, err + } + writer, err := newCompressedJSONLinesWriterInWorkspace( + ctx, + s.workspace, + relativePath, + s.codec, + s.zstdLevel, + ) + if err != nil { + return nil, err + } + + return &jsonlFragmentWriter[T]{ + relativePath: relativePath, + adapt: s.adapt, + writer: writer, + }, nil +} + +type jsonlFragmentWriter[T any] struct { + relativePath string + adapt func(T) any + writer *compressedJSONLinesWriter +} + +func (s *jsonlFragmentWriter[T]) WriteBatch(ctx context.Context, records []T) error { + for _, record := range records { + if err := ctx.Err(); err != nil { + return err + } + if err := s.writer.Write(s.adapt(record)); err != nil { + return err + } + } + return nil +} + +func (s *jsonlFragmentWriter[T]) Prepare(ctx context.Context) (preparedFragment[jsonlFragmentMetadata], error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + prepared, err := s.writer.Prepare() + if err != nil { + return nil, err + } + encoding := prepared.Metadata() + metadata := jsonlFragmentMetadata{ + Path: s.relativePath, + Rows: encoding.Rows, + CompressedBytes: encoding.CompressedBytes, + UncompressedBytes: encoding.UncompressedBytes, + SHA256: encoding.SHA256, + } + return &preparedJSONLFragment{ + fragment: prepared, + metadata: metadata, + }, nil +} + +func (s *jsonlFragmentWriter[T]) Abort() error { + return s.writer.Abort() +} + +type preparedJSONLFragment struct { + fragment *preparedCompressedJSONLinesFragment + metadata jsonlFragmentMetadata +} + +func (s *preparedJSONLFragment) Metadata() jsonlFragmentMetadata { + return s.metadata +} + +func (s *preparedJSONLFragment) Commit(ctx context.Context) error { + return s.fragment.Commit(ctx) +} + +func (s *preparedJSONLFragment) Abort() error { + return s.fragment.Abort() +} + +func jsonlFragmentPath(graphName string, fragmentPhase Phase, shardNumber int, codec CompressionCodec) (string, error) { + if shardNumber <= 0 { + return "", fmt.Errorf("shard number must be > 0") + } + + extension, err := compressionExtension(codec) + if err != nil { + return "", err + } + + var prefix string + switch fragmentPhase { + case PhaseNodes: + prefix = "nodes" + case PhaseEdges: + prefix = "edges" + default: + return "", fmt.Errorf("unsupported fragment phase %q", fragmentPhase) + } + + return path.Join("graphs", graphDirectoryName(graphName), fmt.Sprintf("%s-%06d.jsonl%s", prefix, shardNumber, extension)), nil +} diff --git a/retriever/sink_jsonl_test.go b/retriever/sink_jsonl_test.go new file mode 100644 index 00000000..c5cbf254 --- /dev/null +++ b/retriever/sink_jsonl_test.go @@ -0,0 +1,127 @@ +package retriever + +import ( + "context" + "os" + "path/filepath" + "reflect" + "testing" +) + +func TestJSONLFragmentSinkOwnsWireFormatAndPhysicalMetadata(t *testing.T) { + options := DumpOptions{ + OutputDir: t.TempDir(), + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + } + summary := shardSummary{ + ID: shardID{Graph: "graph/name", Phase: PhaseNodes, Number: 2}, + Rows: 1, + ActionCounts: map[string]int{"preserve": 1}, + } + + metadata := writeTestJSONLShard(t, newJSONLNodeSink(options), summary, []normalizedNode{{ + ID: "1", + Kinds: []string{"User"}, + Properties: map[string]any{"name": "alice"}, + }}) + if metadata.Path != "graphs/graph%2Fname/nodes-000002.jsonl.gz" || metadata.Rows != 1 || metadata.CompressedBytes <= 0 || metadata.UncompressedBytes <= 0 || metadata.SHA256 == "" { + t.Fatalf("JSONL metadata = %+v", metadata) + } + + var records []FragmentNode + if _, err := readCompressedJSONLines(filepath.Join(options.OutputDir, filepath.FromSlash(metadata.Path)), CompressionGzip, func(record FragmentNode) error { + records = append(records, record) + return nil + }); err != nil { + t.Fatalf("read JSONL fragment: %v", err) + } + if !reflect.DeepEqual(records, []FragmentNode{{ + ID: "1", + Kinds: []string{"User"}, + Properties: map[string]any{"name": "alice"}, + }}) { + t.Fatalf("JSONL records = %#v", records) + } +} + +func TestJSONLEdgeSinkOmitsSourceRelationshipID(t *testing.T) { + options := DumpOptions{ + OutputDir: t.TempDir(), + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + } + summary := shardSummary{ID: shardID{Graph: "example", Phase: PhaseEdges, Number: 1}, Rows: 1} + metadata := writeTestJSONLShard(t, newJSONLEdgeSink(options), summary, []normalizedEdge{{ + ID: "source-edge-id", + StartID: "1", + EndID: "2", + Kind: "MemberOf", + }}) + + var records []FragmentEdge + if _, err := readCompressedJSONLines(filepath.Join(options.OutputDir, filepath.FromSlash(metadata.Path)), CompressionGzip, func(record FragmentEdge) error { + records = append(records, record) + return nil + }); err != nil { + t.Fatalf("read JSONL edge fragment: %v", err) + } + if !reflect.DeepEqual(records, []FragmentEdge{{StartID: "1", EndID: "2", Kind: "MemberOf"}}) { + t.Fatalf("JSONL edge records = %#v", records) + } +} + +func TestJSONLFragmentSinkRejectsPhaseAndOutputRejectsRowCountMismatch(t *testing.T) { + options := DumpOptions{ + OutputDir: t.TempDir(), + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + } + sink := newJSONLNodeSink(options) + if _, err := sink.Open(context.Background(), shardID{Graph: "example", Phase: PhaseEdges, Number: 1}); err == nil { + t.Fatalf("expected phase mismatch") + } + + summary := shardSummary{ID: shardID{Graph: "example", Phase: PhaseNodes, Number: 1}, Rows: 2} + output := newShardSinkSet(newJSONLShardSink(sink)) + writer, err := output.OpenShard(context.Background(), summary.ID) + if err != nil { + t.Fatalf("open shard: %v", err) + } + if err := writer.WriteBatch(context.Background(), []normalizedNode{{ID: "1"}}); err != nil { + t.Fatalf("write shard: %v", err) + } + if _, err := writer.Finish(context.Background(), summary); err == nil { + t.Fatalf("expected row count mismatch") + } + + relativePath, err := jsonlFragmentPath(summary.ID.Graph, summary.ID.Phase, summary.ID.Number, options.Compression) + if err != nil { + t.Fatalf("fragment path: %v", err) + } + absolutePath := filepath.Join(options.OutputDir, filepath.FromSlash(relativePath)) + if _, err := os.Stat(absolutePath); !os.IsNotExist(err) { + t.Fatalf("final path exists after prepare failure: %v", err) + } + if _, err := os.Stat(absolutePath + ".tmp"); !os.IsNotExist(err) { + t.Fatalf("staged path exists after prepare failure: %v", err) + } +} + +func writeTestJSONLShard[T any](t *testing.T, sink fragmentSink[T, jsonlFragmentMetadata], summary shardSummary, records []T) jsonlFragmentMetadata { + t.Helper() + + output := newShardSinkSet(newJSONLShardSink(sink)) + writer, err := output.OpenShard(context.Background(), summary.ID) + if err != nil { + t.Fatalf("open JSONL shard: %v", err) + } + if err := writer.WriteBatch(context.Background(), records); err != nil { + t.Fatalf("write JSONL shard: %v", err) + } + committed, err := writer.Finish(context.Background(), summary) + if err != nil { + t.Fatalf("finish JSONL shard: %v", err) + } + return committed.JSONL +} diff --git a/retriever/sink_parquet.go b/retriever/sink_parquet.go new file mode 100644 index 00000000..5f631398 --- /dev/null +++ b/retriever/sink_parquet.go @@ -0,0 +1,275 @@ +package retriever + +import ( + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "hash" + "io" + "path" + + parquetgo "github.com/parquet-go/parquet-go" + parquetzstd "github.com/parquet-go/parquet-go/compress/zstd" +) + +const parquetFragmentFormat = "Parquet" + +type parquetNode struct { + ID string `parquet:"id,zstd"` + Kinds []string `parquet:"kinds,zstd"` + Properties string `parquet:"properties,json,zstd"` +} + +type parquetEdge struct { + ID string `parquet:"id,zstd"` + StartID string `parquet:"start_id,zstd"` + EndID string `parquet:"end_id,zstd"` + Kind string `parquet:"kind,zstd"` + Properties string `parquet:"properties,json,zstd"` +} + +type parquetFragmentMetadata struct { + Path string + Rows int + Bytes int64 + SHA256 string +} + +func (s parquetFragmentMetadata) rowCount() int { + return s.Rows +} + +type parquetFragmentSink[T, P any] struct { + workspace collectionWorkspace + phase Phase + adapt func(T) (P, error) +} + +func newParquetNodeSinkInWorkspace(workspace collectionWorkspace) parquetFragmentSink[normalizedNode, parquetNode] { + return parquetFragmentSink[normalizedNode, parquetNode]{ + workspace: workspace, + phase: PhaseNodes, + adapt: parquetNodeFromNormalized, + } +} + +func newParquetEdgeSinkInWorkspace(workspace collectionWorkspace) parquetFragmentSink[normalizedEdge, parquetEdge] { + return parquetFragmentSink[normalizedEdge, parquetEdge]{ + workspace: workspace, + phase: PhaseEdges, + adapt: parquetEdgeFromNormalized, + } +} + +func (s parquetFragmentSink[T, P]) Open(ctx context.Context, id shardID) (fragmentWriter[T, parquetFragmentMetadata], error) { + if err := ctx.Err(); err != nil { + return nil, err + } + if id.Phase != s.phase { + return nil, fmt.Errorf("Parquet sink for phase %q cannot open shard for phase %q", s.phase, id.Phase) + } + + relativePath, err := parquetFragmentPath(id.Graph, id.Phase, id.Number) + if err != nil { + return nil, err + } + artifact, err := s.workspace.Stage(ctx, relativePath) + if err != nil { + return nil, err + } + hasher := sha256.New() + counter := &countingWriter{writer: io.MultiWriter(artifact, hasher)} + writer := parquetgo.NewGenericWriter[P](counter, parquetgo.Compression(&parquetzstd.Codec{})) + return &parquetFragmentWriter[T, P]{ + relativePath: relativePath, + adapt: s.adapt, + artifact: artifact, + writer: writer, + counter: counter, + hasher: hasher, + }, nil +} + +type parquetFragmentWriter[T, P any] struct { + relativePath string + adapt func(T) (P, error) + artifact stagedWorkspaceFile + writer *parquetgo.GenericWriter[P] + counter *countingWriter + hasher hash.Hash + rows int + state parquetWriterState +} + +type parquetWriterState uint8 + +const ( + parquetWriterOpen parquetWriterState = iota + parquetWriterPrepared + parquetWriterAborted +) + +func (s *parquetFragmentWriter[T, P]) WriteBatch(ctx context.Context, records []T) error { + if s.state != parquetWriterOpen { + return fmt.Errorf("write Parquet fragment after prepare or abort") + } + + rows := make([]P, len(records)) + for index, record := range records { + if err := ctx.Err(); err != nil { + return err + } + row, err := s.adapt(record) + if err != nil { + return fmt.Errorf("adapt Parquet row %d: %w", s.rows+index+1, err) + } + rows[index] = row + } + written, err := s.writer.Write(rows) + s.rows += written + if err != nil { + return fmt.Errorf("write Parquet rows: %w", err) + } + if written != len(rows) { + return fmt.Errorf("write Parquet rows: wrote %d of %d", written, len(rows)) + } + return nil +} + +func (s *parquetFragmentWriter[T, P]) Prepare(ctx context.Context) (preparedFragment[parquetFragmentMetadata], error) { + if s.state != parquetWriterOpen { + return nil, fmt.Errorf("prepare Parquet fragment more than once or after abort") + } + if err := ctx.Err(); err != nil { + return nil, err + } + s.state = parquetWriterPrepared + + if err := s.writer.Close(); err != nil { + s.state = parquetWriterAborted + return nil, cleanupOnError(fmt.Errorf("finish Parquet fragment: %w", err), s.artifact.Abort) + } + if err := s.artifact.Close(); err != nil { + s.state = parquetWriterAborted + return nil, cleanupOnError(fmt.Errorf("close Parquet fragment: %w", err), s.artifact.Abort) + } + + return &preparedParquetFragment{ + artifact: s.artifact, + metadata: parquetFragmentMetadata{ + Path: s.relativePath, + Rows: s.rows, + Bytes: s.counter.count, + SHA256: hex.EncodeToString(s.hasher.Sum(nil)), + }, + }, nil +} + +func (s *parquetFragmentWriter[T, P]) Abort() error { + switch s.state { + case parquetWriterOpen: + s.state = parquetWriterAborted + return s.artifact.Abort() + case parquetWriterAborted: + return nil + default: + return fmt.Errorf("abort Parquet writer after prepare") + } +} + +type preparedParquetFragment struct { + artifact stagedWorkspaceFile + metadata parquetFragmentMetadata + state preparedParquetState +} + +type preparedParquetState uint8 + +const ( + parquetFragmentPrepared preparedParquetState = iota + parquetFragmentCommitted + parquetFragmentAborted +) + +func (s *preparedParquetFragment) Metadata() parquetFragmentMetadata { + return s.metadata +} + +func (s *preparedParquetFragment) Commit(ctx context.Context) error { + if s.state != parquetFragmentPrepared { + return fmt.Errorf("commit Parquet fragment that is not prepared") + } + if err := ctx.Err(); err != nil { + return err + } + if err := s.artifact.Commit(ctx); err != nil { + s.state = parquetFragmentAborted + return cleanupOnError(err, s.artifact.Abort) + } + s.state = parquetFragmentCommitted + return nil +} + +func (s *preparedParquetFragment) Abort() error { + switch s.state { + case parquetFragmentPrepared: + s.state = parquetFragmentAborted + return s.artifact.Abort() + case parquetFragmentAborted: + return nil + default: + return fmt.Errorf("abort committed Parquet fragment") + } +} + +func parquetNodeFromNormalized(node normalizedNode) (parquetNode, error) { + properties, err := parquetProperties(node.Properties) + if err != nil { + return parquetNode{}, err + } + return parquetNode{ID: node.ID, Kinds: node.Kinds, Properties: properties}, nil +} + +func parquetEdgeFromNormalized(edge normalizedEdge) (parquetEdge, error) { + properties, err := parquetProperties(edge.Properties) + if err != nil { + return parquetEdge{}, err + } + return parquetEdge{ + ID: edge.ID, + StartID: edge.StartID, + EndID: edge.EndID, + Kind: edge.Kind, + Properties: properties, + }, nil +} + +func parquetProperties(properties map[string]any) (string, error) { + if len(properties) == 0 { + return "null", nil + } + payload, err := json.Marshal(properties) + if err != nil { + return "", fmt.Errorf("encode properties as JSON: %w", err) + } + return string(payload), nil +} + +func parquetFragmentPath(graphName string, phase Phase, shardNumber int) (string, error) { + if shardNumber <= 0 { + return "", fmt.Errorf("shard number must be > 0") + } + + var prefix string + switch phase { + case PhaseNodes: + prefix = "nodes" + case PhaseEdges: + prefix = "edges" + default: + return "", fmt.Errorf("unsupported fragment phase %q", phase) + } + return path.Join("parquet", "graphs", graphDirectoryName(graphName), fmt.Sprintf("%s-%06d.parquet", prefix, shardNumber)), nil +} diff --git a/retriever/sink_parquet_test.go b/retriever/sink_parquet_test.go new file mode 100644 index 00000000..7f789988 --- /dev/null +++ b/retriever/sink_parquet_test.go @@ -0,0 +1,157 @@ +package retriever + +import ( + "context" + "encoding/json" + "path/filepath" + "reflect" + "testing" + + parquetgo "github.com/parquet-go/parquet-go" +) + +func TestJSONLAndParquetNodeFragmentsAreLogicallyEquivalent(t *testing.T) { + outputDir := t.TempDir() + workspace := newLocalCollectionWorkspace(outputDir, false) + options := DumpOptions{Compression: CompressionGzip, ZstdLevel: DefaultZstdLevel} + output := newShardSinkSet( + newJSONLShardSink(newJSONLNodeSinkInWorkspace(options, workspace)), + newParquetShardSink[normalizedNode](newParquetNodeSinkInWorkspace(workspace)), + ) + id := shardID{Graph: "graph/name", Phase: PhaseNodes, Number: 1} + records := []normalizedNode{ + {ID: "1", Kinds: []string{"Computer", "User"}, Properties: map[string]any{"active": true, "name": "alice", "scores": []any{1.0, 2.0}}}, + {ID: "2", Kinds: []string{"Group"}}, + } + committed := writeTestShard(t, output, id, records) + if committed.Parquet == nil { + t.Fatalf("missing committed Parquet metadata") + } + assertParquetMetadata(t, outputDir, *committed.Parquet) + + var jsonlRows []FragmentNode + jsonlPath := filepath.Join(outputDir, filepath.FromSlash(committed.JSONL.Path)) + if _, err := readCompressedJSONLines[FragmentNode](jsonlPath, options.Compression, func(row FragmentNode) error { + jsonlRows = append(jsonlRows, row) + return nil + }); err != nil { + t.Fatalf("read JSONL nodes: %v", err) + } + parquetRows, err := parquetgo.ReadFile[parquetNode](filepath.Join(outputDir, filepath.FromSlash(committed.Parquet.Path))) + if err != nil { + t.Fatalf("read Parquet nodes: %v", err) + } + if len(jsonlRows) != len(parquetRows) { + t.Fatalf("JSONL rows = %d, Parquet rows = %d", len(jsonlRows), len(parquetRows)) + } + for index, parquetRow := range parquetRows { + jsonlRow := jsonlRows[index] + if parquetRow.ID != jsonlRow.ID || !reflect.DeepEqual(parquetRow.Kinds, jsonlRow.Kinds) || !reflect.DeepEqual(decodeParquetProperties(t, parquetRow.Properties), jsonlRow.Properties) { + t.Fatalf("node row %d differs: JSONL=%+v Parquet=%+v", index, jsonlRow, parquetRow) + } + } +} + +func TestJSONLAndParquetEdgeFragmentsAreLogicallyEquivalentAndParquetRetainsID(t *testing.T) { + outputDir := t.TempDir() + workspace := newLocalCollectionWorkspace(outputDir, false) + options := DumpOptions{Compression: CompressionGzip, ZstdLevel: DefaultZstdLevel} + output := newShardSinkSet( + newJSONLShardSink(newJSONLEdgeSinkInWorkspace(options, workspace)), + newParquetShardSink[normalizedEdge](newParquetEdgeSinkInWorkspace(workspace)), + ) + id := shardID{Graph: "graph/name", Phase: PhaseEdges, Number: 1} + records := []normalizedEdge{ + {ID: "relationship-7", StartID: "1", EndID: "2", Kind: "MemberOf", Properties: map[string]any{"weight": 3.0}}, + {ID: "relationship-8", StartID: "2", EndID: "3", Kind: "Owns"}, + } + committed := writeTestShard(t, output, id, records) + if committed.Parquet == nil { + t.Fatalf("missing committed Parquet metadata") + } + assertParquetMetadata(t, outputDir, *committed.Parquet) + + var jsonlRows []FragmentEdge + jsonlPath := filepath.Join(outputDir, filepath.FromSlash(committed.JSONL.Path)) + if _, err := readCompressedJSONLines[FragmentEdge](jsonlPath, options.Compression, func(row FragmentEdge) error { + jsonlRows = append(jsonlRows, row) + return nil + }); err != nil { + t.Fatalf("read JSONL edges: %v", err) + } + parquetRows, err := parquetgo.ReadFile[parquetEdge](filepath.Join(outputDir, filepath.FromSlash(committed.Parquet.Path))) + if err != nil { + t.Fatalf("read Parquet edges: %v", err) + } + if len(jsonlRows) != len(parquetRows) { + t.Fatalf("JSONL rows = %d, Parquet rows = %d", len(jsonlRows), len(parquetRows)) + } + for index, parquetRow := range parquetRows { + jsonlRow := jsonlRows[index] + if parquetRow.ID != records[index].ID { + t.Fatalf("Parquet edge ID = %q, want %q", parquetRow.ID, records[index].ID) + } + if parquetRow.StartID != jsonlRow.StartID || parquetRow.EndID != jsonlRow.EndID || parquetRow.Kind != jsonlRow.Kind || !reflect.DeepEqual(decodeParquetProperties(t, parquetRow.Properties), jsonlRow.Properties) { + t.Fatalf("edge row %d differs: JSONL=%+v Parquet=%+v", index, jsonlRow, parquetRow) + } + } +} + +func TestParquetPropertyConversionFailureLeavesNoArtifact(t *testing.T) { + outputDir := t.TempDir() + workspace := newLocalCollectionWorkspace(outputDir, false) + sink := newParquetNodeSinkInWorkspace(workspace) + id := shardID{Graph: "example", Phase: PhaseNodes, Number: 1} + writer, err := sink.Open(context.Background(), id) + if err != nil { + t.Fatalf("open sink: %v", err) + } + err = writer.WriteBatch(context.Background(), []normalizedNode{{ID: "1", Properties: map[string]any{"unsupported": func() {}}}}) + if err == nil { + t.Fatalf("expected property conversion failure") + } + if err := writer.Abort(); err != nil { + t.Fatalf("abort failed writer: %v", err) + } + relativePath, err := parquetFragmentPath(id.Graph, id.Phase, id.Number) + if err != nil { + t.Fatalf("fragment path: %v", err) + } + assertPathAbsent(t, filepath.Join(outputDir, filepath.FromSlash(relativePath))) + assertPathAbsent(t, filepath.Join(outputDir, filepath.FromSlash(relativePath))+".tmp") +} + +func writeTestShard[T any](t *testing.T, output shardOutput[T], id shardID, records []T) committedShard { + t.Helper() + writer, err := output.OpenShard(context.Background(), id) + if err != nil { + t.Fatalf("open shard: %v", err) + } + if err := writer.WriteBatch(context.Background(), records); err != nil { + t.Fatalf("write batch: %v", err) + } + committed, err := writer.Finish(context.Background(), shardSummary{ID: id, Rows: len(records)}) + if err != nil { + t.Fatalf("finish shard: %v", err) + } + return committed +} + +func decodeParquetProperties(t *testing.T, properties string) map[string]any { + t.Helper() + var value map[string]any + if err := json.Unmarshal([]byte(properties), &value); err != nil { + t.Fatalf("decode Parquet properties: %v", err) + } + return value +} + +func assertParquetMetadata(t *testing.T, outputDir string, metadata parquetFragmentMetadata) { + t.Helper() + if metadata.Rows <= 0 || metadata.Bytes <= 0 || metadata.SHA256 == "" { + t.Fatalf("invalid Parquet metadata: %+v", metadata) + } + if err := verifyChecksum(filepath.Join(outputDir, filepath.FromSlash(metadata.Path)), metadata.SHA256, metadata.Bytes); err != nil { + t.Fatalf("verify Parquet metadata: %v", err) + } +} diff --git a/retriever/source.go b/retriever/source.go new file mode 100644 index 00000000..2b7dae32 --- /dev/null +++ b/retriever/source.go @@ -0,0 +1,167 @@ +package retriever + +import ( + "context" + "fmt" + + "github.com/specterops/dawgs/graph" + "github.com/specterops/dawgs/query" +) + +type graphEntitySnapshot struct { + NodeCount int64 + EdgeCount int64 +} + +type faucet[T any] interface { + Run(context.Context, entityBatchHandler[T]) (int64, error) +} + +type graphSource interface { + Inventory(context.Context, graph.Graph) (graphEntitySnapshot, error) + Nodes(graph.Graph, int64, int) faucet[*graph.Node] + Edges(graph.Graph, int64, int) faucet[*graph.Relationship] +} + +type databaseGraphSource struct { + db graph.Database +} + +func newDatabaseGraphSource(db graph.Database) databaseGraphSource { + return databaseGraphSource{db: db} +} + +func (s databaseGraphSource) Inventory(ctx context.Context, targetGraph graph.Graph) (graphEntitySnapshot, error) { + var snapshot graphEntitySnapshot + if err := s.db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + + var err error + if snapshot.NodeCount, err = tx.Nodes().Count(); err != nil { + return fmt.Errorf("count nodes: %w", err) + } + + if snapshot.EdgeCount, err = tx.Relationships().Count(); err != nil { + return fmt.Errorf("count relationships: %w", err) + } + + return nil + }); err != nil { + return graphEntitySnapshot{}, err + } + + return snapshot, nil +} + +func (s databaseGraphSource) Nodes(targetGraph graph.Graph, total int64, batchSize int) faucet[*graph.Node] { + return keysetFaucet[*graph.Node]{ + total: total, + batchSize: batchSize, + entityName: "node", + read: func(ctx context.Context, afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Node, error) { + return s.readNodes(ctx, targetGraph, afterID, hasAfterID, limit) + }, + id: func(node *graph.Node) graph.ID { + return node.ID + }, + } +} + +func (s databaseGraphSource) Edges(targetGraph graph.Graph, total int64, batchSize int) faucet[*graph.Relationship] { + return keysetFaucet[*graph.Relationship]{ + total: total, + batchSize: batchSize, + entityName: "relationship", + read: func(ctx context.Context, afterID graph.ID, hasAfterID bool, limit int) ([]*graph.Relationship, error) { + return s.readRelationships(ctx, targetGraph, afterID, hasAfterID, limit) + }, + id: func(relationship *graph.Relationship) graph.ID { + return relationship.ID + }, + } +} + +type keysetFaucet[T any] struct { + total int64 + batchSize int + entityName string + read func(context.Context, graph.ID, bool, int) ([]T, error) + id entityIDFunc[T] +} + +func (s keysetFaucet[T]) Run(ctx context.Context, handle entityBatchHandler[T]) (int64, error) { + return scanEntityBatches(entityScanOptions[T]{ + Total: s.total, + BatchSize: s.batchSize, + EntityName: s.entityName, + Read: func(afterID graph.ID, hasAfterID bool, limit int) ([]T, error) { + return s.read(ctx, afterID, hasAfterID, limit) + }, + ID: s.id, + Handle: handle, + }) +} + +func (s databaseGraphSource) readNodes(ctx context.Context, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Node, error) { + var nodes []*graph.Node + if err := s.db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + nodeQuery := tx.Nodes(). + OrderBy(query.NodeID()). + Limit(batchSize) + + if hasAfterID { + nodeQuery = nodeQuery.Filter(query.GreaterThan(query.NodeID(), afterID)) + } + + return nodeQuery.Fetch(func(cursor graph.Cursor[*graph.Node]) error { + for node := range cursor.Chan() { + nodes = append(nodes, node) + } + + return cursor.Error() + }) + }); err != nil { + if hasAfterID { + return nil, fmt.Errorf("read node batch after ID %d: %w", afterID.Uint64(), err) + } + + return nil, fmt.Errorf("read initial node batch: %w", err) + } + + return nodes, nil +} + +func (s databaseGraphSource) readRelationships(ctx context.Context, targetGraph graph.Graph, afterID graph.ID, hasAfterID bool, batchSize int) ([]*graph.Relationship, error) { + var relationships []*graph.Relationship + if err := s.db.ReadTransaction(ctx, func(tx graph.Transaction) error { + tx = tx.WithGraph(targetGraph) + relationshipQuery := tx.Relationships(). + OrderBy(query.RelationshipID()). + Limit(batchSize) + + if hasAfterID { + relationshipQuery = relationshipQuery.Filter(query.GreaterThan(query.RelationshipID(), afterID)) + } + + return relationshipQuery.Fetch(func(cursor graph.Cursor[*graph.Relationship]) error { + for relationship := range cursor.Chan() { + relationships = append(relationships, relationship) + } + + return cursor.Error() + }) + }); err != nil { + if hasAfterID { + return nil, fmt.Errorf("read relationship batch after ID %d: %w", afterID.Uint64(), err) + } + + return nil, fmt.Errorf("read initial relationship batch: %w", err) + } + + return relationships, nil +} + +func countGraphEntitySnapshot(ctx context.Context, db graph.Database, targetGraph graph.Graph) (graphEntitySnapshot, error) { + return newDatabaseGraphSource(db).Inventory(ctx, targetGraph) +} diff --git a/retriever/source_test.go b/retriever/source_test.go new file mode 100644 index 00000000..4755981b --- /dev/null +++ b/retriever/source_test.go @@ -0,0 +1,194 @@ +package retriever + +import ( + "context" + "os" + "path/filepath" + "reflect" + "testing" + + parquetgo "github.com/parquet-go/parquet-go" + "github.com/specterops/dawgs/graph" +) + +func TestDatabaseGraphSourceCreatesFreshCappedFaucets(t *testing.T) { + database := newScriptedDumpDatabase(map[string]*scriptedDumpGraph{ + "source": { + nodeCount: 2, + nodes: []*graph.Node{ + graph.NewNode(3, nil), + graph.NewNode(1, nil), + graph.NewNode(2, nil), + }, + ignoreNodeLimit: true, + }, + }) + source := newDatabaseGraphSource(database) + target := graph.Graph{Name: "source"} + + snapshot, err := source.Inventory(context.Background(), target) + if err != nil { + t.Fatalf("inventory: %v", err) + } + if snapshot.NodeCount != 2 || snapshot.EdgeCount != 0 { + t.Fatalf("inventory = %+v", snapshot) + } + + for run := range 2 { + var ids []graph.ID + processed, err := source.Nodes(target, snapshot.NodeCount, 10).Run(context.Background(), func(nodes []*graph.Node) error { + for _, node := range nodes { + ids = append(ids, node.ID) + } + return nil + }) + if err != nil { + t.Fatalf("run %d: %v", run+1, err) + } + if processed != 2 || !reflect.DeepEqual(ids, []graph.ID{1, 2}) { + t.Fatalf("run %d processed=%d ids=%v", run+1, processed, ids) + } + } +} + +func TestDumpWithParquetSidecar(t *testing.T) { + outputDir := t.TempDir() + source := &scriptedGraphSource{ + snapshot: graphEntitySnapshot{NodeCount: 2, EdgeCount: 1}, + nodeBatches: [][]*graph.Node{{ + graph.NewNode(1, graph.AsProperties(map[string]any{"name": "alice"}), graph.StringKind("User")), + graph.NewNode(2, nil, graph.StringKind("Group")), + }}, + edgeBatches: [][]*graph.Relationship{{ + graph.NewRelationship(10, 1, 2, graph.AsProperties(map[string]any{"since": 2024}), graph.StringKind("MemberOf")), + }}, + } + options := DefaultDumpOptions(outputDir) + options.Compression = CompressionGzip + options.Parquet = true + options.ShardSize = 10 + options.BatchSize = 10 + + result, err := runDump(context.Background(), source, "scripted-source", []GraphTarget{{Name: "source"}}, options, dumpOverrides{}) + if err != nil { + t.Fatalf("dump: %v", err) + } + if result.ParquetManifestPath == "" || result.ParquetSuccessPath == "" { + t.Fatalf("Parquet publication paths = manifest %q success %q", result.ParquetManifestPath, result.ParquetSuccessPath) + } + if _, err := os.Stat(result.ParquetSuccessPath); err != nil { + t.Fatalf("stat Parquet success marker: %v", err) + } + manifest, err := readParquetManifest(outputDir) + if err != nil { + t.Fatalf("read Parquet manifest: %v", err) + } + if len(manifest.Graphs) != 1 || manifest.Graphs[0].NodeCount != 2 || manifest.Graphs[0].EdgeCount != 1 || len(manifest.Graphs[0].Files) != 2 { + t.Fatalf("Parquet manifest = %+v", manifest) + } + for _, file := range manifest.Graphs[0].Files { + if err := verifyChecksum(filepath.Join(outputDir, filepath.FromSlash(file.Path)), file.SHA256, file.Bytes); err != nil { + t.Fatalf("verify Parquet file %q: %v", file.Path, err) + } + } + edgeFile := manifest.Graphs[0].Files[1] + edges, err := parquetgo.ReadFile[parquetEdge](filepath.Join(outputDir, filepath.FromSlash(edgeFile.Path))) + if err != nil { + t.Fatalf("read Parquet edge file: %v", err) + } + if len(edges) != 1 || edges[0].ID != "10" || edges[0].StartID != "1" || edges[0].EndID != "2" { + t.Fatalf("Parquet edges = %+v", edges) + } + if len(result.Manifest.Graphs) != 1 || len(result.Manifest.Graphs[0].Files) != 2 { + t.Fatalf("JSONL manifest changed by Parquet sidecar: %+v", result.Manifest.Graphs) + } +} + +func TestDumpWithScriptedGraphSource(t *testing.T) { + source := &scriptedGraphSource{ + snapshot: graphEntitySnapshot{NodeCount: 2, EdgeCount: 1}, + nodeBatches: [][]*graph.Node{{ + graph.NewNode(1, graph.AsProperties(map[string]any{"name": "alice"}), graph.StringKind("User")), + graph.NewNode(2, nil, graph.StringKind("Group")), + }}, + edgeBatches: [][]*graph.Relationship{{ + graph.NewRelationship(10, 1, 2, nil, graph.StringKind("MemberOf")), + }}, + } + + result, err := runDump(context.Background(), source, "scripted-source", []GraphTarget{{Name: "source"}}, DumpOptions{ + OutputDir: t.TempDir(), + Scrub: ScrubFull, + Salt: "source-test", + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + ShardSize: 10, + BatchSize: 10, + }, dumpOverrides{}) + if err != nil { + t.Fatalf("dump: %v", err) + } + + if result.NodeCount != 2 || result.EdgeCount != 1 { + t.Fatalf("dump counts: nodes=%d edges=%d", result.NodeCount, result.EdgeCount) + } + if source.inventoryCalls != 1 || source.nodeFaucets != 2 || source.edgeFaucets != 1 { + t.Fatalf("source calls: inventory=%d node_faucets=%d edge_faucets=%d", source.inventoryCalls, source.nodeFaucets, source.edgeFaucets) + } + if len(result.Manifest.Graphs) != 1 || len(result.Manifest.Graphs[0].Files) != 2 { + t.Fatalf("manifest graph = %+v", result.Manifest.Graphs) + } +} + +type scriptedGraphSource struct { + snapshot graphEntitySnapshot + nodeBatches [][]*graph.Node + edgeBatches [][]*graph.Relationship + inventoryCalls int + nodeFaucets int + edgeFaucets int +} + +func (s *scriptedGraphSource) Inventory(context.Context, graph.Graph) (graphEntitySnapshot, error) { + s.inventoryCalls++ + return s.snapshot, nil +} + +func (s *scriptedGraphSource) Nodes(graph.Graph, int64, int) faucet[*graph.Node] { + s.nodeFaucets++ + return scriptedFaucet[*graph.Node]{batches: s.nodeBatches, total: s.snapshot.NodeCount} +} + +func (s *scriptedGraphSource) Edges(graph.Graph, int64, int) faucet[*graph.Relationship] { + s.edgeFaucets++ + return scriptedFaucet[*graph.Relationship]{batches: s.edgeBatches, total: s.snapshot.EdgeCount} +} + +type scriptedFaucet[T any] struct { + batches [][]T + total int64 +} + +func (s scriptedFaucet[T]) Run(ctx context.Context, handle entityBatchHandler[T]) (int64, error) { + var processed int64 + for _, batch := range s.batches { + if err := ctx.Err(); err != nil { + return processed, err + } + + remaining := s.total - processed + if remaining <= 0 { + break + } + if int64(len(batch)) > remaining { + batch = batch[:int(remaining)] + } + + if err := handle(batch); err != nil { + return processed, err + } + processed += int64(len(batch)) + } + + return processed, nil +} diff --git a/retriever/transform.go b/retriever/transform.go new file mode 100644 index 00000000..a242a384 --- /dev/null +++ b/retriever/transform.go @@ -0,0 +1,108 @@ +package retriever + +import "github.com/specterops/dawgs/graph" + +type transformedBatch[T any] struct { + Records []T + ActionCounts []map[string]int +} + +type transformSession interface { + Metadata() ScrubMetadata + NeedsPreparation() bool + PrepareNode(*graph.Node) + TransformNodes([]*graph.Node) transformedBatch[normalizedNode] + TransformEdges([]*graph.Relationship) transformedBatch[normalizedEdge] +} + +func newTransformSession(options DumpOptions) (transformSession, error) { + if options.Scrub != ScrubFull { + return identityTransformSession{}, nil + } + + activeScrubber, err := newScrubber(options.ScrubConfig, options.Salt) + if err != nil { + return nil, err + } + return fullScrubTransformSession{scrubber: activeScrubber}, nil +} + +type identityTransformSession struct{} + +func (identityTransformSession) Metadata() ScrubMetadata { + return ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + } +} + +func (identityTransformSession) NeedsPreparation() bool { + return false +} + +func (identityTransformSession) PrepareNode(*graph.Node) {} + +func (identityTransformSession) TransformNodes(nodes []*graph.Node) transformedBatch[normalizedNode] { + batch := transformedBatch[normalizedNode]{ + Records: make([]normalizedNode, len(nodes)), + ActionCounts: make([]map[string]int, len(nodes)), + } + for index, node := range nodes { + batch.Records[index] = normalizeNode(node) + } + return batch +} + +func (identityTransformSession) TransformEdges(relationships []*graph.Relationship) transformedBatch[normalizedEdge] { + batch := transformedBatch[normalizedEdge]{ + Records: make([]normalizedEdge, len(relationships)), + ActionCounts: make([]map[string]int, len(relationships)), + } + for index, relationship := range relationships { + batch.Records[index] = normalizeEdge(relationship) + } + return batch +} + +type fullScrubTransformSession struct { + scrubber *scrubber +} + +func (s fullScrubTransformSession) Metadata() ScrubMetadata { + return s.scrubber.metadata() +} + +func (fullScrubTransformSession) NeedsPreparation() bool { + return true +} + +func (s fullScrubTransformSession) PrepareNode(node *graph.Node) { + s.scrubber.observeNode(node.Properties.MapOrEmpty()) +} + +func (s fullScrubTransformSession) TransformNodes(nodes []*graph.Node) transformedBatch[normalizedNode] { + batch := transformedBatch[normalizedNode]{ + Records: make([]normalizedNode, len(nodes)), + ActionCounts: make([]map[string]int, len(nodes)), + } + for index, node := range nodes { + record := normalizeNode(node) + record.Properties, batch.ActionCounts[index] = s.scrubber.scrubProperties(record.Properties) + batch.Records[index] = record + } + return batch +} + +func (s fullScrubTransformSession) TransformEdges(relationships []*graph.Relationship) transformedBatch[normalizedEdge] { + batch := transformedBatch[normalizedEdge]{ + Records: make([]normalizedEdge, len(relationships)), + ActionCounts: make([]map[string]int, len(relationships)), + } + for index, relationship := range relationships { + record := normalizeEdge(relationship) + record.Properties, batch.ActionCounts[index] = s.scrubber.scrubProperties(record.Properties) + batch.Records[index] = record + } + return batch +} diff --git a/retriever/transform_test.go b/retriever/transform_test.go new file mode 100644 index 00000000..49e60ee3 --- /dev/null +++ b/retriever/transform_test.go @@ -0,0 +1,174 @@ +package retriever + +import ( + "context" + "reflect" + "testing" + + "github.com/specterops/dawgs/graph" +) + +func TestIdentityTransformSession(t *testing.T) { + session := identityTransformSession{} + if session.NeedsPreparation() { + t.Fatalf("identity transform requires preparation") + } + if metadata := session.Metadata(); !reflect.DeepEqual(metadata, ScrubMetadata{ + Mode: ScrubNone, + NodeActionCounts: map[string]int{}, + EdgeActionCounts: map[string]int{}, + }) { + t.Fatalf("identity metadata = %+v", metadata) + } + + nodes := []*graph.Node{ + graph.NewNode(2, graph.AsProperties(map[string]any{"name": "alice"}), graph.StringKind("User"), graph.StringKind("Admin")), + graph.NewNode(3, nil), + } + if actual := session.TransformNodes(nodes); !reflect.DeepEqual(actual, transformedBatch[normalizedNode]{ + Records: []normalizedNode{ + {ID: "2", Kinds: []string{"Admin", "User"}, Properties: map[string]any{"name": "alice"}}, + {ID: "3", Kinds: []string{}, Properties: map[string]any{}}, + }, + ActionCounts: make([]map[string]int, 2), + }) { + t.Fatalf("identity node batch = %#v", actual) + } + + relationships := []*graph.Relationship{ + graph.NewRelationship(9, 2, 3, graph.AsProperties(map[string]any{"enabled": true}), graph.StringKind("MemberOf")), + } + if actual := session.TransformEdges(relationships); !reflect.DeepEqual(actual, transformedBatch[normalizedEdge]{ + Records: []normalizedEdge{ + {ID: "9", StartID: "2", EndID: "3", Kind: "MemberOf", Properties: map[string]any{"enabled": true}}, + }, + ActionCounts: make([]map[string]int, 1), + }) { + t.Fatalf("identity edge batch = %#v", actual) + } +} + +func TestFullScrubTransformSessionPreparationAndActions(t *testing.T) { + const sourceSID = "S-1-5-21-111111111-222222222-333333333-1001" + + transform, err := newTransformSession(DumpOptions{Scrub: ScrubFull, Salt: "transform-test"}) + if err != nil { + t.Fatalf("new transform: %v", err) + } + if !transform.NeedsPreparation() || transform.Metadata().Mode != ScrubFull { + t.Fatalf("full scrub transform metadata = %+v", transform.Metadata()) + } + + transform.PrepareNode(graph.NewNode(1, graph.AsProperties(map[string]any{"objectid": sourceSID}))) + nodeBatch := transform.TransformNodes([]*graph.Node{ + graph.NewNode(1, graph.AsProperties(map[string]any{ + "description": "private text", + "objectid": sourceSID, + })), + }) + if len(nodeBatch.Records) != 1 { + t.Fatalf("node records = %d", len(nodeBatch.Records)) + } + if !reflect.DeepEqual(nodeBatch.ActionCounts[0], map[string]int{ + string(actionPseudonymize): 1, + string(actionRedact): 1, + }) { + t.Fatalf("node actions = %+v", nodeBatch.ActionCounts[0]) + } + if nodeBatch.Records[0].Properties["objectid"] == sourceSID || nodeBatch.Records[0].Properties["description"] != "[REDACTED]" { + t.Fatalf("scrubbed node = %+v", nodeBatch.Records[0]) + } + + edgeBatch := transform.TransformEdges([]*graph.Relationship{ + graph.NewRelationship(4, 1, 2, graph.AsProperties(map[string]any{"owner_sid": sourceSID}), graph.StringKind("MemberOf")), + }) + if len(edgeBatch.Records) != 1 || !reflect.DeepEqual(edgeBatch.ActionCounts[0], map[string]int{ + string(actionPseudonymize): 1, + }) { + t.Fatalf("edge batch = %#v", edgeBatch) + } + if edgeBatch.Records[0].Properties["owner_sid"] != nodeBatch.Records[0].Properties["objectid"] { + t.Fatalf("prepared registry did not preserve reference equality") + } +} + +func TestDumpGraphTransformsAndObservesEachRecordOnce(t *testing.T) { + source := &scriptedGraphSource{ + snapshot: graphEntitySnapshot{NodeCount: 2, EdgeCount: 1}, + nodeBatches: [][]*graph.Node{{ + graph.NewNode(1, nil, graph.StringKind("User")), + graph.NewNode(2, nil, graph.StringKind("Group")), + }}, + edgeBatches: [][]*graph.Relationship{{ + graph.NewRelationship(3, 1, 2, nil, graph.StringKind("MemberOf")), + }}, + } + transform := &countingTransformSession{ + identity: identityTransformSession{}, + nodeCalls: map[string]int{}, + edgeCalls: map[string]int{}, + } + + options := DumpOptions{ + OutputDir: t.TempDir(), + Scrub: ScrubNone, + Compression: CompressionGzip, + ZstdLevel: DefaultZstdLevel, + ShardSize: 10, + BatchSize: 10, + } + workspace := newLocalCollectionWorkspace(options.OutputDir, false) + graphEntry, schema, metrics, err := dumpGraph( + context.Background(), + source, + GraphTarget{Name: "source"}, + options, + transform, + newShardSinkSet(newJSONLShardSink(newJSONLNodeSinkInWorkspace(options, workspace))), + newShardSinkSet(newJSONLShardSink(newJSONLEdgeSinkInWorkspace(options, workspace))), + nil, + ) + if err != nil { + t.Fatalf("dump graph: %v", err) + } + + if !reflect.DeepEqual(transform.nodeCalls, map[string]int{"1": 1, "2": 1}) || !reflect.DeepEqual(transform.edgeCalls, map[string]int{"3": 1}) { + t.Fatalf("transform calls: nodes=%v edges=%v", transform.nodeCalls, transform.edgeCalls) + } + if graphEntry.NodeCount != 2 || graphEntry.EdgeCount != 1 || metrics.NodeCount != 2 || metrics.EdgeCount != 1 { + t.Fatalf("observed totals: graph=%+v metrics=%+v", graphEntry, metrics) + } + if !reflect.DeepEqual(schema, GraphSchemaMetadata{Name: "source", NodeKinds: []string{"Group", "User"}, EdgeKinds: []string{"MemberOf"}}) { + t.Fatalf("observed schema = %+v", schema) + } +} + +type countingTransformSession struct { + identity identityTransformSession + nodeCalls map[string]int + edgeCalls map[string]int +} + +func (s *countingTransformSession) Metadata() ScrubMetadata { + return s.identity.Metadata() +} + +func (*countingTransformSession) NeedsPreparation() bool { + return false +} + +func (*countingTransformSession) PrepareNode(*graph.Node) {} + +func (s *countingTransformSession) TransformNodes(nodes []*graph.Node) transformedBatch[normalizedNode] { + for _, node := range nodes { + s.nodeCalls[node.ID.String()]++ + } + return s.identity.TransformNodes(nodes) +} + +func (s *countingTransformSession) TransformEdges(relationships []*graph.Relationship) transformedBatch[normalizedEdge] { + for _, relationship := range relationships { + s.edgeCalls[relationship.ID.String()]++ + } + return s.identity.TransformEdges(relationships) +} diff --git a/retriever/workspace.go b/retriever/workspace.go new file mode 100644 index 00000000..79df04ba --- /dev/null +++ b/retriever/workspace.go @@ -0,0 +1,198 @@ +package retriever + +import ( + "context" + "fmt" + "io" + "os" + "path" + "path/filepath" + "strings" +) + +type collectionWorkspace interface { + Root() string + Prepare(context.Context) error + Stage(context.Context, string) (stagedWorkspaceFile, error) + Publish(context.Context, string, []byte) (string, error) +} + +type stagedWorkspaceFile interface { + io.Writer + Close() error + Commit(context.Context) error + Abort() error +} + +type localCollectionWorkspace struct { + root string + force bool +} + +func newLocalCollectionWorkspace(root string, force bool) *localCollectionWorkspace { + return &localCollectionWorkspace{root: root, force: force} +} + +func (s *localCollectionWorkspace) Root() string { + return s.root +} + +func (s *localCollectionWorkspace) Prepare(ctx context.Context) error { + if err := ctx.Err(); err != nil { + return err + } + + info, err := os.Stat(s.root) + if err == nil { + if !info.IsDir() { + return fmt.Errorf("output path %q exists and is not a directory", s.root) + } + + entries, err := os.ReadDir(s.root) + if err != nil { + return fmt.Errorf("read output directory: %w", err) + } + + if len(entries) > 0 { + if !s.force { + return fmt.Errorf("output directory %q is not empty; pass -force to replace it", s.root) + } + + if err := os.RemoveAll(s.root); err != nil { + return fmt.Errorf("replace output directory: %w", err) + } + } + } else if !os.IsNotExist(err) { + return fmt.Errorf("inspect output directory: %w", err) + } + + if err := os.MkdirAll(s.root, 0o755); err != nil { + return fmt.Errorf("create output directory: %w", err) + } + + return nil +} + +func (s *localCollectionWorkspace) Stage(ctx context.Context, relativePath string) (stagedWorkspaceFile, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + + finalPath, err := s.resolve(relativePath) + if err != nil { + return nil, err + } + if err := os.MkdirAll(filepath.Dir(finalPath), 0o755); err != nil { + return nil, fmt.Errorf("create artifact directory: %w", err) + } + + tempPath := finalPath + ".tmp" + file, err := os.OpenFile(tempPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0o600) + if err != nil { + return nil, fmt.Errorf("open staged artifact: %w", err) + } + + return &localStagedWorkspaceFile{ + file: file, + finalPath: finalPath, + tempPath: tempPath, + }, nil +} + +func (s *localCollectionWorkspace) Publish(ctx context.Context, relativePath string, payload []byte) (string, error) { + artifact, err := s.Stage(ctx, relativePath) + if err != nil { + return "", err + } + + if _, err := artifact.Write(payload); err != nil { + return "", cleanupOnError(fmt.Errorf("write staged artifact: %w", err), artifact.Abort) + } + if err := artifact.Close(); err != nil { + return "", cleanupOnError(fmt.Errorf("close staged artifact: %w", err), artifact.Abort) + } + if err := artifact.Commit(ctx); err != nil { + return "", cleanupOnError(err, artifact.Abort) + } + + return s.resolve(relativePath) +} + +func (s *localCollectionWorkspace) resolve(relativePath string) (string, error) { + if relativePath == "" || path.Clean(relativePath) != relativePath || strings.ContainsRune(relativePath, '\\') || !filepath.IsLocal(filepath.FromSlash(relativePath)) { + return "", fmt.Errorf("unsafe workspace path %q", relativePath) + } + + return filepath.Join(s.root, filepath.FromSlash(relativePath)), nil +} + +type localStagedWorkspaceFile struct { + file *os.File + finalPath string + tempPath string + closed bool + committed bool + aborted bool +} + +func (s *localStagedWorkspaceFile) Write(payload []byte) (int, error) { + if s.closed || s.committed || s.aborted { + return 0, fmt.Errorf("write staged artifact after close, commit, or abort") + } + return s.file.Write(payload) +} + +func (s *localStagedWorkspaceFile) Close() error { + if s.committed || s.aborted { + return fmt.Errorf("close staged artifact after commit or abort") + } + if s.closed { + return nil + } + + s.closed = true + return s.file.Close() +} + +func (s *localStagedWorkspaceFile) Commit(ctx context.Context) error { + if s.committed || s.aborted || !s.closed { + return fmt.Errorf("commit artifact that is not staged and closed") + } + if err := ctx.Err(); err != nil { + return err + } + + if err := os.Rename(s.tempPath, s.finalPath); err != nil { + return fmt.Errorf("commit staged artifact: %w", err) + } + s.committed = true + return nil +} + +func (s *localStagedWorkspaceFile) Abort() error { + if s.committed { + return fmt.Errorf("abort committed artifact") + } + if s.aborted { + return nil + } + + s.aborted = true + var closeErr error + if !s.closed { + s.closed = true + closeErr = s.file.Close() + } + return collectErrors(closeErr, removeStagedArtifact(s.tempPath)) +} + +func removeStagedArtifact(path string) error { + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + return err + } + return nil +} + +func prepareOutputDirectory(outputDir string, force bool) error { + return newLocalCollectionWorkspace(outputDir, force).Prepare(context.Background()) +} diff --git a/retriever/workspace_test.go b/retriever/workspace_test.go new file mode 100644 index 00000000..849842a3 --- /dev/null +++ b/retriever/workspace_test.go @@ -0,0 +1,49 @@ +package retriever + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +func TestLocalCollectionWorkspacePublishesAtomically(t *testing.T) { + root := filepath.Join(t.TempDir(), "collection") + workspace := newLocalCollectionWorkspace(root, false) + if err := workspace.Prepare(context.Background()); err != nil { + t.Fatalf("prepare workspace: %v", err) + } + + publishedPath, err := workspace.Publish(context.Background(), "metadata/result.json", []byte("result\n")) + if err != nil { + t.Fatalf("publish artifact: %v", err) + } + if publishedPath != filepath.Join(root, "metadata", "result.json") { + t.Fatalf("published path = %q", publishedPath) + } + if payload, err := os.ReadFile(publishedPath); err != nil { + t.Fatalf("read published artifact: %v", err) + } else if string(payload) != "result\n" { + t.Fatalf("published payload = %q", payload) + } + if _, err := os.Stat(publishedPath + ".tmp"); !os.IsNotExist(err) { + t.Fatalf("staged artifact remains after publish: %v", err) + } +} + +func TestLocalCollectionWorkspaceRejectsUnsafePaths(t *testing.T) { + workspace := newLocalCollectionWorkspace(t.TempDir(), false) + for _, relativePath := range []string{ + "", + "/absolute", + "../parent", + "graphs/../escaped", + `graphs\windows`, + } { + t.Run(relativePath, func(t *testing.T) { + if _, err := workspace.Stage(context.Background(), relativePath); err == nil { + t.Fatalf("expected workspace path %q to be rejected", relativePath) + } + }) + } +}