From c3d12adcf7cdd46c8ca9eb6d45f20c485f27861e Mon Sep 17 00:00:00 2001
From: Sven Klemm <sven@timescale.com>
Date: Tue, 8 Oct 2024 22:27:47 +0200
Subject: [PATCH] Add handling for plpgsql LOOP and EXIT

---
 src/pgspot/visitors.py     |  9 +++++++++
 testdata/expected/loop.out |  6 ++++++
 testdata/loop.sql          | 11 +++++++++++
 3 files changed, 26 insertions(+)
 create mode 100644 testdata/expected/loop.out
 create mode 100644 testdata/loop.sql

diff --git a/src/pgspot/visitors.py b/src/pgspot/visitors.py
index fcc0d1c..dc4d6b6 100644
--- a/src/pgspot/visitors.py
+++ b/src/pgspot/visitors.py
@@ -138,6 +138,12 @@ def visit(self, node):
                                 self.state,
                                 "SELECT " + value["expr"]["PLpgSQL_expr"]["query"],
                             )
+                    case "PLpgSQL_stmt_exit":
+                        if "cond" in value:
+                            visit_sql(
+                                self.state,
+                                "SELECT " + value["cond"]["PLpgSQL_expr"]["query"],
+                            )
                     case "PLpgSQL_stmt_if":
                         if "cond" in value:
                             visit_sql(
@@ -169,6 +175,9 @@ def visit(self, node):
                             )
                         if "body" in value:
                             self.visit(value["body"])
+                    case "PLpgSQL_stmt_loop":
+                        if "body" in value:
+                            self.visit(value["body"])
                     case "PLpgSQL_stmt_raise":
                         if "params" in value:
                             for item in value["params"]:
diff --git a/testdata/expected/loop.out b/testdata/expected/loop.out
new file mode 100644
index 0000000..62c8fdb
--- /dev/null
+++ b/testdata/expected/loop.out
@@ -0,0 +1,6 @@
+PS002: Unsafe function creation: l() at line 1
+PS005: Function without explicit search_path: l() at line 1
+PS016: Unqualified function call: f1 at line 1
+
+ Errors: 0 Warnings: 3 Unknown: 0 
+
diff --git a/testdata/loop.sql b/testdata/loop.sql
new file mode 100644
index 0000000..8e735a2
--- /dev/null
+++ b/testdata/loop.sql
@@ -0,0 +1,11 @@
+create or replace function l() returns void as $func$
+declare
+begin
+ 	loop
+ 	end loop;
+ 	loop
+ 		PERFORM f1('select 1');
+    exit when true;
+ 	end loop;
+end
+$func$ language plpgsql;