From d8ffd790228587388866dc0ba19c39922ea29252 Mon Sep 17 00:00:00 2001
From: Arjun Lall <lallarjun@gmail.com>
Date: Fri, 20 Dec 2024 11:45:42 -0500
Subject: [PATCH] Add support for WITH CTE remapping

---
 src/query_handler_test.go |  5 +++++
 src/select_remapper.go    | 10 ++++++++++
 2 files changed, 15 insertions(+)

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
index 96df196..bead98b 100644
--- a/src/query_handler_test.go
+++ b/src/query_handler_test.go
@@ -512,6 +512,11 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"gss_authenticated", "encrypted"},
 			"values":      {},
 		},
+		// WITH
+		"WITH RECURSIVE simple_cte AS (SELECT oid, rolname FROM pg_roles WHERE rolname = 'postgres' UNION ALL SELECT oid, rolname FROM pg_roles) SELECT * FROM simple_cte": {
+			"description": {"oid", "rolname"},
+			"values":      {"10", "bemidb"},
+		},
 	}
 
 	for query, responses := range responsesByQuery {
diff --git a/src/select_remapper.go b/src/select_remapper.go
index f449b23..fa22c1c 100644
--- a/src/select_remapper.go
+++ b/src/select_remapper.go
@@ -88,6 +88,16 @@ func (selectRemapper *SelectRemapper) remapSelectStatement(selectStatement *pgQu
 		selectStatement = selectRemapper.remapperWhere.RemapWhereExpressions(selectStatement, selectStatement.WhereClause, indentLevel)
 	}
 
+	// WITH
+	if selectStatement.WithClause != nil {
+		selectRemapper.traceTreeTraversal("WITH CTE's", indentLevel)
+		for _, cte := range selectStatement.WithClause.Ctes {
+			if cteSelect := cte.GetCommonTableExpr().Ctequery.GetSelectStmt(); cteSelect != nil {
+				selectRemapper.remapSelectStatement(cteSelect, indentLevel+1) // recursive
+			}
+		}
+	}
+
 	// FROM
 	if len(selectStatement.FromClause) > 0 {
 		for i, fromNode := range selectStatement.FromClause {