Fix bug: diskquota stop working after removing any extension (#379) (#25

) Fix the bug caused by #220: After the user removes any extension, the bgworker in the current database will be stopped. Cherry-picked-from: 8210b78 Co-authored-by: Zhang Hao <[email protected]>
arenadata · Oct 2, 2023 · 87b5dc3 · 87b5dc3
1 parent 3b06e37
commit 87b5dc3
Show file tree

Hide file tree

Showing 4 changed files with 60 additions and 3 deletions.
diff --git a/src/gp_activetable.c b/src/gp_activetable.c
@@ -190,15 +190,14 @@ object_access_hook_QuotaStmt(ObjectAccessType access, Oid classId, Oid objectId,
 {
 	if (prev_object_access_hook) (*prev_object_access_hook)(access, classId, objectId, subId, arg);
 
-	// if is 'drop extension diskquota'
+	/* if is 'drop extension diskquota' */
 	if (classId == ExtensionRelationId && access == OAT_DROP)
 	{
 		if (get_extension_oid("diskquota", true) == objectId)
 		{
 			invalidate_database_rejectmap(MyDatabaseId);
+			diskquota_stop_worker();
 		}
-
-		diskquota_stop_worker();
 		return;
 	}
 

diff --git a/tests/regress/diskquota_schedule b/tests/regress/diskquota_schedule
@@ -42,5 +42,6 @@ test: test_tablespace_diff_schema
 test: test_worker_schedule
 test: test_worker_schedule_exception
 test: test_dbname_encoding
+test: test_drop_any_extension
 test: test_drop_extension
 test: reset_config
diff --git a/tests/regress/expected/test_drop_any_extension.out b/tests/regress/expected/test_drop_any_extension.out
@@ -0,0 +1,34 @@
+CREATE DATABASE test_drop_db;
+\c test_drop_db
+CREATE EXTENSION diskquota;
+CREATE EXTENSION gp_inject_fault;
+SELECT diskquota.init_table_size_table();
+ init_table_size_table 
+-----------------------
+
+(1 row)
+
+SELECT diskquota.set_schema_quota(current_schema, '1MB');
+ set_schema_quota 
+------------------
+
+(1 row)
+
+CREATE TABLE t(i int);
+NOTICE:  Table doesn't have 'DISTRIBUTED BY' clause -- Using column named 'i' as the Greenplum Database data distribution key for this table.
+HINT:  The 'DISTRIBUTED BY' clause determines the distribution of data. Make sure column(s) chosen are the optimal data distribution key to minimize skew.
+DROP EXTENSION gp_inject_fault;
+-- expect success
+INSERT INTO t SELECT generate_series(1, 100000);
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+-- expect fail
+INSERT INTO t SELECT generate_series(1, 100000);
+ERROR:  schema's disk space quota exceeded with name: public
+DROP EXTENSION diskquota;
+\c contrib_regression
+DROP DATABASE test_drop_db;
diff --git a/tests/regress/sql/test_drop_any_extension.sql b/tests/regress/sql/test_drop_any_extension.sql
@@ -0,0 +1,23 @@
+CREATE DATABASE test_drop_db;
+
+\c test_drop_db
+
+CREATE EXTENSION diskquota;
+CREATE EXTENSION gp_inject_fault;
+SELECT diskquota.init_table_size_table();
+
+SELECT diskquota.set_schema_quota(current_schema, '1MB');
+CREATE TABLE t(i int);
+
+DROP EXTENSION gp_inject_fault;
+
+-- expect success
+INSERT INTO t SELECT generate_series(1, 100000);
+SELECT diskquota.wait_for_worker_new_epoch();
+-- expect fail
+INSERT INTO t SELECT generate_series(1, 100000);
+
+DROP EXTENSION diskquota;
+
+\c contrib_regression
+DROP DATABASE test_drop_db;