From 39e44d28a53d76074a45bc528189f4db65974c04 Mon Sep 17 00:00:00 2001
From: Sven Klemm <sven@timescale.com>
Date: Sat, 24 Aug 2024 13:12:44 +0200
Subject: [PATCH] Add some helper code for PLpgSQL AST

pglast returns PLpgSQL AST as a nested dict. This patch add some glue code
to convert the nested dict into a tree of objects to make it more
convenient to work with.
---
 src/pgspot/plpgsql.py | 46 +++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)
 create mode 100644 src/pgspot/plpgsql.py

diff --git a/src/pgspot/plpgsql.py b/src/pgspot/plpgsql.py
new file mode 100644
index 0000000..8c2909d
--- /dev/null
+++ b/src/pgspot/plpgsql.py
@@ -0,0 +1,46 @@
+#
+# pglast returns PLpgSQL AST as a nested dict. This file contains some glue code
+# to convert the nested dict into a tree of objects.
+#
+# Disable pylint warnings about class names cause we are trying to match the
+# AST names used by PLpgSQL parser.
+# pylint: disable-msg=C0103
+# pylint: disable-msg=R0903
+
+
+class PLpgSQLNode:
+    def __init__(self, raw):
+        self.type = list(raw.keys())[0]
+        self.lineno = None
+        for k, v in raw[self.type].items():
+            setattr(self, k, build_node(v))
+
+    def __str__(self):
+        return f"{self.type}({self.lineno})"
+
+    def __repr__(self):
+        fields = self.__dict__.copy()
+        fields.pop("type")
+        return f"{self.type}({fields})"
+
+
+class PLpgSQL_var(PLpgSQLNode):
+    def __init__(self, raw):
+        self.refname = None
+        self.datatype = None
+        super().__init__(raw)
+
+    def __str__(self):
+        return f"Variable(name={self.refname} type={self.datatype})"
+
+
+def build_node(node):
+    if isinstance(node, list):
+        return [build_node(n) for n in node]
+    if isinstance(node, dict):
+        name = list(node.keys())[0]
+        if globals().get(name) is not None:
+            return globals()[name](node)
+        return PLpgSQLNode(node)
+
+    return node