From b006580111d0b93136a4226cf1974e28a00cb9b3 Mon Sep 17 00:00:00 2001
From: Pierre Senellart <pierre@senellart.com>
Date: Fri, 26 Jul 2024 16:49:38 +0800
Subject: [PATCH] Fix handling of resjunk TargetEntry (used in GROUP BY) of
 subqueries

---
 src/provsql.c             | 51 ++++++++++++++++++++++++++++++++++-----
 test/expected/resjunk.out | 13 ++++++++++
 test/schedule.common      |  2 +-
 test/sql/resjunk.sql      | 33 +++++++++++++++++++++++++
 4 files changed, 92 insertions(+), 7 deletions(-)
 create mode 100644 test/expected/resjunk.out
 create mode 100644 test/sql/resjunk.sql

diff --git a/src/provsql.c b/src/provsql.c
index 473f7ad..85cd5b5 100644
--- a/src/provsql.c
+++ b/src/provsql.c
@@ -203,6 +203,7 @@ static List *get_provenance_attributes(const constants_t *constants, Query *q)
       if(new_subquery != NULL) {
         int i=0;
         int *offset = (int *)palloc(old_targetlist_length * sizeof(int));
+        unsigned varattnoprovsql;
 
         r->subquery = new_subquery;
 
@@ -229,8 +230,16 @@ static List *get_provenance_attributes(const constants_t *constants, Query *q)
           reduce_varattno_by_offset(q->targetList, rteid, offset);
         }
 
+        varattnoprovsql=0;
+        for(ListCell *cell = list_head(new_subquery->targetList); cell!=NULL; cell=my_lnext(new_subquery->targetList, cell)) {
+          TargetEntry *te = (TargetEntry*) lfirst(cell);
+          ++varattnoprovsql;
+          if(!strcmp(te->resname,PROVSQL_COLUMN_NAME))
+            break;
+        }
+
         r->eref->colnames = lappend(r->eref->colnames, makeString(pstrdup(PROVSQL_COLUMN_NAME)));
-        prov_atts=lappend(prov_atts,make_provenance_attribute(constants, q, r, rteid, list_length(new_subquery->targetList)));
+        prov_atts=lappend(prov_atts,make_provenance_attribute(constants, q, r, rteid, varattnoprovsql));
         fix_type_of_aggregation_result(constants, q, rteid, r->subquery->targetList);
       }
     }
@@ -583,6 +592,7 @@ static Expr *make_aggregation_expression(
     agg->args=list_make1(te_inner);
     agg->aggkind=AGGKIND_NORMAL;
     agg->location=-1;
+    agg->aggno=agg->aggtransno=-1;
 
     agg->aggargtypes = list_make1_oid(constants->OID_TYPE_UUID);
 
@@ -959,11 +969,40 @@ static void add_to_select(
   Query *q,
   Expr *provenance)
 {
-  TargetEntry *te = makeNode(TargetEntry);
-  te->expr = provenance;
-  te->resno = list_length(q->targetList) + 1;
-  te->resname = (char *)PROVSQL_COLUMN_NAME;
-  q->targetList = lappend(q->targetList, te);
+  TargetEntry *newte = makeNode(TargetEntry);
+  bool inserted=false;
+  unsigned resno=0;
+
+  newte->expr = provenance;
+  newte->resname = (char *)PROVSQL_COLUMN_NAME;
+
+  /* Make sure to insert before all resjunk Target Entry */
+  for (ListCell *cell = list_head(q->targetList); cell != NULL;)
+  {
+    TargetEntry *te = (TargetEntry*) lfirst(cell);
+
+    if(!inserted)
+      ++resno;
+
+    if(te->resjunk) {
+      if(!inserted) {
+        newte->resno=resno;
+        q->targetList = list_insert_nth(q->targetList, resno-1, newte);
+        cell=list_nth_cell(q->targetList, resno);
+        te = (TargetEntry*) lfirst(cell);
+        inserted=true;
+      }
+
+      ++te->resno;
+    }
+
+    cell = my_lnext(q->targetList, cell);
+  }
+
+  if(!inserted) {
+    newte->resno = resno+1;
+    q->targetList = lappend(q->targetList, newte);
+  }
 }
 
 typedef struct provenance_mutator_context
diff --git a/test/expected/resjunk.out b/test/expected/resjunk.out
new file mode 100644
index 0000000..599ab57
--- /dev/null
+++ b/test/expected/resjunk.out
@@ -0,0 +1,13 @@
+\set ECHO none
+remove_provenance
+
+(1 row)
+city
+Berlin
+Berlin
+New York
+New York
+Paris
+Paris
+Paris
+(7 rows)
diff --git a/test/schedule.common b/test/schedule.common
index e26d37f..3140b61 100644
--- a/test/schedule.common
+++ b/test/schedule.common
@@ -11,7 +11,7 @@ test: provenance_in_from identify_token subquery create_provenance_mapping
 test: security formula counting
 
 # Test of various ProvSQL features and SQL language capabilities
-test: deterministic union_all union nested_union union_nary distinct group_by_provenance except null unsupported_features no_attribute
+test: deterministic union_all union nested_union union_nary distinct group_by_provenance except null unsupported_features no_attribute resjunk
 test: aggregation
 test: agg_distinct agg_order_by
 test: create_as
diff --git a/test/sql/resjunk.sql b/test/sql/resjunk.sql
new file mode 100644
index 0000000..76d79cf
--- /dev/null
+++ b/test/sql/resjunk.sql
@@ -0,0 +1,33 @@
+\set ECHO none
+\pset format unaligned
+SET search_path TO provsql_test, provsql;
+
+CREATE TABLE resjunk AS
+SELECT
+    city,
+    provenance ()
+FROM (
+    SELECT
+        p1.city
+    FROM
+        personnel p1
+        JOIN personnel p2 ON p1.city = p2.city
+    GROUP BY
+        p1.city,
+        p2.name,
+        p2.id) t;
+
+SELECT
+    remove_provenance ('resjunk');
+
+ALTER TABLE resjunk
+    DROP COLUMN provenance;
+
+SELECT
+    *
+FROM
+    resjunk
+ORDER BY
+    city;
+
+DROP TABLE resjunk;