Merge pull request #2558 from dolthub/jennifer/fix

fix type resolution in select of create view stmt

Author: jennifersp

postgres/parser/parser/parse.go

@@ -129,15 +129,15 @@ func (p *Parser) parseOneWithDepth(depth int, sql string) (Statement, error) {
 	return stmts[0], nil
 }
 
-func (p *Parser) scanOneStmt() (sql string, tokens []sqlSymType, done bool) {
+func (p *Parser) scanOneStmt(parse func(sqlStr string, tokens []sqlSymType) error) bool {
 	var lval sqlSymType
-	tokens = p.tokBuf[:0]
+	tokens := p.tokBuf[:0]
 
 	// Scan the first token.
 	for {
 		p.scanner.scan(&lval)
 		if lval.id == 0 {
-			return "", nil, true
+			return true
 		}
 		if lval.id != ';' {
 			break
@@ -150,12 +150,21 @@ func (p *Parser) scanOneStmt() (sql string, tokens []sqlSymType, done bool) {
 	tokens = append(tokens, lval)
 	for {
 		if lval.id == ERROR {
-			return p.scanner.in[startPos:], tokens, true
+			_ = parse(p.scanner.in[startPos:], tokens)
+			return true
 		}
 		posBeforeScan := p.scanner.pos
 		p.scanner.scan(&lval)
 		if lval.id == 0 || lval.id == ';' {
-			return p.scanner.in[startPos:posBeforeScan], tokens, (lval.id == 0)
+			err := parse(p.scanner.in[startPos:posBeforeScan], tokens)
+			if lval.id == 0 || (err != nil && !strings.Contains(err.Error(), "EOF")) {
+				// done scanning all statements OR due to non EOF error
+				return true
+			} else if err == nil {
+				// done scanning single statement
+				return false
+			}
+			// continue scanning if it's EOF error
 		}
 		lval.pos -= startPos
 		tokens = append(tokens, lval)
@@ -166,19 +175,27 @@ func (p *Parser) parseWithDepth(depth int, sql string, nakedIntType *types.T) (S
 	stmts := Statements(p.stmtBuf[:0])
 	p.scanner.init(sql)
 	defer p.scanner.cleanup()
+	var err error
 	for {
-		sql, tokens, done := p.scanOneStmt()
-		stmt, err := p.parse(depth+1, sql, tokens, nakedIntType)
-		if err != nil {
-			return nil, err
-		}
-		if stmt.AST != nil {
-			stmts = append(stmts, stmt)
-		}
+		done := p.scanOneStmt(func(sqlStr string, tokens []sqlSymType) error {
+			var stmt Statement
+			stmt, err = p.parse(depth+1, sqlStr, tokens, nakedIntType)
+			if err != nil {
+				// if it's EOF syntax error, try running again
+				return err
+			}
+			if stmt.AST != nil {
+				stmts = append(stmts, stmt)
+			}
+			return nil
+		})
 		if done {
 			break
 		}
 	}
+	if err != nil {
+		return nil, err
+	}
 	return stmts, nil
 }
 
@@ -259,7 +276,9 @@ func HasMultipleStatements(sql string) bool {
 	defer p.scanner.cleanup()
 	count := 0
 	for {
-		_, _, done := p.scanOneStmt()
+		done := p.scanOneStmt(func(sqlStr string, tokens []sqlSymType) error {
+			return nil
+		})
 		if done {
 			break
 		}

postgres/parser/parser/sql.y

@@ -1153,11 +1153,11 @@ func (u *sqlSymUnion) vacuumTableAndColsList() tree.VacuumTableAndColsList {
 %type <str> table_alias_name constraint_name target_name opt_from_ref_table
 %type <*tree.UnresolvedObjectName> collation_name 
 %type <str> db_object_name_component
-%type <*tree.UnresolvedObjectName> table_name standalone_index_name sequence_name type_name routine_name aggregate_name
+%type <*tree.UnresolvedObjectName> table_name standalone_index_name sequence_name type_name routine_name aggregate_name partition_name
 %type <*tree.UnresolvedObjectName> view_name db_object_name simple_db_object_name complex_db_object_name  opt_collate
 %type <*tree.UnresolvedObjectName> db_object_name_no_keywords simple_db_object_name_no_keywords complex_db_object_name_no_keywords
 %type <[]*tree.UnresolvedObjectName> type_name_list sequence_name_list
-%type <str> schema_name opt_schema_name opt_schema opt_version tablespace_name partition_name
+%type <str> schema_name opt_schema_name opt_schema opt_version tablespace_name
 %type <[]string> schema_name_list role_spec_list opt_role_list opt_owned_by_list
 %type <*tree.UnresolvedName> table_pattern complex_table_pattern
 %type <*tree.UnresolvedName> column_path prefixed_column_path column_path_with_star
@@ -1201,7 +1201,7 @@ func (u *sqlSymUnion) vacuumTableAndColsList() tree.VacuumTableAndColsList {
 %type <[]*tree.Order> sortby_list
 %type <tree.IndexParams> constraint_index_params
 %type <tree.IndexElemList> index_params index_params_name_only opt_index_params_name_only opt_include_index_cols partition_index_params exclude_elems
-%type <tree.NameList> name_list opt_name_list privilege_list
+%type <tree.NameList> name_list opt_name_list privilege_list sconst_as_name_list
 %type <[]int32> opt_array_bounds
 %type <tree.From> from_clause
 %type <tree.TableExprs> from_list opt_from_list
@@ -4611,6 +4611,10 @@ begin_end_block:
   {
     $$.val = &tree.BeginEndBlock{Statements: $3.stmts()}
   }
+| BEGIN ATOMIC RETURN a_expr ';' END
+  {
+    $$.val = &tree.BeginEndBlock{Statements: []tree.Statement{&tree.Return{Expr: $4.expr()}}}
+  }
 
 opt_schema:
   /* EMPTY */
@@ -8726,6 +8730,44 @@ create_trigger_stmt:
       Args: $20.nameList(),
     }
   }
+| CREATE opt_constraint TRIGGER trigger_name trigger_time trigger_events ON table_name opt_from_ref_table
+  opt_trigger_deferrable_mode opt_trigger_relations opt_for_each opt_when EXECUTE function_or_procedure routine_name '(' sconst_as_name_list ')'
+  {
+    $$.val = &tree.CreateTrigger{
+      Replace: false,
+      Constraint: $2.bool(),
+      Name: tree.Name($4),
+      Time: $5.triggerTime(),
+      Events: $6.triggerEvents(),
+      OnTable: $8.unresolvedObjectName().ToTableName(),
+      RefTable: tree.Name($9),
+      Deferrable: $10.triggerDeferrableMode(),
+      Relations: $11.triggerRelations(),
+      ForEachRow: $12.bool(),
+      When: $13.expr(),
+      FuncName: $16.unresolvedObjectName(),
+      Args: $18.nameList(),
+    }
+  }
+| CREATE OR REPLACE opt_constraint TRIGGER trigger_name trigger_time trigger_events ON table_name opt_from_ref_table
+  opt_trigger_deferrable_mode opt_trigger_relations opt_for_each opt_when EXECUTE function_or_procedure routine_name '(' sconst_as_name_list ')'
+  {
+    $$.val = &tree.CreateTrigger{
+      Replace: true,
+      Constraint: $4.bool(),
+      Name: tree.Name($6),
+      Time: $7.triggerTime(),
+      Events: $8.triggerEvents(),
+      OnTable: $10.unresolvedObjectName().ToTableName(),
+      RefTable: tree.Name($11),
+      Deferrable: $12.triggerDeferrableMode(),
+      Relations: $13.triggerRelations(),
+      ForEachRow: $14.bool(),
+      When: $15.expr(),
+      FuncName: $18.unresolvedObjectName(),
+      Args: $20.nameList(),
+    }
+  }
 
 function_or_procedure:
   FUNCTION
@@ -9573,28 +9615,28 @@ alter_table_all_in_tablespace_stmt:
   }
 
 alter_table_parition_stmt:
-  ALTER TABLE table_name ATTACH PARTITION partition_name partition_of
+  ALTER TABLE relation_expr ATTACH PARTITION partition_name partition_of
   {
     $$.val = &tree.AlterTablePartition{
-      Name: $3.unresolvedObjectName(), IfExists: false, Partition: tree.Name($6), Spec: $7.partitionBoundSpec(),
+      Name: $3.unresolvedObjectName(), IfExists: false, Partition: $6.unresolvedObjectName(), Spec: $7.partitionBoundSpec(),
     }
   }
-| ALTER TABLE IF EXISTS table_name ATTACH PARTITION partition_name partition_of
+| ALTER TABLE IF EXISTS relation_expr ATTACH PARTITION partition_name partition_of
   {
     $$.val = &tree.AlterTablePartition{
-      Name: $5.unresolvedObjectName(), IfExists: true, Partition: tree.Name($8), Spec: $9.partitionBoundSpec(),
+      Name: $5.unresolvedObjectName(), IfExists: true, Partition: $8.unresolvedObjectName(), Spec: $9.partitionBoundSpec(),
     }
   }
-| ALTER TABLE table_name DETACH PARTITION partition_name detach_partition_type
+| ALTER TABLE relation_expr DETACH PARTITION partition_name detach_partition_type
   {
     $$.val = &tree.AlterTablePartition{
-      Name: $3.unresolvedObjectName(), IfExists: false, Partition: tree.Name($6), IsDetach: true, DetachType: $7.detachPartition(),
+      Name: $3.unresolvedObjectName(), IfExists: false, Partition: $6.unresolvedObjectName(), IsDetach: true, DetachType: $7.detachPartition(),
     }
   }
-| ALTER TABLE IF EXISTS table_name DETACH PARTITION partition_name detach_partition_type
+| ALTER TABLE IF EXISTS relation_expr DETACH PARTITION partition_name detach_partition_type
   {
     $$.val = &tree.AlterTablePartition{
-      Name: $5.unresolvedObjectName(), IfExists: true, Partition: tree.Name($8), IsDetach: true, DetachType: $9.detachPartition(),
+      Name: $5.unresolvedObjectName(), IfExists: true, Partition: $8.unresolvedObjectName(), IsDetach: true, DetachType: $9.detachPartition(),
     }
   }
 
@@ -14223,6 +14265,16 @@ name_list:
     $$.val = append($1.nameList(), tree.Name($3))
   }
 
+sconst_as_name_list:
+  SCONST
+  {
+    $$.val = tree.NameList{tree.Name($1)}
+  }
+| sconst_as_name_list ',' SCONST
+  {
+    $$.val = append($1.nameList(), tree.Name($3))
+  }
+
 sequence_name_list:
   sequence_name
   {
@@ -14421,7 +14473,7 @@ cursor_name:           name
 
 tablespace_name:       name
 
-partition_name:        name
+partition_name:        db_object_name
 
 routine_name:         db_object_name
 

postgres/parser/sem/tree/alter_table.go

@@ -951,7 +951,7 @@ var _ Statement = &AlterTablePartition{}
 type AlterTablePartition struct {
 	Name       *UnresolvedObjectName
 	IfExists   bool
-	Partition  Name
+	Partition  *UnresolvedObjectName
 	Spec       PartitionBoundSpec
 	IsDetach   bool
 	DetachType DetachPartition
@@ -966,7 +966,7 @@ func (node *AlterTablePartition) Format(ctx *FmtCtx) {
 	node.Name.Format(ctx)
 	if node.IsDetach {
 		ctx.WriteString(" DETACH PARTITION ")
-		ctx.FormatNode(&node.Partition)
+		node.Name.Format(ctx)
 		switch node.DetachType {
 		case DetachPartitionNone:
 		case DetachPartitionConcurrently:
@@ -976,7 +976,7 @@ func (node *AlterTablePartition) Format(ctx *FmtCtx) {
 		}
 	} else {
 		ctx.WriteString(" ATTACH PARTITION ")
-		ctx.FormatNode(&node.Partition)
+		node.Name.Format(ctx)
 		ctx.WriteByte(' ')
 		ctx.FormatNode(&node.Spec)
 	}

postgres/parser/sem/tree/create_function.go

@@ -280,14 +280,12 @@ type BeginEndBlock struct {
 }
 
 func (node *BeginEndBlock) Format(ctx *FmtCtx) {
-	ctx.WriteString("BEGIN ATOMIC")
-	for i, s := range node.Statements {
-		if i != 0 {
-			ctx.WriteString("; ")
-		}
+	ctx.WriteString("BEGIN ATOMIC ")
+	for _, s := range node.Statements {
 		ctx.FormatNode(s)
+		ctx.WriteString("; ")
 	}
-	ctx.WriteString(" END")
+	ctx.WriteString("END")
 }
 
 var _ Statement = &Return{}

server/analyzer/init.go

@@ -55,7 +55,16 @@ const (
 
 // Init adds additional rules to the analyzer to handle Doltgres-specific functionality.
 func Init() {
+	// OnceBeforeDefault runs before AlwaysBeforeDefault in GMS
+	analyzer.OnceBeforeDefault = append([]analyzer.Rule{
+		{Id: ruleId_ResolveType, Apply: ResolveType}, // ResolveType rule must run before simplifyFilters rule in GMS
+		{Id: ruleId_ApplyTablesForAnalyzeAllTables, Apply: applyTablesForAnalyzeAllTables},
+		{Id: ruleId_ConvertDropPrimaryKeyConstraint, Apply: convertDropPrimaryKeyConstraint}},
+		analyzer.OnceBeforeDefault...)
+
 	analyzer.AlwaysBeforeDefault = append(analyzer.AlwaysBeforeDefault,
+		// ResolveType rule must run in this batch in addition to OnceBeforeDefault batch
+		// because of custom batch set optimization in GMS skipping OnceBeforeDefault batch for some nodes.
 		analyzer.Rule{Id: ruleId_ResolveType, Apply: ResolveType},
 		analyzer.Rule{Id: ruleId_TypeSanitizer, Apply: TypeSanitizer},
 		analyzer.Rule{Id: ruleId_ResolveValuesTypes, Apply: ResolveValuesTypes},
@@ -70,11 +79,6 @@ func Init() {
 		analyzer.Rule{Id: ruleId_ResolveProcedureDefaults, Apply: ResolveProcedureDefaults},
 	)
 
-	analyzer.OnceBeforeDefault = append([]analyzer.Rule{
-		{Id: ruleId_ApplyTablesForAnalyzeAllTables, Apply: applyTablesForAnalyzeAllTables},
-		{Id: ruleId_ConvertDropPrimaryKeyConstraint, Apply: convertDropPrimaryKeyConstraint}},
-		analyzer.OnceBeforeDefault...)
-
 	// We remove several validation rules and substitute our own
 	analyzer.OnceBeforeDefault = insertAnalyzerRules(analyzer.OnceBeforeDefault, analyzer.ValidateCreateTableId, true,
 		analyzer.Rule{Id: ruleId_ValidateCreateTable, Apply: validateCreateTable})

server/ast/convert.go

@@ -191,6 +191,8 @@ func Convert(postgresStmt parser.Statement) (vitess.Statement, error) {
 		return nodeReparentDatabase(ctx, stmt)
 	case *tree.Restore:
 		return nodeRestore(ctx, stmt)
+	case *tree.Return:
+		return nodeReturn(ctx, stmt)
 	case *tree.Revoke:
 		return nodeRevoke(ctx, stmt)
 	case *tree.RevokeRole: