Skip to content

Commit 3cd3ada

Browse files
lfittllfittl
committed
Add Summary helper to wrap pg_query_summary function
1 parent 926ef9b commit 3cd3ada

File tree

3 files changed

+173
-0
lines changed

3 files changed

+173
-0
lines changed

‎parser/parser.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -309,3 +309,18 @@ func IsUtilityStmt(input string) (result []bool, err error) {
309309

310310
return
311311
}
312+
313+
// SummaryToProtobuf - Extracts summary information from SQL statement
314+
func SummaryToProtobuf(input string, truncateLimit int) ([]byte, error) {
315+
inputC := C.CString(input)
316+
defer C.free(unsafe.Pointer(inputC))
317+
318+
resultC := C.pg_query_summary(inputC, C.int(0), C.int(truncateLimit))
319+
defer C.pg_query_free_summary_parse_result(resultC)
320+
321+
if resultC.error != nil {
322+
return nil, newPgQueryError(resultC.error)
323+
}
324+
325+
return toBytes(C.GoStringN(resultC.summary.data, C.int(resultC.summary.len))), nil
326+
}

‎pg_query.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -92,3 +92,17 @@ func SplitWithParser(input string, trimSpace bool) (result []string, err error)
9292
func IsUtilityStmt(input string) (result []bool, err error) {
9393
return parser.IsUtilityStmt(input)
9494
}
95+
96+
// Summary - Extracts summary information from SQL statement
97+
//
98+
// Optionally, you can pass a positive numbered truncateLimit to return a
99+
// "smart" truncated version of the input statement that is at most limit length.
100+
func Summary(input string, truncateLimit int) (result *SummaryResult, err error) {
101+
protobufSummary, err := parser.SummaryToProtobuf(input, truncateLimit)
102+
if err != nil {
103+
return
104+
}
105+
result = &SummaryResult{}
106+
err = proto.Unmarshal(protobufSummary, result)
107+
return
108+
}

‎summary_test.go‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,144 @@
1+
//go:build cgo
2+
// +build cgo
3+
4+
package pg_query_test
5+
6+
import (
7+
"testing"
8+
9+
"google.golang.org/protobuf/proto"
10+
11+
pg_query "github.com/pganalyze/pg_query_go/v6"
12+
)
13+
14+
var summaryTests = []struct {
15+
input string
16+
truncateLimit int
17+
expected *pg_query.SummaryResult
18+
}{
19+
// Basic SELECT with filter column
20+
{
21+
input: "SELECT * FROM users WHERE id = 1",
22+
truncateLimit: -1,
23+
expected: &pg_query.SummaryResult{
24+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select}},
25+
FilterColumns: []*pg_query.SummaryResult_FilterColumn{{Column: "id"}},
26+
StatementTypes: []string{"SelectStmt"},
27+
},
28+
},
29+
// Query truncation
30+
{
31+
input: "SELECT id, name, email FROM users WHERE id = 1",
32+
truncateLimit: 30,
33+
expected: &pg_query.SummaryResult{
34+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select}},
35+
FilterColumns: []*pg_query.SummaryResult_FilterColumn{{Column: "id"}},
36+
StatementTypes: []string{"SelectStmt"},
37+
TruncatedQuery: "SELECT ... FROM users WHERE...",
38+
},
39+
},
40+
// No truncation when disabled (truncateLimit = -1)
41+
{
42+
input: "SELECT * FROM users",
43+
truncateLimit: -1,
44+
expected: &pg_query.SummaryResult{
45+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select}},
46+
StatementTypes: []string{"SelectStmt"},
47+
},
48+
},
49+
// JOIN with table aliases
50+
{
51+
input: "SELECT * FROM users u JOIN orders o ON u.id = o.user_id",
52+
truncateLimit: -1,
53+
expected: &pg_query.SummaryResult{
54+
Tables: []*pg_query.SummaryResult_Table{
55+
{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select},
56+
{Name: "orders", TableName: "orders", Context: pg_query.SummaryResult_Select},
57+
},
58+
Aliases: map[string]string{"u": "users", "o": "orders"},
59+
StatementTypes: []string{"SelectStmt"},
60+
},
61+
},
62+
// Aggregate functions (count, max)
63+
{
64+
input: "SELECT count(*), max(id) FROM users",
65+
truncateLimit: -1,
66+
expected: &pg_query.SummaryResult{
67+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select}},
68+
Functions: []*pg_query.SummaryResult_Function{
69+
{Name: "count", FunctionName: "count", Context: pg_query.SummaryResult_Call},
70+
{Name: "max", FunctionName: "max", Context: pg_query.SummaryResult_Call},
71+
},
72+
StatementTypes: []string{"SelectStmt"},
73+
},
74+
},
75+
// Common Table Expression (CTE)
76+
{
77+
input: "WITH active_users AS (SELECT * FROM users WHERE active = true) SELECT * FROM active_users",
78+
truncateLimit: -1,
79+
expected: &pg_query.SummaryResult{
80+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select}},
81+
CteNames: []string{"active_users"},
82+
FilterColumns: []*pg_query.SummaryResult_FilterColumn{{Column: "active"}},
83+
StatementTypes: []string{"SelectStmt"},
84+
},
85+
},
86+
// Schema-qualified table name
87+
{
88+
input: "SELECT * FROM public.users",
89+
truncateLimit: -1,
90+
expected: &pg_query.SummaryResult{
91+
Tables: []*pg_query.SummaryResult_Table{{Name: "public.users", SchemaName: "public", TableName: "users", Context: pg_query.SummaryResult_Select}},
92+
StatementTypes: []string{"SelectStmt"},
93+
},
94+
},
95+
// INSERT statement
96+
{
97+
input: "INSERT INTO users (name) VALUES ('test')",
98+
truncateLimit: -1,
99+
expected: &pg_query.SummaryResult{
100+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_DML}},
101+
StatementTypes: []string{"InsertStmt", "SelectStmt"},
102+
},
103+
},
104+
// UPDATE statement
105+
{
106+
input: "UPDATE users SET name = 'test' WHERE id = 1",
107+
truncateLimit: -1,
108+
expected: &pg_query.SummaryResult{
109+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_DML}},
110+
StatementTypes: []string{"UpdateStmt"},
111+
},
112+
},
113+
// DELETE statement
114+
{
115+
input: "DELETE FROM users WHERE id = 1",
116+
truncateLimit: -1,
117+
expected: &pg_query.SummaryResult{
118+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_DML}},
119+
StatementTypes: []string{"DeleteStmt"},
120+
},
121+
},
122+
}
123+
124+
func TestSummary(t *testing.T) {
125+
for _, tt := range summaryTests {
126+
t.Run(tt.input, func(t *testing.T) {
127+
result, err := pg_query.Summary(tt.input, tt.truncateLimit)
128+
if err != nil {
129+
t.Fatalf("Summary() error: %s", err)
130+
}
131+
132+
if !proto.Equal(result, tt.expected) {
133+
t.Errorf("Summary() mismatch:\n got: %v\n expected: %v", result, tt.expected)
134+
}
135+
})
136+
}
137+
}
138+
139+
func TestSummaryError(t *testing.T) {
140+
_, err := pg_query.Summary("SELECT * FROM", -1)
141+
if err == nil {
142+
t.Error("expected error for invalid SQL")
143+
}
144+
}

0 commit comments

Comments
 (0)