Add functions for turning restoring on/off and setting license key

These functions improve usability and take all the proper steps to
set restoring on/off (and stop/start background workers in the process)
and to set the license key via a function rather than a guc modification.

sql/CMakeLists.txt

@@ -40,6 +40,7 @@ set(SOURCE_FILES
   gapfill.sql
   maintenance_utils.sql
   partialize_finalize.sql
+  restoring.sql
 )
 
 # These files should be pre-pended to update scripts so that they are

sql/edition.sql

@@ -19,3 +19,19 @@ AS '@MODULE_PATHNAME@', 'ts_print_tsl_license_expiration_info' LANGUAGE C;
 
 CREATE OR REPLACE FUNCTION _timescaledb_internal.license_edition() RETURNS TEXT
 AS '@MODULE_PATHNAME@', 'ts_license_edition' LANGUAGE C;
+
+CREATE OR REPLACE FUNCTION _timescaledb_internal.current_db_set_license_key(new_key TEXT) RETURNS TEXT AS 
+$BODY$
+DECLARE 
+    db text; 
+BEGIN
+    SELECT current_database() INTO db;
+    EXECUTE format('ALTER DATABASE %I SET timescaledb.license_key = %L', db, new_key);
+    EXECUTE format('SET SESSION timescaledb.license_key = %L', new_key);
+    PERFORM _timescaledb_internal.restart_background_workers();
+    RETURN new_key;
+END
+$BODY$ 
+LANGUAGE PLPGSQL;
+
+

sql/restoring.sql

@@ -0,0 +1,32 @@
+-- This file and its contents are licensed under the Apache License 2.0.
+-- Please see the included NOTICE for copyright information and
+-- LICENSE-APACHE for a copy of the license.
+
+CREATE OR REPLACE FUNCTION timescaledb_pre_restore() RETURNS BOOL AS 
+$BODY$
+DECLARE 
+    db text; 
+BEGIN
+    SELECT current_database() INTO db;
+    EXECUTE format($$ALTER DATABASE %I SET timescaledb.restoring ='on'$$, db);
+    SET SESSION timescaledb.restoring = 'on';
+    PERFORM _timescaledb_internal.stop_background_workers();
+    RETURN true;
+END
+$BODY$ 
+LANGUAGE PLPGSQL;
+
+
+CREATE OR REPLACE FUNCTION timescaledb_post_restore() RETURNS BOOL AS 
+$BODY$
+DECLARE 
+    db text; 
+BEGIN
+    SELECT current_database() INTO db;
+    EXECUTE format($$ALTER DATABASE %I SET timescaledb.restoring ='off'$$, db);
+    SET SESSION timescaledb.restoring='off';
+    PERFORM _timescaledb_internal.start_background_workers();
+    RETURN true;
+END
+$BODY$ 
+LANGUAGE PLPGSQL;

test/expected/bgw_launcher.out

@@ -330,6 +330,32 @@ SELECT ((current_setting('server_version_num')::int < 100000) OR wait_greater(:'
  t
 (1 row)
 
+-- Make sure running pre_restore function stops background workers
+SELECT timescaledb_pre_restore();
+ timescaledb_pre_restore 
+-------------------------
+ t
+(1 row)
+
+SELECT wait_worker_counts(1,0,0,0);
+ wait_worker_counts 
+--------------------
+ t
+(1 row)
+
+--And post_restore starts them
+SELECT timescaledb_post_restore();
+ timescaledb_post_restore 
+--------------------------
+ t
+(1 row)
+
+SELECT wait_worker_counts(1,0,1,0);
+ wait_worker_counts 
+--------------------
+ t
+(1 row)
+
 -- Make sure dropping the extension means that the scheduler is stopped
 BEGIN;
 DROP EXTENSION timescaledb;

test/expected/edition.out

@@ -37,3 +37,39 @@ select * from timescaledb_information.license;
  apache  | f       | infinity
 (1 row)
 
+SELECT _timescaledb_internal.current_db_set_license_key('CommunityLicense');
+ current_db_set_license_key 
+----------------------------
+ CommunityLicense
+(1 row)
+
+select * from timescaledb_information.license;
+  edition  | expired | expiration_time 
+-----------+---------+-----------------
+ community | f       | infinity
+(1 row)
+
+create or replace function get_sqlstate(in_text TEXT) RETURNS TEXT AS
+$$
+BEGIN
+    BEGIN
+        EXECUTE in_text;
+    EXCEPTION WHEN others THEN GET STACKED DIAGNOSTICS in_text = RETURNED_SQLSTATE;
+    END;
+    RETURN in_text;
+END;
+$$
+LANGUAGE PLPGSQL;
+SELECT get_sqlstate($$SELECT _timescaledb_internal.current_db_set_license_key('ApacheOnly')$$);
+ get_sqlstate 
+--------------
+ 42501
+(1 row)
+
+select * from timescaledb_information.license;
+  edition  | expired | expiration_time 
+-----------+---------+-----------------
+ community | f       | infinity
+(1 row)
+
+drop function get_sqlstate(TEXT);

test/expected/extension.out

@@ -46,5 +46,7 @@ ORDER BY proname;
  show_tablespaces
  time_bucket
  time_bucket_gapfill
-(32 rows)
+ timescaledb_post_restore
+ timescaledb_pre_restore
+(34 rows)
 

test/expected/pg_dump.out

@@ -262,11 +262,23 @@ SELECT * FROM _timescaledb_catalog.chunk_constraint;
 -- environmental variables that originally called this psql command. Sadly
 -- vars passed to psql do not work in \! commands so we can't do it that way.
 \! utils/pg_dump_aux_dump.sh dump/pg_dump.sql
-ALTER DATABASE :TEST_DBNAME SET timescaledb.restoring='on';
+\c :TEST_DBNAME  
+SET client_min_messages = ERROR;
+CREATE EXTENSION timescaledb CASCADE;
+RESET client_min_messages;
+SELECT timescaledb_pre_restore();
+ timescaledb_pre_restore 
+-------------------------
+ t
+(1 row)
+
+SHOW timescaledb.restoring;
+ timescaledb.restoring 
+-----------------------
+ on
+(1 row)
+
 \! utils/pg_dump_aux_restore.sh dump/pg_dump.sql
-\c :TEST_DBNAME
--- Set to OFF for future DB sessions.
-ALTER DATABASE :TEST_DBNAME SET timescaledb.restoring='off';
 -- Inserting with restoring ON in current session causes tuples to be
 -- inserted on main table, but this should be protected by the insert
 -- blocking trigger.
@@ -275,9 +287,19 @@ INSERT INTO "test_schema"."two_Partitions"("timeCustom", device_id, series_0, se
 VALUES (1357894000000000000, 'dev5', 1.5, 2);
 ERROR:  cannot INSERT into hypertable "two_Partitions" during restore
 \set ON_ERROR_STOP 1
--- Now set to OFF in current session to reenable TimescaleDB. Check
--- INSERTs below.
-SET timescaledb.restoring='off';
+-- Now run our post-restore function.
+SELECT timescaledb_post_restore();
+ timescaledb_post_restore 
+--------------------------
+ t
+(1 row)
+
+SHOW timescaledb.restoring;
+ timescaledb.restoring 
+-----------------------
+ off
+(1 row)
+
 --should be same as count above
 SELECT count(*) = :num_dependent_objects as dependent_objects_match
   FROM pg_depend
@@ -497,3 +519,31 @@ WHERE   refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass AND
  _timescaledb_catalog.tablespace_id_seq
 (8 rows)
 
+        
+-- Make sure we can't run our restoring functions as a normal perm user as that would disable functionality for the whole db
+\c :TEST_DBNAME :ROLE_DEFAULT_PERM_USER
+-- Hides error messages in cases where error messages differ between Postgres versions
+create or replace function get_sqlstate(in_text TEXT) RETURNS TEXT AS
+$$
+BEGIN
+    BEGIN
+        EXECUTE in_text;
+    EXCEPTION WHEN others THEN GET STACKED DIAGNOSTICS in_text = RETURNED_SQLSTATE;
+    END;
+    RETURN in_text;
+END;
+$$
+LANGUAGE PLPGSQL;
+SELECT get_sqlstate('SELECT timescaledb_pre_restore()');
+ get_sqlstate 
+--------------
+ 42501
+(1 row)
+
+SELECT get_sqlstate('SELECT timescaledb_post_restore()');
+ get_sqlstate 
+--------------
+ 42501
+(1 row)
+
+drop function get_sqlstate(TEXT);

test/sql/bgw_launcher.sql

@@ -150,6 +150,13 @@ SELECT wait_worker_counts(1,0,1,0);
 
 SELECT ((current_setting('server_version_num')::int < 100000) OR wait_greater(:'orig_backend_start', :'TEST_DBNAME_2')) as wait_greater;
 
+-- Make sure running pre_restore function stops background workers
+SELECT timescaledb_pre_restore();
+SELECT wait_worker_counts(1,0,0,0);
+--And post_restore starts them
+SELECT timescaledb_post_restore();
+SELECT wait_worker_counts(1,0,1,0);
+
 
 -- Make sure dropping the extension means that the scheduler is stopped
 BEGIN;

test/sql/edition.sql

@@ -22,3 +22,26 @@ SET timescaledb.license_key='ApacheOnly';
 SELECT allow_downgrade_to_apache();
 SET timescaledb.license_key='ApacheOnly';
 select * from timescaledb_information.license;
+
+SELECT _timescaledb_internal.current_db_set_license_key('CommunityLicense');
+select * from timescaledb_information.license;
+
+-- Default perm user shouldn't be able to change the license key.
+\c :TEST_DBNAME :ROLE_DEFAULT_PERM_USER
+
+-- Hides error messages in cases where error messages differe between Postgres versions
+create or replace function get_sqlstate(in_text TEXT) RETURNS TEXT AS
+$$
+BEGIN
+    BEGIN
+        EXECUTE in_text;
+    EXCEPTION WHEN others THEN GET STACKED DIAGNOSTICS in_text = RETURNED_SQLSTATE;
+    END;
+    RETURN in_text;
+END;
+$$
+LANGUAGE PLPGSQL;
+
+SELECT get_sqlstate($$SELECT _timescaledb_internal.current_db_set_license_key('ApacheOnly')$$);
+select * from timescaledb_information.license;
+drop function get_sqlstate(TEXT);

test/sql/pg_dump.sql

@@ -79,12 +79,16 @@ SELECT * FROM _timescaledb_catalog.chunk_constraint;
 -- environmental variables that originally called this psql command. Sadly
 -- vars passed to psql do not work in \! commands so we can't do it that way.
 \! utils/pg_dump_aux_dump.sh dump/pg_dump.sql
-ALTER DATABASE :TEST_DBNAME SET timescaledb.restoring='on';
-\! utils/pg_dump_aux_restore.sh dump/pg_dump.sql
-\c :TEST_DBNAME
 
--- Set to OFF for future DB sessions.
-ALTER DATABASE :TEST_DBNAME SET timescaledb.restoring='off';
+\c :TEST_DBNAME  
+SET client_min_messages = ERROR;
+CREATE EXTENSION timescaledb CASCADE;
+RESET client_min_messages;
+SELECT timescaledb_pre_restore();
+SHOW timescaledb.restoring;
+
+\! utils/pg_dump_aux_restore.sh dump/pg_dump.sql
+
 
 -- Inserting with restoring ON in current session causes tuples to be
 -- inserted on main table, but this should be protected by the insert
@@ -94,9 +98,9 @@ INSERT INTO "test_schema"."two_Partitions"("timeCustom", device_id, series_0, se
 VALUES (1357894000000000000, 'dev5', 1.5, 2);
 \set ON_ERROR_STOP 1
 
--- Now set to OFF in current session to reenable TimescaleDB. Check
--- INSERTs below.
-SET timescaledb.restoring='off';
+-- Now run our post-restore function.
+SELECT timescaledb_post_restore();
+SHOW timescaledb.restoring;
 
 --should be same as count above
 SELECT count(*) = :num_dependent_objects as dependent_objects_match
@@ -142,3 +146,28 @@ WHERE   refclassid = 'pg_catalog.pg_extension'::pg_catalog.regclass AND
         deptype = 'e' AND
         classid='pg_catalog.pg_class'::pg_catalog.regclass
         AND objid NOT IN (select unnest(extconfig) from pg_extension where extname='timescaledb');
+
+        
+
+
+
+
+-- Make sure we can't run our restoring functions as a normal perm user as that would disable functionality for the whole db
+\c :TEST_DBNAME :ROLE_DEFAULT_PERM_USER
+-- Hides error messages in cases where error messages differ between Postgres versions
+create or replace function get_sqlstate(in_text TEXT) RETURNS TEXT AS
+$$
+BEGIN
+    BEGIN
+        EXECUTE in_text;
+    EXCEPTION WHEN others THEN GET STACKED DIAGNOSTICS in_text = RETURNED_SQLSTATE;
+    END;
+    RETURN in_text;
+END;
+$$
+LANGUAGE PLPGSQL;
+
+SELECT get_sqlstate('SELECT timescaledb_pre_restore()');
+SELECT get_sqlstate('SELECT timescaledb_post_restore()');
+
+drop function get_sqlstate(TEXT);