| 21 | ) |
| 22 | |
| 23 | func TestInsertSelect(t *testing.T) { |
| 24 | tests := []struct { |
| 25 | name string |
| 26 | input string |
| 27 | wantQuery bool |
| 28 | wantCols int |
| 29 | }{ |
| 30 | {"with columns", "INSERT INTO t1 (a) SELECT a FROM t2", true, 1}, |
| 31 | {"without columns", "INSERT INTO t1 SELECT * FROM t2", true, 0}, |
| 32 | {"multiple columns and WHERE", "INSERT INTO t1 (a, b) SELECT a, b FROM t2 WHERE x > 1", true, 2}, |
| 33 | {"with UNION", "INSERT INTO t1 SELECT a FROM t2 UNION SELECT a FROM t3", true, 0}, |
| 34 | {"VALUES still works", "INSERT INTO t1 VALUES (1)", false, 0}, |
| 35 | {"VALUES with columns", "INSERT INTO t1 (a, b) VALUES (1, 2)", false, 2}, |
| 36 | } |
| 37 | |
| 38 | for _, tt := range tests { |
| 39 | t.Run(tt.name, func(t *testing.T) { |
| 40 | tokens := tokenizeSQL(t, tt.input) |
| 41 | p := NewParser() |
| 42 | result, err := p.Parse(tokens) |
| 43 | if err != nil { |
| 44 | t.Fatalf("Parse(%q) error: %v", tt.input, err) |
| 45 | } |
| 46 | if len(result.Statements) < 1 { |
| 47 | t.Fatalf("expected at least 1 statement, got %d", len(result.Statements)) |
| 48 | } |
| 49 | |
| 50 | // For UNION case, the top-level might be SetOperation |
| 51 | if tt.name == "with UNION" { |
| 52 | // The top-level should be InsertStatement with a SetOperation query |
| 53 | insert, ok := result.Statements[0].(*ast.InsertStatement) |
| 54 | if !ok { |
| 55 | t.Fatalf("expected InsertStatement, got %T", result.Statements[0]) |
| 56 | } |
| 57 | if insert.Query == nil { |
| 58 | t.Fatal("expected Query to be set for UNION case") |
| 59 | } |
| 60 | if _, ok := insert.Query.(*ast.SetOperation); !ok { |
| 61 | t.Fatalf("expected Query to be *SetOperation, got %T", insert.Query) |
| 62 | } |
| 63 | return |
| 64 | } |
| 65 | |
| 66 | insert, ok := result.Statements[0].(*ast.InsertStatement) |
| 67 | if !ok { |
| 68 | t.Fatalf("expected InsertStatement, got %T", result.Statements[0]) |
| 69 | } |
| 70 | |
| 71 | if (insert.Query != nil) != tt.wantQuery { |
| 72 | t.Errorf("Query present = %v, want %v", insert.Query != nil, tt.wantQuery) |
| 73 | } |
| 74 | if len(insert.Columns) != tt.wantCols { |
| 75 | t.Errorf("columns = %d, want %d", len(insert.Columns), tt.wantCols) |
| 76 | } |
| 77 | |
| 78 | // Verify SQL() roundtrip doesn't panic |
| 79 | _ = insert.SQL() |
| 80 | }) |