Support schema-qualified table names in INSERT/UPDATE

Author: Copilot

ast.go

@@ -2721,9 +2721,11 @@ type InsertStatement struct {
 	InsertOrIgnore   Pos // position of IGNORE keyword after INSERT OR
 	Into             Pos // position of INTO keyword
 
-	Table *Ident // table name
-	As    Pos    // position of AS keyword
-	Alias *Ident // optional alias
+	Schema *Ident // optional schema name
+	Dot    Pos    // position of DOT between schema and table name
+	Table  *Ident // table name
+	As     Pos    // position of AS keyword
+	Alias  *Ident // optional alias
 
 	ColumnsLparen Pos      // position of column list left paren
 	Columns       []*Ident // optional column list
@@ -2748,6 +2750,7 @@ func (s *InsertStatement) Clone() *InsertStatement {
 	}
 	other := *s
 	other.WithClause = s.WithClause.Clone()
+	other.Schema = s.Schema.Clone()
 	other.Table = s.Table.Clone()
 	other.Alias = s.Alias.Clone()
 	other.Columns = cloneIdents(s.Columns)
@@ -2783,7 +2786,12 @@ func (s *InsertStatement) String() string {
 		}
 	}
 
-	fmt.Fprintf(&buf, " INTO %s", s.Table.String())
+	buf.WriteString(" INTO ")
+	if s.Schema != nil {
+		buf.WriteString(s.Schema.String())
+		buf.WriteString(".")
+	}
+	buf.WriteString(s.Table.String())
 	if s.Alias != nil {
 		fmt.Fprintf(&buf, " AS %s", s.Alias.String())
 	}

parser.go

@@ -1475,10 +1475,18 @@ func (p *Parser) parseInsertStatement(inTrigger bool, withClause *WithClause) (_
 	}
 	stmt.Into, _, _ = p.scan()
 
-	// Parse table name & optional alias.
+	// Parse table name & optional schema & alias.
 	if stmt.Table, err = p.parseIdent("table name"); err != nil {
 		return &stmt, err
 	}
+	// Check for schema.table syntax
+	if p.peek() == DOT {
+		stmt.Schema = stmt.Table
+		stmt.Dot, _, _ = p.scan()
+		if stmt.Table, err = p.parseIdent("table name"); err != nil {
+			return &stmt, err
+		}
+	}
 	if p.peek() == AS {
 		stmt.As, _, _ = p.scan()
 		if stmt.Alias, err = p.parseIdent("alias"); err != nil {
@@ -1731,7 +1739,7 @@ func (p *Parser) parseUpdateStatement(inTrigger bool, withClause *WithClause) (_
 		return nil, p.errorExpected(p.pos, p.tok, "table name")
 	}
 	ident, _ := p.parseIdent("table name")
-	if stmt.Table, err = p.parseQualifiedTableName(ident, false, false, false); err != nil {
+	if stmt.Table, err = p.parseQualifiedTableName(ident, true, false, false); err != nil {
 		return &stmt, err
 	}
 

parser_test.go

@@ -3685,6 +3685,34 @@ func TestParser_ParseStatement(t *testing.T) {
 			},
 		})
 
+		// Test schema-qualified table name
+		AssertParseStatement(t, `INSERT INTO main.tableau VALUES (1, 2), (3, 4)`, &sql.InsertStatement{
+			Insert: pos(0),
+			Into:   pos(7),
+			Schema: &sql.Ident{NamePos: pos(12), Name: "main"},
+			Dot:    pos(16),
+			Table:  &sql.Ident{NamePos: pos(17), Name: "tableau"},
+			Values: pos(25),
+			ValueLists: []*sql.ExprList{
+				{
+					Lparen: pos(32),
+					Exprs: []sql.Expr{
+						&sql.NumberLit{ValuePos: pos(33), Value: "1"},
+						&sql.NumberLit{ValuePos: pos(36), Value: "2"},
+					},
+					Rparen: pos(37),
+				},
+				{
+					Lparen: pos(40),
+					Exprs: []sql.Expr{
+						&sql.NumberLit{ValuePos: pos(41), Value: "3"},
+						&sql.NumberLit{ValuePos: pos(44), Value: "4"},
+					},
+					Rparen: pos(45),
+				},
+			},
+		})
+
 		AssertParseStatementError(t, `INSERT`, `1:6: expected INTO, found 'EOF'`)
 		AssertParseStatementError(t, `INSERT OR`, `1:9: expected ROLLBACK, REPLACE, ABORT, FAIL, or IGNORE, found 'EOF'`)
 		AssertParseStatementError(t, `INSERT INTO`, `1:11: expected table name, found 'EOF'`)
@@ -3856,6 +3884,22 @@ func TestParser_ParseStatement(t *testing.T) {
 			}},
 		})
 
+		// Test schema-qualified table name
+		AssertParseStatement(t, `UPDATE main.tableau SET n2=n1`, &sql.UpdateStatement{
+			Update: pos(0),
+			Table: &sql.QualifiedTableName{
+				Schema: &sql.Ident{NamePos: pos(7), Name: "main"},
+				Dot:    pos(11),
+				Name:   &sql.Ident{NamePos: pos(12), Name: "tableau"},
+			},
+			Set: pos(20),
+			Assignments: []*sql.Assignment{{
+				Columns: []*sql.Ident{{NamePos: pos(24), Name: "n2"}},
+				Eq:      pos(26),
+				Expr:    &sql.Ident{NamePos: pos(27), Name: "n1"},
+			}},
+		})
+
 		AssertParseStatementError(t, `UPDATE`, `1:6: expected table name, found 'EOF'`)
 		AssertParseStatementError(t, `UPDATE OR`, `1:9: expected ROLLBACK, REPLACE, ABORT, FAIL, or IGNORE, found 'EOF'`)
 		AssertParseStatementError(t, `UPDATE tbl`, `1:10: expected SET, found 'EOF'`)