Diskquota stops working after removing any extension.

This happens because during the refactoring in c37cfd3
the diskquota_stop_worker function was incorrectly placed outside
the check that the diskquota extension was being removed.
I corrected this and added tests to check the correctness
of operation for soft and hard restrictions.

src/gp_activetable.c

@@ -191,13 +191,9 @@ 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 (classId == ExtensionRelationId && access == OAT_DROP)
+	if (classId == ExtensionRelationId && access == OAT_DROP && get_extension_oid("diskquota", true) == objectId)
 	{
-		if (get_extension_oid("diskquota", true) == objectId)
-		{
-			invalidate_database_rejectmap(MyDatabaseId);
-		}
-
+		invalidate_database_rejectmap(MyDatabaseId);
 		diskquota_stop_worker();
 		return;
 	}

tests/regress/diskquota_schedule

@@ -31,6 +31,8 @@ test: test_fetch_table_stat
 test: test_appendonly
 test: test_rejectmap
 test: test_clean_rejectmap_after_drop
+test: test_clean_hard_rejectmap_after_drop_other_extension
+test: test_clean_soft_rejectmap_after_drop_other_extension
 test: test_rejectmap_mul_db
 test: test_ctas_pause
 test: test_ctas_role

tests/regress/expected/test_clean_hard_rejectmap_after_drop_other_extension.out

@@ -0,0 +1,69 @@
+CREATE DATABASE test_clean_soft_rejectmap_after_drop_other_extension;
+\c test_clean_soft_rejectmap_after_drop_other_extension
+CREATE EXTENSION diskquota;
+CREATE EXTENSION gp_inject_fault;
+SELECT diskquota.init_table_size_table();
+ init_table_size_table 
+-----------------------
+
+(1 row)
+
+\! gpconfig -c "diskquota.hard_limit" -v "on" > /dev/null
+\! gpstop -u > /dev/null
+CREATE ROLE r;
+SELECT diskquota.set_role_quota('r', '1MB');
+ set_role_quota 
+----------------
+
+(1 row)
+
+CREATE TABLE t(i int) DISTRIBUTED BY (i);
+ALTER TABLE t OWNER TO r;
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+INSERT INTO t SELECT i FROM generate_series(1, 100000000) i; -- fail. Expect that current role is in the rejectmap
+ERROR:  role's disk space quota exceeded with name: 34523  (seg0 127.0.0.1:6002 pid=23690)
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+SELECT diskquota.status();
+              status              
+----------------------------------
+ ("soft limits",on)
+ ("hard limits",on)
+ ("current binary version",2.2.1)
+ ("current schema version",2.2)
+(4 rows)
+
+DROP EXTENSION gp_inject_fault;
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect failed
+ERROR:  role's disk space quota exceeded with name: r
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+SELECT diskquota.status();
+              status              
+----------------------------------
+ ("soft limits",on)
+ ("hard limits",on)
+ ("current binary version",2.2.1)
+ ("current schema version",2.2)
+(4 rows)
+
+DROP EXTENSION diskquota;
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect ok
+\c contrib_regression
+DROP DATABASE test_clean_soft_rejectmap_after_drop_other_extension;
+DROP ROLE r;
+\! gpconfig -c "diskquota.hard_limit" -v "off" > /dev/null
+\! gpstop -u > /dev/null

tests/regress/expected/test_clean_soft_rejectmap_after_drop_other_extension.out

@@ -0,0 +1,55 @@
+CREATE DATABASE test_clean_soft_rejectmap_after_drop_other_extension;
+\c test_clean_soft_rejectmap_after_drop_other_extension
+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) DISTRIBUTED BY (i);
+INSERT INTO t SELECT i FROM generate_series(1, 100000) i; -- ok. diskquota soft limit does not check when first write
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+SELECT diskquota.status();
+              status              
+----------------------------------
+ ("soft limits",on)
+ ("hard limits",off)
+ ("current binary version",2.2.1)
+ ("current schema version",2.2)
+(4 rows)
+
+DROP EXTENSION gp_inject_fault;
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect failed
+ERROR:  schema's disk space quota exceeded with name: public
+SELECT diskquota.wait_for_worker_new_epoch();
+ wait_for_worker_new_epoch 
+---------------------------
+ t
+(1 row)
+
+SELECT diskquota.status();
+              status              
+----------------------------------
+ ("soft limits",on)
+ ("hard limits",off)
+ ("current binary version",2.2.1)
+ ("current schema version",2.2)
+(4 rows)
+
+DROP EXTENSION diskquota;
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect ok
+\c contrib_regression
+DROP DATABASE test_clean_soft_rejectmap_after_drop_other_extension;

tests/regress/sql/test_clean_hard_rejectmap_after_drop_other_extension.sql

@@ -0,0 +1,36 @@
+CREATE DATABASE test_clean_soft_rejectmap_after_drop_other_extension;
+
+\c test_clean_soft_rejectmap_after_drop_other_extension
+CREATE EXTENSION diskquota;
+CREATE EXTENSION gp_inject_fault;
+SELECT diskquota.init_table_size_table();
+
+\! gpconfig -c "diskquota.hard_limit" -v "on" > /dev/null
+\! gpstop -u > /dev/null
+
+CREATE ROLE r;
+SELECT diskquota.set_role_quota('r', '1MB');
+CREATE TABLE t(i int) DISTRIBUTED BY (i);
+ALTER TABLE t OWNER TO r;
+SELECT diskquota.wait_for_worker_new_epoch();
+
+INSERT INTO t SELECT i FROM generate_series(1, 100000000) i; -- fail. Expect that current role is in the rejectmap
+
+SELECT diskquota.wait_for_worker_new_epoch();
+SELECT diskquota.status();
+DROP EXTENSION gp_inject_fault;
+
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect failed
+
+SELECT diskquota.wait_for_worker_new_epoch();
+SELECT diskquota.status();
+DROP EXTENSION diskquota;
+
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect ok
+
+\c contrib_regression
+DROP DATABASE test_clean_soft_rejectmap_after_drop_other_extension;
+DROP ROLE r;
+
+\! gpconfig -c "diskquota.hard_limit" -v "off" > /dev/null
+\! gpstop -u > /dev/null

tests/regress/sql/test_clean_soft_rejectmap_after_drop_other_extension.sql

@@ -0,0 +1,25 @@
+CREATE DATABASE test_clean_soft_rejectmap_after_drop_other_extension;
+
+\c test_clean_soft_rejectmap_after_drop_other_extension
+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) DISTRIBUTED BY (i);
+
+INSERT INTO t SELECT i FROM generate_series(1, 100000) i; -- ok. diskquota soft limit does not check when first write
+
+SELECT diskquota.wait_for_worker_new_epoch();
+SELECT diskquota.status();
+DROP EXTENSION gp_inject_fault;
+
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect failed
+
+SELECT diskquota.wait_for_worker_new_epoch();
+SELECT diskquota.status();
+DROP EXTENSION diskquota;
+
+INSERT INTO t SELECT i FROM generate_series(1, 100) i; -- expect ok
+
+\c contrib_regression
+DROP DATABASE test_clean_soft_rejectmap_after_drop_other_extension;