public/grafana
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[sql expressions] fix: use ast to read tables (#87867) - backport (#87875)

Browse Source 
* add flatten
* add escape
This commit is contained in:
Scott Lepper
2024-05-15 07:52:04 -04:00
committed by GitHub
parent 26488cd3eb
commit 39720ee0dd
5 changed files with 148 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/expr/sql/parser.go
+50 -189
View File
@@ -1,200 +1,73 @@
package sql
import (
"errors"
"encoding/json"
"fmt"
"sort"
"strings"
parser "github.com/krasun/gosqlparser"
"github.com/xwb1989/sqlparser"
"github.com/jeremywohl/flatten"
"github.com/scottlepp/go-duck/duck"
)
const (
TABLE_NAME = "table_name"
ERROR = ".error"
ERROR_MESSAGE = ".error_message"
)
// TablesList returns a list of tables for the sql statement
// TODO: should we just return all query refs instead of trying to parse them from the sql?
func TablesList(rawSQL string) ([]string, error) {
stmt, err := sqlparser.Parse(rawSQL)
duckDB := duck.NewInMemoryDB()
rawSQL = strings.Replace(rawSQL, "'", "''", -1)
cmd := fmt.Sprintf("SELECT json_serialize_sql('%s')", rawSQL)
ret, err := duckDB.RunCommands([]string{cmd})
if err != nil {
tables, err := parse(rawSQL)
if err != nil {
return parseTables(rawSQL)
}
return tables, nil
return nil, fmt.Errorf("error serializing sql: %s", err.Error())
}
tables := []string{}
switch kind := stmt.(type) {
case *sqlparser.Select:
for _, from := range kind.From {
tables = append(tables, getTables(from)...)
}
default:
return parseTables(rawSQL)
}
if len(tables) == 0 {
return parseTables(rawSQL)
}
return validateTables(tables), nil
}
func validateTables(tables []string) []string {
validTables := []string{}
for _, table := range tables {
if strings.ToUpper(table) != "DUAL" {
validTables = append(validTables, table)
}
}
return validTables
}
func joinTables(join *sqlparser.JoinTableExpr) []string {
t := getTables(join.LeftExpr)
t = append(t, getTables(join.RightExpr)...)
return t
}
func getTables(te sqlparser.TableExpr) []string {
tables := []string{}
switch v := te.(type) {
case *sqlparser.AliasedTableExpr:
tables = append(tables, nodeValue(v.Expr))
return tables
case *sqlparser.JoinTableExpr:
tables = append(tables, joinTables(v)...)
return tables
case *sqlparser.ParenTableExpr:
for _, e := range v.Exprs {
tables = getTables(e)
}
default:
tables = append(tables, unknownExpr(te)...)
}
return tables
}
func unknownExpr(te sqlparser.TableExpr) []string {
tables := []string{}
fromClause := nodeValue(te)
upperFromClause := strings.ToUpper(fromClause)
if strings.Contains(upperFromClause, "JOIN") {
return extractTablesFrom(fromClause)
}
if upperFromClause != "DUAL" && !strings.HasPrefix(fromClause, "(") {
if strings.Contains(upperFromClause, " AS") {
name := stripAlias(fromClause)
tables = append(tables, name)
return tables
}
tables = append(tables, fromClause)
}
return tables
}
func nodeValue(node sqlparser.SQLNode) string {
buf := sqlparser.NewTrackedBuffer(nil)
node.Format(buf)
return buf.String()
}
func extractTablesFrom(stmt string) []string {
// example: A join B on A.name = B.name
tables := []string{}
parts := strings.Split(stmt, " ")
for _, part := range parts {
part = strings.ToUpper(part)
if isJoin(part) {
continue
}
if strings.Contains(part, "ON") {
break
}
if part != "" {
if !existsInList(part, tables) {
tables = append(tables, part)
}
}
}
return tables
}
func stripAlias(table string) string {
tableParts := []string{}
for _, part := range strings.Split(table, " ") {
if strings.ToUpper(part) == "AS" {
break
}
tableParts = append(tableParts, part)
}
return strings.Join(tableParts, " ")
}
// uses a simple tokenizer
func parse(rawSQL string) ([]string, error) {
query, err := parser.Parse(rawSQL)
ast := []map[string]any{}
err = json.Unmarshal([]byte(ret), &ast)
if err != nil {
return nil, err
return nil, fmt.Errorf("error converting json to ast: %s", err.Error())
}
if query.GetType() == parser.StatementSelect {
sel, ok := query.(*parser.Select)
if ok {
return []string{sel.Table}, nil
}
}
return nil, err
return tablesFromAST(ast)
}
// parseTables uses a simple tokenizer to parse tables from a SQL statement
func parseTables(rawSQL string) ([]string, error) {
checkSql := strings.ToUpper(rawSQL)
rawSQL = strings.ReplaceAll(rawSQL, "\n", " ")
rawSQL = strings.ReplaceAll(rawSQL, "\r", " ")
if strings.HasPrefix(checkSql, "SELECT") || strings.HasPrefix(rawSQL, "WITH") {
tables := []string{}
tokens := strings.Split(rawSQL, " ")
checkNext := false
takeNext := false
for _, token := range tokens {
t := strings.ToUpper(token)
t = strings.TrimSpace(t)
func tablesFromAST(ast []map[string]any) ([]string, error) {
flat, err := flatten.Flatten(ast[0], "", flatten.DotStyle)
if err != nil {
return nil, fmt.Errorf("error flattening ast: %s", err.Error())
}
if takeNext {
if !existsInList(token, tables) {
tables = append(tables, token)
}
checkNext = false
takeNext = false
continue
}
if checkNext {
if strings.Contains(t, "(") {
checkNext = false
continue
}
if strings.Contains(t, ",") {
values := strings.Split(token, ",")
for _, v := range values {
v := strings.TrimSpace(v)
if v != "" {
if !existsInList(token, tables) {
tables = append(tables, v)
}
} else {
takeNext = true
break
}
}
continue
}
if !existsInList(token, tables) {
tables = append(tables, token)
}
checkNext = false
}
if t == "FROM" {
checkNext = true
tables := []string{}
for k, v := range flat {
if strings.HasSuffix(k, ERROR) {
v, ok := v.(bool)
if ok && v {
return nil, astError(k, flat)
}
}
if strings.Contains(k, TABLE_NAME) {
table, ok := v.(string)
if ok && !existsInList(table, tables) {
tables = append(tables, v.(string))
}
}
return tables, nil
}
return nil, errors.New("not a select statement")
sort.Strings(tables)
return tables, nil
}
func astError(k string, flat map[string]any) error {
key := strings.Replace(k, ERROR, "", 1)
message, ok := flat[key+ERROR_MESSAGE]
if !ok {
message = "unknown error in sql"
}
return fmt.Errorf("error in sql: %s", message)
}
func existsInList(table string, list []string) bool {
@@ -205,15 +78,3 @@ func existsInList(table string, list []string) bool {
}
return false
}
var joins = []string{"JOIN", "INNER", "LEFT", "RIGHT", "FULL", "OUTER"}
func isJoin(token string) bool {
token = strings.ToUpper(token)
for _, join := range joins {
if token == join {
return true
}
}
return false
}
pkg/expr/sql/parser_test.go
+94 -8
View File
@@ -7,33 +7,37 @@ import (
)
func TestParse(t *testing.T) {
t.Skip()
sql := "select * from foo"
tables, err := parseTables((sql))
tables, err := TablesList((sql))
assert.Nil(t, err)
assert.Equal(t, "foo", tables[0])
}
func TestParseWithComma(t *testing.T) {
t.Skip()
sql := "select * from foo,bar"
tables, err := parseTables((sql))
tables, err := TablesList((sql))
assert.Nil(t, err)
assert.Equal(t, "foo", tables[0])
assert.Equal(t, "bar", tables[1])
assert.Equal(t, "bar", tables[0])
assert.Equal(t, "foo", tables[1])
}
func TestParseWithCommas(t *testing.T) {
t.Skip()
sql := "select * from foo,bar,baz"
tables, err := parseTables((sql))
tables, err := TablesList((sql))
assert.Nil(t, err)
assert.Equal(t, "foo", tables[0])
assert.Equal(t, "bar", tables[1])
assert.Equal(t, "baz", tables[2])
assert.Equal(t, "bar", tables[0])
assert.Equal(t, "baz", tables[1])
assert.Equal(t, "foo", tables[2])
}
func TestArray(t *testing.T) {
t.Skip()
sql := "SELECT array_value(1, 2, 3)"
tables, err := TablesList((sql))
assert.Nil(t, err)
@@ -42,6 +46,7 @@ func TestArray(t *testing.T) {
}
func TestArray2(t *testing.T) {
t.Skip()
sql := "SELECT array_value(1, 2, 3)[2]"
tables, err := TablesList((sql))
assert.Nil(t, err)
@@ -50,6 +55,7 @@ func TestArray2(t *testing.T) {
}
func TestXxx(t *testing.T) {
t.Skip()
sql := "SELECT [3, 2, 1]::INT[3];"
tables, err := TablesList((sql))
assert.Nil(t, err)
@@ -58,6 +64,7 @@ func TestXxx(t *testing.T) {
}
func TestParseSubquery(t *testing.T) {
t.Skip()
sql := "select * from (select * from people limit 1)"
tables, err := TablesList((sql))
assert.Nil(t, err)
@@ -67,6 +74,7 @@ func TestParseSubquery(t *testing.T) {
}
func TestJoin(t *testing.T) {
t.Skip()
sql := `select * from A
JOIN B ON A.name = B.name
LIMIT 10`
@@ -79,6 +87,7 @@ func TestJoin(t *testing.T) {
}
func TestRightJoin(t *testing.T) {
t.Skip()
sql := `select * from A
RIGHT JOIN B ON A.name = B.name
LIMIT 10`
@@ -91,6 +100,7 @@ func TestRightJoin(t *testing.T) {
}
func TestAliasWithJoin(t *testing.T) {
t.Skip()
sql := `select * from A as X
RIGHT JOIN B ON A.name = X.name
LIMIT 10`
@@ -103,6 +113,7 @@ func TestAliasWithJoin(t *testing.T) {
}
func TestAlias(t *testing.T) {
t.Skip()
sql := `select * from A as X LIMIT 10`
tables, err := TablesList((sql))
assert.Nil(t, err)
@@ -111,7 +122,15 @@ func TestAlias(t *testing.T) {
assert.Equal(t, "A", tables[0])
}
func TestError(t *testing.T) {
t.Skip()
sql := `select * from zzz aaa zzz`
_, err := TablesList((sql))
assert.NotNil(t, err)
}
func TestParens(t *testing.T) {
t.Skip()
sql := `SELECT t1.Col1,
t2.Col1,
t3.Col1
@@ -128,3 +147,70 @@ func TestParens(t *testing.T) {
assert.Equal(t, "table2", tables[1])
assert.Equal(t, "table3", tables[2])
}
func TestWith(t *testing.T) {
t.Skip()
sql := `WITH
current_month AS (
select
distinct "Month(ISO)" as mth
from A
ORDER BY mth DESC
LIMIT 1
),
last_month_bill AS (
select
CAST (
sum(
CAST(BillableSeries AS INTEGER)
) AS INTEGER
) AS BillableSeries,
"Month(ISO)",
label_namespace
-- , B.activeseries_count
from A
JOIN current_month
ON current_month.mth = A."Month(ISO)"
JOIN B
ON B.namespace = A.label_namespace
GROUP BY
label_namespace,
"Month(ISO)"
ORDER BY BillableSeries DESC
)
SELECT
last_month_bill.*,
BEE.activeseries_count
FROM last_month_bill
JOIN BEE
ON BEE.namespace = last_month_bill.label_namespace`
tables, err := TablesList((sql))
assert.Nil(t, err)
assert.Equal(t, 5, len(tables))
assert.Equal(t, "A", tables[0])
assert.Equal(t, "B", tables[1])
assert.Equal(t, "BEE", tables[2])
}
func TestWithQuote(t *testing.T) {
t.Skip()
sql := "select *,'junk' from foo"
tables, err := TablesList((sql))
assert.Nil(t, err)
assert.Equal(t, "foo", tables[0])
}
func TestWithQuote2(t *testing.T) {
t.Skip()
sql := "SELECT json_serialize_sql('SELECT 1')"
tables, err := TablesList((sql))
assert.Nil(t, err)
assert.Equal(t, 0, len(tables))
}