MDEV-34757 Check leaf_tables_saved in partition pruning in UPDATE and DELETE

mysql-test/main/ps_mem_leaks.result

@@ -147,7 +147,79 @@ CALL p1(2);
 # Clean up
 DROP TABLE t1;
 DROP PROCEDURE p1;
-# End of 10.11 tests
+#
+# MDEV-34757: assertion of (mem_root->flags & 4) == 0 fails in 2nd ps execution with partition pruning
+#
+CREATE TABLE t1 (id INT, value INT);
+CREATE TABLE t2 (id INT);
+PREPARE stmt FROM 'EXPLAIN UPDATE t1 SET value = (SELECT 1 FROM t2 WHERE id = t1.id)';
+EXECUTE stmt;
+id	select_type	table	type	possible_keys	key	key_len	ref	rows	Extra
+1	SIMPLE	NULL	NULL	NULL	NULL	NULL	NULL	NULL	Impossible WHERE
+INSERT INTO t1 VALUES (1,10),(2,10),(3,10);
+INSERT INTO t2 VALUES (1),(2);
+EXECUTE stmt;
+id	select_type	table	type	possible_keys	key	key_len	ref	rows	Extra
+1	PRIMARY	t1	ALL	NULL	NULL	NULL	NULL	3	
+2	DEPENDENT SUBQUERY	t2	ALL	NULL	NULL	NULL	NULL	2	Using where
+SELECT * FROM t1;
+id	value
+1	10
+2	10
+3	10
+DEALLOCATE PREPARE stmt;
+DROP TABLE t1, t2;
+set @var1=5;
+set @var2=4;
+create table t1 (a int) partition by list(a) (
+partition p0 values in (null,1,2),
+partition p1 values in (3,4)
+);
+create table t2 (a int);
+insert into t1 values (1),(2),(3),(4);
+insert into t2 values (4);
+PREPARE stmt FROM 'UPDATE t1 t1 SET a = (SELECT 1 FROM t2 WHERE a = t1.a) where a = ?';
+execute stmt using @var1;
+execute stmt using @var2;
+deallocate prepare stmt;
+drop table t1, t2;
+set @var1=5;
+set @var2=4;
+create table t1 (a int) partition by list(a) (
+partition p0 values in (null,1,2),
+partition p1 values in (3,4)
+);
+create table t2 (a int);
+insert into t1 values (1),(2),(3),(4);
+insert into t2 values (4);
+PREPARE stmt FROM 'EXPLAIN UPDATE t1 t1 SET a = (SELECT 1 FROM t2 WHERE a = t1.a) where a = ?';
+execute stmt using @var1;
+execute stmt using @var2;
+deallocate prepare stmt;
+drop table t1, t2;
+set @var1=1;
+set @var2=2;
+CREATE TABLE t1 ( id INT(10), value INT(10) );
+CREATE TABLE t2 ( id INT(10) );
+insert into t1 values (1,10),(2,10),(3,10);
+insert into t2 values (1),(2);
+PREPARE stmt FROM 'UPDATE t1 t1 SET value = (SELECT 1 FROM t2 WHERE id = t1.id) WHERE  ?=?';
+execute stmt using @var1, @var2;
+execute stmt using @var1, @var1;
+deallocate prepare stmt;
+DROP TABLE t1,t2;
+set @var1=1;
+set @var2=2;
+CREATE TABLE t1 ( id INT(10), value INT(10) );
+CREATE TABLE t2 ( id INT(10) );
+insert into t1 values (1,10),(2,10),(3,10);
+insert into t2 values (1),(2);
+PREPARE stmt FROM 'EXPLAIN UPDATE t1 t1 SET value = (SELECT 1 FROM t2 WHERE id = t1.id) WHERE  ?=?';
+execute stmt using @var1, @var2;
+execute stmt using @var1, @var1;
+deallocate prepare stmt;
+DROP TABLE t1,t2;
+# End of 11.1 tests
 #
 # MDEV-34649: Memory leaks on running DELETE statement in PS mode with positional parameters
 #

mysql-test/main/ps_mem_leaks.test

@@ -5,6 +5,7 @@
 # The cmake  option -DWITH_PROTECT_STATEMENT_MEMROOT is used only
 # for debug build
 --source include/have_debug.inc
+--source include/have_partition.inc
 
 --echo #
 --echo # MDEV-32369: Memory leak when executing PS for query with IN subquery
@@ -166,7 +167,89 @@ CALL p1(2);
 DROP TABLE t1;
 DROP PROCEDURE p1;
 
---echo # End of 10.11 tests
+--echo #
+--echo # MDEV-34757: assertion of (mem_root->flags & 4) == 0 fails in 2nd ps execution with partition pruning
+--echo #
+# same as the first MDEV-34444 testcase but with explain
+CREATE TABLE t1 (id INT, value INT);
+CREATE TABLE t2 (id INT);
+
+PREPARE stmt FROM 'EXPLAIN UPDATE t1 SET value = (SELECT 1 FROM t2 WHERE id = t1.id)';
+EXECUTE stmt;
+INSERT INTO t1 VALUES (1,10),(2,10),(3,10);
+INSERT INTO t2 VALUES (1),(2);
+EXECUTE stmt;
+SELECT * FROM t1;
+DEALLOCATE PREPARE stmt;
+DROP TABLE t1, t2;
+
+# 2nd ps mem leak; partition pruning
+set @var1=5;
+set @var2=4;
+create table t1 (a int) partition by list(a) (
+  partition p0 values in (null,1,2),
+  partition p1 values in (3,4)
+);
+create table t2 (a int);
+insert into t1 values (1),(2),(3),(4);
+insert into t2 values (4);
+PREPARE stmt FROM 'UPDATE t1 t1 SET a = (SELECT 1 FROM t2 WHERE a = t1.a) where a = ?';
+execute stmt using @var1;
+execute stmt using @var2;
+deallocate prepare stmt;
+drop table t1, t2;
+
+# same but with explain
+set @var1=5;
+set @var2=4;
+create table t1 (a int) partition by list(a) (
+  partition p0 values in (null,1,2),
+  partition p1 values in (3,4)
+);
+create table t2 (a int);
+insert into t1 values (1),(2),(3),(4);
+insert into t2 values (4);
+PREPARE stmt FROM 'EXPLAIN UPDATE t1 t1 SET a = (SELECT 1 FROM t2 WHERE a = t1.a) where a = ?';
+# We are testing 2nd ps assertion failure, not explain output
+--disable_result_log
+execute stmt using @var1;
+execute stmt using @var2;
+--enable_result_log
+deallocate prepare stmt;
+drop table t1, t2;
+
+# top level impossible where
+set @var1=1;
+set @var2=2;
+CREATE TABLE t1 ( id INT(10), value INT(10) );
+CREATE TABLE t2 ( id INT(10) );
+insert into t1 values (1,10),(2,10),(3,10);
+insert into t2 values (1),(2);
+
+PREPARE stmt FROM 'UPDATE t1 t1 SET value = (SELECT 1 FROM t2 WHERE id = t1.id) WHERE  ?=?';
+execute stmt using @var1, @var2;
+execute stmt using @var1, @var1;
+deallocate prepare stmt;
+DROP TABLE t1,t2;
+
+# top level impossible where, with explain
+set @var1=1;
+set @var2=2;
+CREATE TABLE t1 ( id INT(10), value INT(10) );
+CREATE TABLE t2 ( id INT(10) );
+insert into t1 values (1,10),(2,10),(3,10);
+insert into t2 values (1),(2);
+
+PREPARE stmt FROM 'EXPLAIN UPDATE t1 t1 SET value = (SELECT 1 FROM t2 WHERE id = t1.id) WHERE  ?=?';
+# We are testing 2nd ps assertion failure, not explain output
+--disable_result_log
+execute stmt using @var1, @var2;
+execute stmt using @var1, @var1;
+--enable_result_log
+deallocate prepare stmt;
+DROP TABLE t1,t2;
+
+--echo # End of 11.1 tests
 
 --echo #
 --echo # MDEV-34649: Memory leaks on running DELETE statement in PS mode with positional parameters

sql/sql_delete.cc

@@ -494,6 +494,12 @@ bool Sql_cmd_delete::delete_from_single_table(THD *thd)
     if (thd->binlog_for_noop_dml(transactional_table))
       DBUG_RETURN(1);
 
+    if (!thd->lex->current_select->leaf_tables_saved)
+    {
+      thd->lex->current_select->save_leaf_tables(thd);
+      thd->lex->current_select->leaf_tables_saved= true;
+    }
+
     my_ok(thd, 0);
     DBUG_RETURN(0);
   }

sql/sql_update.cc

@@ -463,6 +463,12 @@ bool Sql_cmd_update::update_single_table(THD *thd)
     if (thd->binlog_for_noop_dml(transactional_table))
       DBUG_RETURN(1);
 
+    if (!thd->lex->current_select->leaf_tables_saved)
+    {
+      thd->lex->current_select->save_leaf_tables(thd);
+      thd->lex->current_select->leaf_tables_saved= true;
+    }
+
     my_ok(thd);				// No matching records
     DBUG_RETURN(0);
   }