intruduce rowmark_add_tableoids_walker(), fix rowmark_add_tableoids()

funbringer · funbringer · commit 86f6f8bdf1b4 · 2016-09-07T12:23:05.000+03:00
diff --git a/expected/pg_pathman.out b/expected/pg_pathman.out
@@ -1768,3 +1768,41 @@ NOTICE:  100 rows copied from test_fkey_0
               10
 (1 row)
 
+/* Check rowmarks */
+SELECT create_hash_partitions('test_fkey', 'id', 5);
+ create_hash_partitions 
+------------------------
+                      5
+(1 row)
+
+EXPLAIN (COSTS OFF)
+SELECT * FROM test_fkey
+WHERE id = (SELECT id FROM test_fkey ORDER BY id OFFSET 10 LIMIT 1 FOR UPDATE)
+FOR SHARE;
+                              QUERY PLAN                               
+-----------------------------------------------------------------------
+ LockRows
+   InitPlan 1 (returns $1)
+     ->  Limit
+           ->  LockRows
+                 ->  Sort
+                       Sort Key: test_fkey_0_1.id
+                       ->  Append
+                             ->  Seq Scan on test_fkey_0 test_fkey_0_1
+                             ->  Seq Scan on test_fkey_1 test_fkey_1_1
+                             ->  Seq Scan on test_fkey_2 test_fkey_2_1
+                             ->  Seq Scan on test_fkey_3 test_fkey_3_1
+                             ->  Seq Scan on test_fkey_4 test_fkey_4_1
+   ->  Append
+         ->  Seq Scan on test_fkey_0
+               Filter: (id = $1)
+         ->  Seq Scan on test_fkey_1
+               Filter: (id = $1)
+         ->  Seq Scan on test_fkey_2
+               Filter: (id = $1)
+         ->  Seq Scan on test_fkey_3
+               Filter: (id = $1)
+         ->  Seq Scan on test_fkey_4
+               Filter: (id = $1)
+(23 rows)
+
diff --git a/sql/pg_pathman.sql b/sql/pg_pathman.sql
@@ -665,3 +665,10 @@ SELECT create_hash_partitions('test_fkey', 'id', 10);
 INSERT INTO test_fkey VALUES(1, 'wrong');
 INSERT INTO test_fkey VALUES(1, 'test');
 SELECT drop_partitions('test_fkey');
+
+/* Check rowmarks */
+SELECT create_hash_partitions('test_fkey', 'id', 5);
+EXPLAIN (COSTS OFF)
+SELECT * FROM test_fkey
+WHERE id = (SELECT id FROM test_fkey ORDER BY id OFFSET 10 LIMIT 1 FOR UPDATE)
+FOR SHARE;
diff --git a/src/utils.c b/src/utils.c
@@ -39,6 +39,7 @@ static void change_varnos_in_restrinct_info(RestrictInfo *rinfo,
 static bool change_varno_walker(Node *node, change_varno_context *context);
 static List *get_tableoids_list(List *tlist);
 static void lock_rows_visitor(Plan *plan, void *context);
+static bool rowmark_add_tableoids_walker(Node *node, void *context);
 
 
 /*
@@ -492,6 +493,57 @@ plan_tree_walker(Plan *plan,
 	visitor(plan, context);
 }
 
+static bool
+rowmark_add_tableoids_walker(Node *node, void *context)
+{
+	if (node == NULL)
+		return false;
+
+	if (IsA(node, Query))
+	{
+		Query	   *parse = (Query *) node;
+		ListCell   *lc;
+
+		/* Generate 'tableoid' for partitioned table rowmark */
+		foreach (lc, parse->rowMarks)
+		{
+			RowMarkClause  *rc = (RowMarkClause *) lfirst(lc);
+			Oid				parent = getrelid(rc->rti, parse->rtable);
+			Var			   *var;
+			TargetEntry	   *tle;
+			char			resname[64];
+
+			/* Check that table is partitioned */
+			if (!get_pathman_relation_info(parent))
+				continue;
+
+			var = makeVar(rc->rti,
+						  TableOidAttributeNumber,
+						  OIDOID,
+						  -1,
+						  InvalidOid,
+						  0);
+
+			/* Use parent's Oid as TABLEOID_STR's key (%u) */
+			snprintf(resname, sizeof(resname), TABLEOID_STR("%u"), parent);
+
+			tle = makeTargetEntry((Expr *) var,
+								  list_length(parse->targetList) + 1,
+								  pstrdup(resname),
+								  true);
+
+			/* There's no problem here since new attribute is junk */
+			parse->targetList = lappend(parse->targetList, tle);
+		}
+
+		return query_tree_walker((Query *) node,
+								 rowmark_add_tableoids_walker,
+								 NULL, 0);
+	}
+
+	return expression_tree_walker(node, rowmark_add_tableoids_walker, NULL);
+}
+
 /*
  * Add missing 'TABLEOID_STR%u' junk attributes for inherited partitions
  *
@@ -504,56 +556,7 @@ plan_tree_walker(Plan *plan,
 void
 rowmark_add_tableoids(Query *parse)
 {
-	ListCell *lc;
-
-	check_stack_depth();
-
-	foreach(lc, parse->rtable)
-	{
-		RangeTblEntry *rte = (RangeTblEntry *) lfirst(lc);
-
-		switch(rte->rtekind)
-		{
-			case RTE_SUBQUERY:
-				rowmark_add_tableoids(rte->subquery);
-				break;
-
-			default:
-				break;
-		}
-	}
-
-	/* Generate 'tableoid' for partitioned table rowmark */
-	foreach (lc, parse->rowMarks)
-	{
-		RowMarkClause  *rc = (RowMarkClause *) lfirst(lc);
-		Oid				parent = getrelid(rc->rti, parse->rtable);
-		Var			   *var;
-		TargetEntry	   *tle;
-		char			resname[64];
-
-		/* Check that table is partitioned */
-		if (!get_pathman_relation_info(parent))
-			continue;
-
-		var = makeVar(rc->rti,
-					  TableOidAttributeNumber,
-					  OIDOID,
-					  -1,
-					  InvalidOid,
-					  0);
-
-		/* Use parent's Oid as TABLEOID_STR's key (%u) */
-		snprintf(resname, sizeof(resname), TABLEOID_STR("%u"), parent);
-
-		tle = makeTargetEntry((Expr *) var,
-							  list_length(parse->targetList) + 1,
-							  pstrdup(resname),
-							  true);
-
-		/* There's no problem here since new attribute is junk */
-		parse->targetList = lappend(parse->targetList, tle);
-	}
+	rowmark_add_tableoids_walker((Node *) parse, NULL);
 }
 
 /*
