ADBDEV-4266: Diskquota stops working after removing any extension

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;
 	}
 

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

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;

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;