From 27965f29fdabc2f0ac6759fb033a466b343f8626 Mon Sep 17 00:00:00 2001
From: Sven Klemm <sven@timescale.com>
Date: Mon, 6 Jan 2025 16:18:14 +0100
Subject: [PATCH] Fix RETURN QUERY EXECUTE handling

Fixes #211
---
 src/pgspot/visitors.py             | 18 ++++++++++++++++++
 testdata/expected/return_query.out |  5 +++++
 testdata/return_query.sql          | 12 ++++++++++++
 3 files changed, 35 insertions(+)
 create mode 100644 testdata/expected/return_query.out
 create mode 100644 testdata/return_query.sql

diff --git a/src/pgspot/visitors.py b/src/pgspot/visitors.py
index dc4d6b6..b72c37d 100644
--- a/src/pgspot/visitors.py
+++ b/src/pgspot/visitors.py
@@ -191,6 +191,24 @@ def visit(self, node):
                                 self.state,
                                 "SELECT " + value["expr"]["PLpgSQL_expr"]["query"],
                             )
+                    case "PLpgSQL_stmt_return_query":
+                        if "dynquery" in value:
+                            query = (
+                                "SELECT " + value["dynquery"]["PLpgSQL_expr"]["query"]
+                            )
+                            parsed = parse_sql(query)[0].stmt.targetList[0].val
+
+                            # When the query is a string literal we can analyze it's content
+                            if (
+                                isinstance(parsed, ast.A_Const)
+                                and parsed.isnull is False
+                                and isinstance(parsed.val, ast.String)
+                            ):
+                                visit_sql(
+                                    self.state,
+                                    parsed.val.sval,
+                                )
+
                     case "PLpgSQL_stmt_while":
                         if "cond" in value:
                             visit_sql(
diff --git a/testdata/expected/return_query.out b/testdata/expected/return_query.out
new file mode 100644
index 0000000..9f364a6
--- /dev/null
+++ b/testdata/expected/return_query.out
@@ -0,0 +1,5 @@
+PS005: Function without explicit search_path: ret_query(y integer) at line 2
+PS016: Unqualified function call: format at line 2
+
+ Errors: 0 Warnings: 2 Unknown: 0 
+
diff --git a/testdata/return_query.sql b/testdata/return_query.sql
new file mode 100644
index 0000000..cc1c1e9
--- /dev/null
+++ b/testdata/return_query.sql
@@ -0,0 +1,12 @@
+
+create function ret_query() returns table(y int)
+as $func$
+declare
+  a text;
+begin
+    return query execute $$select format('abc',1,2)::int$$;
+    return query execute a;
+    return query execute a + 'def';
+end;
+$func$ language plpgsql;
+