Take only comment inside CREATE query

Author: oxisto

internal/engine/sqlite/convert.go

@@ -112,7 +112,7 @@ func (c *cc) convertCreate_table_stmtContext(n *parser.Create_table_stmtContext)
 	stmt := &ast.CreateTableStmt{
 		Name:        parseTableName(n),
 		IfNotExists: n.EXISTS_() != nil,
-		Comment:     comment(tokenStream, n, false),
+		Comment:     comment(tokenStream, n.OPEN_PAR().GetSymbol()),
 	}
 
 	for _, idef := range n.AllColumn_def() {
@@ -126,27 +126,22 @@ func (c *cc) convertCreate_table_stmtContext(n *parser.Create_table_stmtContext)
 				Colname:   identifier(def.Column_name().GetText()),
 				IsNotNull: hasNotNullConstraint(def.AllColumn_constraint()),
 				TypeName:  &ast.TypeName{Name: typeName},
-				Comment:   comment(tokenStream, def, true),
+				Comment:   comment(tokenStream, def.GetStop()),
 			})
 
 		}
 	}
 	return stmt
 }
 
-// comment returns the comment associated with the given context. The parameter right indicates whether the comment is
-// to the right or to the left of the context.
-func comment(tokenStream *antlr.CommonTokenStream, ctx antlr.ParserRuleContext, right bool) string {
+// comment returns the comment associated with the given context.
+func comment(tokenStream *antlr.CommonTokenStream, from antlr.Token) string {
 	var (
 		hiddenTokens []antlr.Token
 		comment      string
 	)
 
-	if right {
-		hiddenTokens = tokenStream.GetHiddenTokensToRight(ctx.GetStop().GetTokenIndex()+1, antlr.TokenHiddenChannel)
-	} else {
-		hiddenTokens = tokenStream.GetHiddenTokensToLeft(ctx.GetStart().GetTokenIndex(), antlr.TokenHiddenChannel)
-	}
+	hiddenTokens = tokenStream.GetHiddenTokensToRight(from.GetTokenIndex()+1, antlr.TokenHiddenChannel)
 
 	for _, token := range hiddenTokens {
 		// Filter for single-line comments

internal/engine/sqlite/parse.go

@@ -50,7 +50,6 @@ func (p *Parser) Parse(r io.Reader) ([]ast.Statement, error) {
 	pp.AddErrorListener(el)
 	// pp.BuildParseTrees = true
 	tree := pp.Parse()
-
 	if el.err != "" {
 		return nil, errors.New(el.err)
 	}
@@ -84,7 +83,6 @@ func (p *Parser) Parse(r io.Reader) ([]ast.Statement, error) {
 			loc = stmt.GetStop().GetStop() + 2
 		}
 	}
-
 	return stmts, nil
 }