Add support for table functions in CREATE TABLE AS statements (#214)

This commit adds support for using table functions like remoteSecure()
and remote() in CREATE TABLE AS statements when column definitions are
present, and fixes support for AS SELECT without parentheses.

The parser now supports table creations such as:
- CREATE TABLE (...) AS remoteSecure('host', 'db', 'table', 'user', 'pass')
- CREATE TABLE (...) AS remote('host', 'db', 'table')
- CREATE TABLE AS SELECT ... (without parentheses)
- CREATE TABLE (...) AS SELECT ... (without parentheses)

Author: swgillespie

parser/ast.go

@@ -1770,18 +1770,19 @@ func (c *CreateDatabase) Accept(visitor ASTVisitor) error {
 }
 
 type CreateTable struct {
-	CreatePos    Pos // position of CREATE|ATTACH keyword
-	StatementEnd Pos
-	OrReplace    bool
-	Name         *TableIdentifier
-	IfNotExists  bool
-	UUID         *UUID
-	OnCluster    *ClusterClause
-	TableSchema  *TableSchemaClause
-	Engine       *EngineExpr
-	SubQuery     *SubQuery
-	HasTemporary bool
-	Comment      *StringLiteral
+	CreatePos      Pos // position of CREATE|ATTACH keyword
+	StatementEnd   Pos
+	OrReplace      bool
+	Name           *TableIdentifier
+	IfNotExists    bool
+	UUID           *UUID
+	OnCluster      *ClusterClause
+	TableSchema    *TableSchemaClause
+	Engine         *EngineExpr
+	SubQuery       *SubQuery
+	TableFunction  *TableFunctionExpr
+	HasTemporary   bool
+	Comment        *StringLiteral
 }
 
 func (c *CreateTable) Pos() Pos {
@@ -1829,6 +1830,10 @@ func (c *CreateTable) String() string {
 		builder.WriteString(" AS ")
 		builder.WriteString(c.SubQuery.String())
 	}
+	if c.TableFunction != nil {
+		builder.WriteString(" AS ")
+		builder.WriteString(c.TableFunction.String())
+	}
 	if c.Comment != nil {
 		builder.WriteString(" COMMENT ")
 		builder.WriteString(c.Comment.String())
@@ -1867,6 +1872,11 @@ func (c *CreateTable) Accept(visitor ASTVisitor) error {
 			return err
 		}
 	}
+	if c.TableFunction != nil {
+		if err := c.TableFunction.Accept(visitor); err != nil {
+			return err
+		}
+	}
 	return visitor.VisitCreateTable(c)
 }
 
@@ -3098,7 +3108,7 @@ func (t *TableSchemaClause) String() string {
 		builder.WriteString(t.AliasTable.String())
 	}
 	if t.TableFunction != nil {
-		builder.WriteByte(' ')
+		builder.WriteString(" AS ")
 		builder.WriteString(t.TableFunction.String())
 	}
 	return builder.String()

parser/parser_table.go

@@ -233,12 +233,39 @@ func (p *Parser) parseCreateTable(pos Pos, orReplace bool) (*CreateTable, error)
 	}
 
 	if p.tryConsumeKeywords(KeywordAs) {
-		subQuery, err := p.parseSubQuery(p.Pos())
-		if err != nil {
-			return nil, err
+		// After AS, we can have: SELECT/WITH (with or without parens), or table_function(...)
+		// Check if it's a SELECT/WITH query (explicitly check keywords/paren before ident)
+		if p.matchKeyword(KeywordSelect) || p.matchKeyword(KeywordWith) || p.matchTokenKind(TokenKindLParen) {
+			// It's a SELECT or WITH query (with or without parentheses)
+			subQuery, err := p.parseSubQuery(p.Pos())
+			if err != nil {
+				return nil, err
+			}
+			createTable.SubQuery = subQuery
+			createTable.StatementEnd = subQuery.End()
+		} else if p.matchTokenKind(TokenKindIdent) {
+			// It's a table function: remote(...), remoteSecure(...), etc.
+			ident, err := p.parseIdent()
+			if err != nil {
+				return nil, err
+			}
+			if p.matchTokenKind(TokenKindLParen) {
+				argsExpr, err := p.parseTableArgList(p.Pos())
+				if err != nil {
+					return nil, err
+				}
+				tableFunc := &TableFunctionExpr{
+					Name: ident,
+					Args: argsExpr,
+				}
+				createTable.TableFunction = tableFunc
+				createTable.StatementEnd = tableFunc.End()
+			} else {
+				return nil, fmt.Errorf("expected ( after identifier in AS clause, got %q", p.lastTokenKind())
+			}
+		} else {
+			return nil, fmt.Errorf("expected SELECT, WITH or identifier after AS, got %q", p.lastTokenKind())
 		}
-		createTable.SubQuery = subQuery
-		createTable.StatementEnd = subQuery.End()
 	}
 
 	comment, err := p.tryParseComment()
@@ -382,51 +409,53 @@ func (p *Parser) parseTableSchemaClause(pos Pos) (*TableSchemaClause, error) {
 			SchemaEnd: rightParenPos,
 			Columns:   columns,
 		}, nil
-	case p.tryConsumeKeywords(KeywordAs):
+	case p.matchKeyword(KeywordAs) && !p.peekKeyword(KeywordSelect) && !p.peekKeyword(KeywordWith) && !p.peekTokenKind(TokenKindLParen):
+		// Handle AS only if followed by identifier (not SELECT/WITH/LPAREN)
+		// This handles: AS ident, AS ident.ident, AS ident(...)
+		// CREATE TABLE will handle: AS SELECT, AS WITH, AS (SELECT ...)
+		p.tryConsumeKeywords(KeywordAs)
+
+		ident, err := p.parseIdent()
+		if err != nil {
+			return nil, err
+		}
 		switch {
-		case p.matchTokenKind(TokenKindIdent):
-			ident, err := p.parseIdent()
+		case p.matchTokenKind(TokenKindDot):
+			// it's a database.table
+			dotIdent, err := p.tryParseDotIdent(p.Pos())
 			if err != nil {
 				return nil, err
 			}
-			switch {
-			case p.matchTokenKind(TokenKindDot):
-				// it's a database.table
-				dotIdent, err := p.tryParseDotIdent(p.Pos())
-				if err != nil {
-					return nil, err
-				}
-				return &TableSchemaClause{
-					SchemaPos: pos,
-					SchemaEnd: dotIdent.End(),
-					AliasTable: &TableIdentifier{
-						Database: ident,
-						Table:    dotIdent,
-					},
-				}, nil
-			case p.matchTokenKind(TokenKindLParen):
-				// it's a table function
-				argsExpr, err := p.parseTableArgList(pos)
-				if err != nil {
-					return nil, err
-				}
-				return &TableSchemaClause{
-					SchemaPos: pos,
-					SchemaEnd: p.End(),
-					TableFunction: &TableFunctionExpr{
-						Name: ident,
-						Args: argsExpr,
-					},
-				}, nil
-			default:
-				return &TableSchemaClause{
-					SchemaPos: pos,
-					SchemaEnd: p.End(),
-					AliasTable: &TableIdentifier{
-						Table: ident,
-					},
-				}, nil
+			return &TableSchemaClause{
+				SchemaPos: pos,
+				SchemaEnd: dotIdent.End(),
+				AliasTable: &TableIdentifier{
+					Database: ident,
+					Table:    dotIdent,
+				},
+			}, nil
+		case p.matchTokenKind(TokenKindLParen):
+			// it's a table function
+			argsExpr, err := p.parseTableArgList(pos)
+			if err != nil {
+				return nil, err
 			}
+			return &TableSchemaClause{
+				SchemaPos: pos,
+				SchemaEnd: p.End(),
+				TableFunction: &TableFunctionExpr{
+					Name: ident,
+					Args: argsExpr,
+				},
+			}, nil
+		default:
+			return &TableSchemaClause{
+				SchemaPos: pos,
+				SchemaEnd: p.End(),
+				AliasTable: &TableIdentifier{
+					Table: ident,
+				},
+			}, nil
 		}
 	}
 	// no schema is ok for MATERIALIZED VIEW

parser/testdata/ddl/create_table_as_remote_function.sql

@@ -0,0 +1,11 @@
+-- CREATE TABLE with columns AS table function (remoteSecure)
+CREATE TABLE test_remote
+(
+    id UInt64,
+    name String,
+    value Int32
+)
+AS remoteSecure('host.example.com', 'source_db', 'source_table', 'user', 'password');
+
+-- Simpler test case with remote()
+CREATE TABLE test_table (id UInt64, name String) AS remote('localhost', 'db', 'source_table');

parser/testdata/ddl/format/create_table_as_remote_function.sql

@@ -0,0 +1,17 @@
+-- Origin SQL:
+-- CREATE TABLE with columns AS table function (remoteSecure)
+CREATE TABLE test_remote
+(
+    id UInt64,
+    name String,
+    value Int32
+)
+AS remoteSecure('host.example.com', 'source_db', 'source_table', 'user', 'password');
+
+-- Simpler test case with remote()
+CREATE TABLE test_table (id UInt64, name String) AS remote('localhost', 'db', 'source_table');
+
+
+-- Format SQL:
+CREATE TABLE test_remote (id UInt64, name String, value Int32) AS remoteSecure('host.example.com', 'source_db', 'source_table', 'user', 'password');
+CREATE TABLE test_table (id UInt64, name String) AS remote('localhost', 'db', 'source_table');

parser/testdata/ddl/output/attach_table_basic.sql.golden.json

@@ -480,6 +480,7 @@
       }
     },
     "SubQuery": null,
+    "TableFunction": null,
     "HasTemporary": false,
     "Comment": null
   }

parser/testdata/ddl/output/create_distributed_table.sql.golden.json

@@ -140,6 +140,7 @@
       "OrderBy": null
     },
     "SubQuery": null,
+    "TableFunction": null,
     "HasTemporary": false,
     "Comment": null
   }

parser/testdata/ddl/output/create_or_replace.sql.golden.json

@@ -317,6 +317,7 @@
       }
     },
     "SubQuery": null,
+    "TableFunction": null,
     "HasTemporary": false,
     "Comment": {
       "LiteralPos": 372,