From 3a843c7cfb2c4d4a11d88f6dabc0bcb2804bc768 Mon Sep 17 00:00:00 2001
From: surister <surister98@gmail.com>
Date: Tue, 11 Jun 2024 15:23:02 +0200
Subject: [PATCH] Implement AstBuilder to extract metadata from the query

---
 .../cratedb_sqlparse/AstBuilder.py            | 50 +++++++++++++++++++
 1 file changed, 50 insertions(+)
 create mode 100644 cratedb_sqlparse_py/cratedb_sqlparse/AstBuilder.py

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