Skip to content

Commit 5defd21

Browse files
lfittllfittl
authored
Add Summary helper to wrap pg_query_summary function (#147)
1 parent 926ef9b commit 5defd21

3 files changed

Lines changed: 164 additions & 0 deletions

File tree

‎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: 135 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,135 @@
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+
// JOIN with table aliases
41+
{
42+
input: "SELECT * FROM users u JOIN orders o ON u.id = o.user_id",
43+
truncateLimit: -1,
44+
expected: &pg_query.SummaryResult{
45+
Tables: []*pg_query.SummaryResult_Table{
46+
{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select},
47+
{Name: "orders", TableName: "orders", Context: pg_query.SummaryResult_Select},
48+
},
49+
Aliases: map[string]string{"u": "users", "o": "orders"},
50+
StatementTypes: []string{"SelectStmt"},
51+
},
52+
},
53+
// Aggregate functions (count, max)
54+
{
55+
input: "SELECT count(*), max(id) FROM users",
56+
truncateLimit: -1,
57+
expected: &pg_query.SummaryResult{
58+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select}},
59+
Functions: []*pg_query.SummaryResult_Function{
60+
{Name: "count", FunctionName: "count", Context: pg_query.SummaryResult_Call},
61+
{Name: "max", FunctionName: "max", Context: pg_query.SummaryResult_Call},
62+
},
63+
StatementTypes: []string{"SelectStmt"},
64+
},
65+
},
66+
// Common Table Expression (CTE)
67+
{
68+
input: "WITH active_users AS (SELECT * FROM users WHERE active = true) SELECT * FROM active_users",
69+
truncateLimit: -1,
70+
expected: &pg_query.SummaryResult{
71+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_Select}},
72+
CteNames: []string{"active_users"},
73+
FilterColumns: []*pg_query.SummaryResult_FilterColumn{{Column: "active"}},
74+
StatementTypes: []string{"SelectStmt"},
75+
},
76+
},
77+
// Schema-qualified table name
78+
{
79+
input: "SELECT * FROM public.users",
80+
truncateLimit: -1,
81+
expected: &pg_query.SummaryResult{
82+
Tables: []*pg_query.SummaryResult_Table{{Name: "public.users", SchemaName: "public", TableName: "users", Context: pg_query.SummaryResult_Select}},
83+
StatementTypes: []string{"SelectStmt"},
84+
},
85+
},
86+
// INSERT statement
87+
{
88+
input: "INSERT INTO users (name) VALUES ('test')",
89+
truncateLimit: -1,
90+
expected: &pg_query.SummaryResult{
91+
Tables: []*pg_query.SummaryResult_Table{{Name: "users", TableName: "users", Context: pg_query.SummaryResult_DML}},
92+
StatementTypes: []string{"InsertStmt", "SelectStmt"},
93+
},
94+
},
95+
// UPDATE statement
96+
{
97+
input: "UPDATE users SET name = 'test' WHERE id = 1",
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{"UpdateStmt"},
102+
},
103+
},
104+
// DELETE statement
105+
{
106+
input: "DELETE FROM users 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{"DeleteStmt"},
111+
},
112+
},
113+
}
114+
115+
func TestSummary(t *testing.T) {
116+
for _, tt := range summaryTests {
117+
t.Run(tt.input, func(t *testing.T) {
118+
result, err := pg_query.Summary(tt.input, tt.truncateLimit)
119+
if err != nil {
120+
t.Fatalf("Summary() error: %s", err)
121+
}
122+
123+
if !proto.Equal(result, tt.expected) {
124+
t.Errorf("Summary() mismatch:\n got: %v\n expected: %v", result, tt.expected)
125+
}
126+
})
127+
}
128+
}
129+
130+
func TestSummaryError(t *testing.T) {
131+
_, err := pg_query.Summary("SELECT * FROM", -1)
132+
if err == nil {
133+
t.Error("expected error for invalid SQL")
134+
}
135+
}

0 commit comments

Comments
 (0)