diff --git a/cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py b/cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py new file mode 100644 index 0000000..dd0f796 --- /dev/null +++ b/cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py @@ -0,0 +1,50 @@ +import typing as t + +from cratedb_sqlparse.generated_parser.SqlBaseParser import SqlBaseParser +from cratedb_sqlparse.generated_parser.SqlBaseParserVisitor import SqlBaseParserVisitor + + +class AstBuilder(SqlBaseParserVisitor): + """ + The class implements the antlr4 visitor pattern similar to how we do it in CrateDB + https://github.com/crate/crate/blob/master/libs/sql-parser/src/main/java/io/crate/sql/parser/AstBuilder.java + + The biggest difference is that in CrateDB, `AstBuilder`, visitor methods + return a specialized Statement visitor. + + Sqlparse just extracts whatever data it needs from the context and injects it to the current + visited statement, enriching its metadata. + """ + + @property + def stmt(self): + if not hasattr(self, "_stmt"): + raise Exception("You should call `enrich` first, that is the entrypoint.") + return self._stmt + + @stmt.setter + def stmt(self, value): + self._stmt = value + + def enrich(self, stmt) -> None: + self.stmt = stmt + self.visit(self.stmt.ctx) + + def visitTableName(self, ctx: SqlBaseParser.TableNameContext): + fqn = ctx.qname() + parts = self.get_text(fqn).replace('"', "").split(".") + + if len(parts) == 1: + name = parts[0] + schema = None + else: + schema, name = parts + + self.stmt.metadata.table_name = name + self.stmt.metadata.schema = schema + + def get_text(self, node) -> t.Optional[str]: + """Gets the text representation of the node or None if it doesn't have one""" + if node: + return node.getText() + return node