Skip to content

Commit b38c746

Browse files
tae2089tae2089
committed
perf: batch PostgreSQL graph persistence
1 parent 8486b25 commit b38c746

34 files changed

Lines changed: 1095 additions & 287 deletions

‎cmd/ccg/main_test.go‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import (
55
"log/slog"
66
"os"
77
"path/filepath"
8+
"strconv"
89
"strings"
910
"testing"
1011
"time"
@@ -250,8 +251,8 @@ func TestRunMigrations_SQLiteDownRestoresNullableColumns(t *testing.T) {
250251
if err != nil {
251252
t.Fatalf("create migrator: %v", err)
252253
}
253-
if err := migrator.Steps(-6); err != nil {
254-
t.Fatalf("run down migration: %v", err)
254+
if err := migrator.Migrate(1); err != nil {
255+
t.Fatalf("migrate down to version 1: %v", err)
255256
}
256257

257258
var version migration.MigrationSchemaVersion
@@ -295,8 +296,8 @@ func TestRunMigrations_SQLiteDownFromVersionThreeDropsPolicyTables(t *testing.T)
295296
if err != nil {
296297
t.Fatalf("create migrator: %v", err)
297298
}
298-
if err := migrator.Steps(-5); err != nil {
299-
t.Fatalf("run down migration: %v", err)
299+
if err := migrator.Migrate(2); err != nil {
300+
t.Fatalf("migrate down to version 2: %v", err)
300301
}
301302

302303
var version migration.MigrationSchemaVersion
@@ -409,7 +410,8 @@ func TestEnsureSchemaVersion_LogsRuntimeSchemaPassAndFail(t *testing.T) {
409410
t.Fatalf("ensure schema version: %v", err)
410411
}
411412
passLog := logs.String()
412-
for _, want := range []string{"database runtime schema check passed", "driver=sqlite", "required_version=7", "auto_migrated=true"} {
413+
requiredVersionLog := "required_version=" + strconv.Itoa(migration.RequiredSchemaVersion)
414+
for _, want := range []string{"database runtime schema check passed", "driver=sqlite", requiredVersionLog, "auto_migrated=true"} {
413415
if !strings.Contains(passLog, want) {
414416
t.Fatalf("expected runtime schema pass log to contain %q, got %q", want, passLog)
415417
}
@@ -423,7 +425,7 @@ func TestEnsureSchemaVersion_LogsRuntimeSchemaPassAndFail(t *testing.T) {
423425
t.Fatal("expected parity failure")
424426
}
425427
failLog := logs.String()
426-
for _, want := range []string{"database runtime schema check failed", "driver=sqlite", "required_version=7", "auto_migrated=false", "error.message="} {
428+
for _, want := range []string{"database runtime schema check failed", "driver=sqlite", requiredVersionLog, "auto_migrated=false", "error.message="} {
427429
if !strings.Contains(failLog, want) {
428430
t.Fatalf("expected runtime schema failure log to contain %q, got %q", want, failLog)
429431
}

‎internal/adapters/inbound/cli/build_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ func setupBuildTest(t *testing.T) (*Deps, *bytes.Buffer, *bytes.Buffer, *gorm.DB
2222
t.Helper()
2323
deps, stdout, stderr := newTestDeps()
2424

25-
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: gormlogger.Discard})
25+
db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "ccg.db")), &gorm.Config{Logger: gormlogger.Discard})
2626
if err != nil {
2727
t.Fatal(err)
2828
}

‎internal/adapters/inbound/cli/update_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ func setupUpdateTest(t *testing.T) (*Deps, *bytes.Buffer, *bytes.Buffer, *gorm.D
2626
t.Helper()
2727
deps, stdout, stderr := newTestDeps()
2828

29-
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: gormlogger.Discard})
29+
db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "ccg.db")), &gorm.Config{Logger: gormlogger.Discard})
3030
if err != nil {
3131
t.Fatal(err)
3232
}
@@ -61,7 +61,7 @@ func setupUpdateGraphOnlyTest(t *testing.T) (*Deps, *bytes.Buffer, *bytes.Buffer
6161
t.Helper()
6262
deps, stdout, stderr := newTestDeps()
6363

64-
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: gormlogger.Discard})
64+
db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "ccg.db")), &gorm.Config{Logger: gormlogger.Discard})
6565
if err != nil {
6666
t.Fatal(err)
6767
}

‎internal/adapters/outbound/graphgorm/annotation_bulk_test.go‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
package graphgorm
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"sync/atomic"
7+
"testing"
8+
"time"
9+
10+
"gorm.io/driver/sqlite"
11+
"gorm.io/gorm"
12+
"gorm.io/gorm/logger"
13+
14+
requestctx "github.com/tae2089/code-context-graph/internal/ctx"
15+
"github.com/tae2089/code-context-graph/internal/domain/graph"
16+
)
17+
18+
type annotationSQLCounter struct {
19+
statements atomic.Int64
20+
}
21+
22+
func (c *annotationSQLCounter) LogMode(logger.LogLevel) logger.Interface { return c }
23+
func (c *annotationSQLCounter) Info(context.Context, string, ...any) {}
24+
func (c *annotationSQLCounter) Warn(context.Context, string, ...any) {}
25+
func (c *annotationSQLCounter) Error(context.Context, string, ...any) {}
26+
func (c *annotationSQLCounter) Trace(_ context.Context, _ time.Time, fc func() (string, int64), _ error) {
27+
c.statements.Add(1)
28+
_, _ = fc()
29+
}
30+
31+
func (c *annotationSQLCounter) reset() {
32+
c.statements.Store(0)
33+
}
34+
35+
func setupAnnotationSQLMeasurement(t *testing.T, count int) (*Store, context.Context, []*graph.Annotation, *annotationSQLCounter) {
36+
t.Helper()
37+
counter := &annotationSQLCounter{}
38+
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: counter})
39+
if err != nil {
40+
t.Fatalf("open SQLite: %v", err)
41+
}
42+
store := New(db)
43+
if err := store.AutoMigrate(); err != nil {
44+
t.Fatalf("auto migrate: %v", err)
45+
}
46+
ctx := requestctx.WithNamespace(context.Background(), "annotation-sql-measurement")
47+
nodes := make([]graph.Node, count)
48+
for i := range nodes {
49+
nodes[i] = graph.Node{
50+
QualifiedName: fmt.Sprintf("sample.F%d", i),
51+
Kind: graph.NodeKindFunction,
52+
Name: fmt.Sprintf("F%d", i),
53+
FilePath: "sample.go",
54+
StartLine: i + 1,
55+
EndLine: i + 1,
56+
Language: "go",
57+
}
58+
}
59+
if err := store.UpsertNodes(ctx, nodes); err != nil {
60+
t.Fatalf("upsert nodes: %v", err)
61+
}
62+
stored, err := store.GetNodesByFile(ctx, "sample.go")
63+
if err != nil {
64+
t.Fatalf("load nodes: %v", err)
65+
}
66+
annotations := make([]*graph.Annotation, len(stored))
67+
for i := range stored {
68+
annotations[i] = &graph.Annotation{
69+
NodeID: stored[i].ID,
70+
Summary: fmt.Sprintf("summary %d", i),
71+
Tags: []graph.DocTag{{
72+
Kind: graph.TagIntent,
73+
Value: fmt.Sprintf("intent %d", i),
74+
Ordinal: 0,
75+
}},
76+
}
77+
}
78+
counter.reset()
79+
return store, ctx, annotations, counter
80+
}
81+
82+
func TestUpsertAnnotations_BatchesSQLAndPreservesTags(t *testing.T) {
83+
store, ctx, annotations, counter := setupAnnotationSQLMeasurement(t, 20)
84+
started := time.Now()
85+
if err := store.UpsertAnnotations(ctx, annotations); err != nil {
86+
t.Fatalf("upsert annotations: %v", err)
87+
}
88+
statements := counter.statements.Load()
89+
t.Logf("bulk annotation candidate: annotations=%d statements=%d elapsed=%s", len(annotations), statements, time.Since(started))
90+
if statements > 5 {
91+
t.Fatalf("bulk annotation statements = %d, want <= 5", statements)
92+
}
93+
94+
for i, annotation := range annotations {
95+
got, err := store.GetAnnotation(ctx, annotation.NodeID)
96+
if err != nil {
97+
t.Fatalf("get annotation %d: %v", i, err)
98+
}
99+
if got == nil || got.Summary != annotation.Summary {
100+
t.Fatalf("annotation %d = %+v, want summary %q", i, got, annotation.Summary)
101+
}
102+
if len(got.Tags) != 1 || got.Tags[0].Value != annotation.Tags[0].Value {
103+
t.Fatalf("annotation %d tags = %+v, want %+v", i, got.Tags, annotation.Tags)
104+
}
105+
}
106+
107+
for i, annotation := range annotations {
108+
annotation.Summary = fmt.Sprintf("updated summary %d", i)
109+
annotation.Tags = []graph.DocTag{{Kind: graph.TagDomainRule, Value: fmt.Sprintf("rule %d", i), Ordinal: 0}}
110+
}
111+
counter.reset()
112+
if err := store.UpsertAnnotations(ctx, annotations); err != nil {
113+
t.Fatalf("update annotations: %v", err)
114+
}
115+
if statements := counter.statements.Load(); statements > 5 {
116+
t.Fatalf("bulk annotation update statements = %d, want <= 5", statements)
117+
}
118+
for i, annotation := range annotations {
119+
got, err := store.GetAnnotation(ctx, annotation.NodeID)
120+
if err != nil {
121+
t.Fatalf("get updated annotation %d: %v", i, err)
122+
}
123+
if got == nil || got.Summary != annotation.Summary || len(got.Tags) != 1 || got.Tags[0].Value != annotation.Tags[0].Value {
124+
t.Fatalf("updated annotation %d = %+v, want %+v", i, got, annotation)
125+
}
126+
}
127+
}
128+
129+
func TestUpsertAnnotations_RejectsForeignNodeBeforeWrites(t *testing.T) {
130+
store, ctx, annotations, _ := setupAnnotationSQLMeasurement(t, 1)
131+
foreignCtx := requestctx.WithNamespace(context.Background(), "foreign")
132+
foreignNodes := []graph.Node{{
133+
QualifiedName: "foreign.F",
134+
Kind: graph.NodeKindFunction,
135+
Name: "F",
136+
FilePath: "foreign.go",
137+
StartLine: 1,
138+
EndLine: 1,
139+
Language: "go",
140+
}}
141+
if err := store.UpsertNodes(foreignCtx, foreignNodes); err != nil {
142+
t.Fatalf("upsert foreign node: %v", err)
143+
}
144+
storedForeign, err := store.GetNodesByFile(foreignCtx, "foreign.go")
145+
if err != nil || len(storedForeign) != 1 {
146+
t.Fatalf("load foreign node: nodes=%+v err=%v", storedForeign, err)
147+
}
148+
foreignAnnotation := &graph.Annotation{NodeID: storedForeign[0].ID, Summary: "foreign"}
149+
150+
err = store.UpsertAnnotations(ctx, []*graph.Annotation{annotations[0], foreignAnnotation})
151+
if err == nil {
152+
t.Fatal("expected foreign-node batch to be rejected")
153+
}
154+
got, getErr := store.GetAnnotation(ctx, annotations[0].NodeID)
155+
if getErr != nil {
156+
t.Fatalf("get owned annotation after rejection: %v", getErr)
157+
}
158+
if got != nil {
159+
t.Fatalf("owned annotation was written before batch rejection: %+v", got)
160+
}
161+
}

‎internal/adapters/outbound/graphgorm/deletegraph_query_test.go‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
package graphgorm
2+
3+
import (
4+
"context"
5+
"strings"
6+
"sync"
7+
"testing"
8+
"time"
9+
10+
"gorm.io/driver/sqlite"
11+
"gorm.io/gorm"
12+
"gorm.io/gorm/logger"
13+
14+
requestctx "github.com/tae2089/code-context-graph/internal/ctx"
15+
"github.com/tae2089/code-context-graph/internal/domain/graph"
16+
)
17+
18+
type deleteGraphSQLRecorder struct {
19+
mu sync.Mutex
20+
statements []string
21+
}
22+
23+
func (r *deleteGraphSQLRecorder) LogMode(logger.LogLevel) logger.Interface { return r }
24+
func (r *deleteGraphSQLRecorder) Info(context.Context, string, ...any) {}
25+
func (r *deleteGraphSQLRecorder) Warn(context.Context, string, ...any) {}
26+
func (r *deleteGraphSQLRecorder) Error(context.Context, string, ...any) {}
27+
func (r *deleteGraphSQLRecorder) Trace(_ context.Context, _ time.Time, fc func() (string, int64), _ error) {
28+
sql, _ := fc()
29+
r.mu.Lock()
30+
r.statements = append(r.statements, sql)
31+
r.mu.Unlock()
32+
}
33+
34+
func (r *deleteGraphSQLRecorder) reset() {
35+
r.mu.Lock()
36+
r.statements = nil
37+
r.mu.Unlock()
38+
}
39+
40+
func (r *deleteGraphSQLRecorder) searchDocumentDelete() string {
41+
r.mu.Lock()
42+
defer r.mu.Unlock()
43+
for _, statement := range r.statements {
44+
if strings.Contains(statement, "DELETE FROM `search_documents`") {
45+
return statement
46+
}
47+
}
48+
return ""
49+
}
50+
51+
func TestDeleteGraph_DeletesSearchDocumentsByNamespace(t *testing.T) {
52+
recorder := &deleteGraphSQLRecorder{}
53+
db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{Logger: recorder})
54+
if err != nil {
55+
t.Fatalf("open SQLite: %v", err)
56+
}
57+
store := New(db)
58+
if err := store.AutoMigrate(); err != nil {
59+
t.Fatalf("auto migrate: %v", err)
60+
}
61+
if err := db.AutoMigrate(&graph.SearchDocument{}); err != nil {
62+
t.Fatalf("migrate search documents: %v", err)
63+
}
64+
65+
ctx := requestctx.WithNamespace(context.Background(), "delete-me")
66+
if err := store.UpsertNodes(ctx, []graph.Node{{
67+
QualifiedName: "pkg.DeleteMe",
68+
Kind: graph.NodeKindFunction,
69+
Name: "DeleteMe",
70+
FilePath: "delete.go",
71+
StartLine: 1,
72+
EndLine: 1,
73+
Language: "go",
74+
}}); err != nil {
75+
t.Fatalf("upsert node: %v", err)
76+
}
77+
node, err := store.GetNode(ctx, "pkg.DeleteMe")
78+
if err != nil {
79+
t.Fatalf("get node: %v", err)
80+
}
81+
if err := db.Create(&graph.SearchDocument{
82+
Namespace: "delete-me",
83+
NodeID: node.ID,
84+
Content: "delete me",
85+
Language: "go",
86+
}).Error; err != nil {
87+
t.Fatalf("create search document: %v", err)
88+
}
89+
90+
recorder.reset()
91+
if err := store.DeleteGraph(ctx); err != nil {
92+
t.Fatalf("DeleteGraph: %v", err)
93+
}
94+
95+
statement := recorder.searchDocumentDelete()
96+
if statement == "" {
97+
t.Fatal("search_documents DELETE statement not recorded")
98+
}
99+
if !strings.Contains(statement, "WHERE namespace = \"delete-me\"") {
100+
t.Fatalf("search_documents DELETE = %q, want direct namespace predicate", statement)
101+
}
102+
if strings.Contains(statement, "SELECT `id` FROM `nodes`") {
103+
t.Fatalf("search_documents DELETE = %q, must not select node IDs", statement)
104+
}
105+
}

0 commit comments

Comments
 (0)