/*
*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0
*/
\pset pager off
CREATE EXTENSION pg_tle;
-- create semi-privileged role to manipulate pg_tle artifacts
CREATE ROLE dbadmin;
GRANT pgtle_admin TO dbadmin;
-- create unprivileged role to create trusted extensions
CREATE ROLE dbstaff;
-- create alt unprivileged role to create trusted extensions
CREATE ROLE dbstaff2;
-- create completely unprivileged role
CREATE ROLE dbguest;
GRANT CREATE, USAGE ON SCHEMA PUBLIC TO pgtle_admin;
GRANT CREATE, USAGE ON SCHEMA PUBLIC TO dbstaff;
DO
$$
  DECLARE
    objname text;
    sql text;
  BEGIN
    SELECT current_database() INTO objname;
    EXECUTE format('GRANT CREATE ON DATABASE %I TO dbstaff;', objname);
  END;
$$ LANGUAGE plpgsql;
-- create function that can be executed by superuser only
CREATE OR REPLACE FUNCTION superuser_only()
RETURNS INT AS $$
(
  SELECT 51
) $$ LANGUAGE sql;
REVOKE EXECUTE ON FUNCTION superuser_only() FROM PUBLIC;
SET search_path TO pgtle,public;
-- installation of artifacts requires semi-privileged role
SET SESSION AUTHORIZATION dbadmin;
SELECT CURRENT_USER;
 current_user 
--------------
 dbadmin
(1 row)

SELECT pgtle.install_extension
(
 'test123',
 '1.0',
 'Test TLE Functions',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test123_func()
  RETURNS INT AS $$
  (
    SELECT 42
  )$$ LANGUAGE sql;
$_pgtle_$
);
 install_extension 
-------------------
 t
(1 row)

-- install a trusted extension that calls functions requiring superuser privilege
SELECT pgtle.install_extension
(
 'test_no_switch_to_superuser_when_trusted',
 '1.0',
 'Test TLE Functions',
$_pgtle_$
  SELECT superuser_only();
$_pgtle_$
);
 install_extension 
-------------------
 t
(1 row)

SET search_path TO public;
-- superuser can create and use extensions that do not require superuser privilege
RESET SESSION AUTHORIZATION;
CREATE EXTENSION test123;
SELECT test123_func();
 test123_func 
--------------
           42
(1 row)

DROP EXTENSION test123;
SET SESSION AUTHORIZATION dbstaff;
SELECT CURRENT_USER;
 current_user 
--------------
 dbstaff
(1 row)

-- unprivileged role can create and use extensions that do not require superuser privilege
CREATE EXTENSION test123;
SELECT test123_func();
 test123_func 
--------------
           42
(1 row)

-- unprivileged role can not create a trusted extension that requires superuser privilege
-- fails
CREATE EXTENSION test_no_switch_to_superuser_when_trusted;
ERROR:  permission denied for function superuser_only
-- switch to dbstaff2
SET SESSION AUTHORIZATION dbstaff2;
SELECT CURRENT_USER;
 current_user 
--------------
 dbstaff2
(1 row)

-- fails
DROP EXTENSION test123;
ERROR:  must be owner of extension test123
-- suceeds
SELECT test123_func();
 test123_func 
--------------
           42
(1 row)

-- fails
DROP FUNCTION test123_func();
ERROR:  must be owner of function test123_func
-- switch to dbguest
SET SESSION AUTHORIZATION dbguest;
SELECT CURRENT_USER;
 current_user 
--------------
 dbguest
(1 row)

-- fails
DROP EXTENSION test123;
ERROR:  must be owner of extension test123
-- suceeds
SELECT test123_func();
 test123_func 
--------------
           42
(1 row)

-- fails
DROP FUNCTION test123_func();
ERROR:  must be owner of function test123_func
SET search_path TO pgtle, public;
-- installation of artifacts requires semi-privileged role
SET SESSION AUTHORIZATION dbadmin;
SELECT CURRENT_USER;
 current_user 
--------------
 dbadmin
(1 row)

SELECT pgtle.install_update_path
(
 'test123',
 '1.0',
 '1.1',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test123_func_2()
  RETURNS INT AS $$
  (
    SELECT 424242
  )$$ LANGUAGE sql;
$_pgtle_$
);
 install_update_path 
---------------------
 t
(1 row)

SET search_path TO public;
-- unprivileged role can modify and use trusted extension
SET SESSION AUTHORIZATION dbstaff;
SELECT CURRENT_USER;
 current_user 
--------------
 dbstaff
(1 row)

ALTER EXTENSION test123 UPDATE TO '1.1';
SELECT test123_func_2();
 test123_func_2 
----------------
         424242
(1 row)

SELECT * FROM pgtle.extension_update_paths('test123');
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 
(2 rows)

SELECT * FROM pgtle.available_extensions() ORDER BY name;
                   name                   | default_version |      comment       
------------------------------------------+-----------------+--------------------
 test123                                  | 1.0             | Test TLE Functions
 test_no_switch_to_superuser_when_trusted | 1.0             | Test TLE Functions
(2 rows)

SELECT * FROM pgtle.available_extension_versions() ORDER BY name;
                   name                   | version | superuser | trusted | relocatable | schema | requires |      comment       
------------------------------------------+---------+-----------+---------+-------------+--------+----------+--------------------
 test123                                  | 1.0     | f         | f       | f           |        | {pg_tle} | Test TLE Functions
 test123                                  | 1.1     | f         | f       | f           |        | {pg_tle} | Test TLE Functions
 test_no_switch_to_superuser_when_trusted | 1.0     | f         | f       | f           |        | {pg_tle} | Test TLE Functions
(3 rows)

DROP EXTENSION test123;
-- negative tests, run as unprivileged role
-- should fail
-- attempt to create a function in pgtle directly
CREATE OR REPLACE FUNCTION pgtle.foo()
RETURNS TEXT AS $$
SELECT 'ok'
$$ LANGUAGE sql;
ERROR:  permission denied for schema pgtle
-- create a function in public and then attempt alter to pg_tle
-- this works
CREATE OR REPLACE FUNCTION public.pg_tlefoo()
RETURNS TEXT AS $$
SELECT 'ok'
$$ LANGUAGE sql;
-- but this should fail
ALTER FUNCTION public.pg_tlefoo() SET SCHEMA pgtle;
ERROR:  pgtle schema reserved for pg_tle functions
-- clean up, should work
DROP FUNCTION public.pg_tlefoo();
-- negative tests, run as superuser
RESET SESSION AUTHORIZATION;
-- attempt to shadow existing file-based extension
-- fail
SELECT pgtle.install_extension
(
 'plpgsql',
 '1.0',
 'Test TLE Functions',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test123_func()
  RETURNS INT AS $$
  (
    SELECT 42
  )$$ LANGUAGE sql;
$_pgtle_$
);
ERROR:  control file already exists for the plpgsql extension
-- attempt to alter a pg_tle extension function
-- fail
ALTER FUNCTION pgtle.install_extension
(
  name text,
  version text,
  description text,
  ext text,
  requires text[]
)
SET search_path TO 'public';
ERROR:  altering pg_tle functions in pgtle schema not allowed
-- test uninstall extensions by a specific version
SELECT pgtle.install_extension
(
 'new_ext',
 '1.0',
 'Test TLE Functions',
$_pgtle_$
  CREATE FUNCTION fun()
  RETURNS INT AS $$ SELECT 1; $$ LANGUAGE SQL;
$_pgtle_$
);
 install_extension 
-------------------
 t
(1 row)

SELECT pgtle.install_update_path
(
 'new_ext',
 '1.0',
 '1.1',
$_pgtle_$
  CREATE OR REPLACE FUNCTION fun()
  RETURNS INT AS $$ SELECT 2; $$ LANGUAGE SQL;
$_pgtle_$
);
 install_update_path 
---------------------
 t
(1 row)

SELECT pgtle.install_update_path
(
 'new_ext',
 '1.1',
 '1.0',
$_pgtle_$
  CREATE OR REPLACE FUNCTION fun()
  RETURNS INT AS $$ SELECT 1; $$ LANGUAGE SQL;
$_pgtle_$
);
 install_update_path 
---------------------
 t
(1 row)

-- check available extension versions -- should be 1.0 and 1.1
SELECT *
FROM pgtle.available_extension_versions() x WHERE x.name = 'new_ext';
  name   | version | superuser | trusted | relocatable | schema | requires |      comment       
---------+---------+-----------+---------+-------------+--------+----------+--------------------
 new_ext | 1.0     | f         | f       | f           |        | {pg_tle} | Test TLE Functions
 new_ext | 1.1     | f         | f       | f           |        | {pg_tle} | Test TLE Functions
(2 rows)

-- check avaialble version update paths -- should be 1.0<=>1.1
SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 1.1--1.0
(2 rows)

SELECT pgtle.uninstall_extension('new_ext', '1.1');
 uninstall_extension 
---------------------
 t
(1 row)

-- check avaialble versions, should only be 1.0
SELECT *
FROM pgtle.available_extension_versions() x WHERE x.name = 'new_ext';
  name   | version | superuser | trusted | relocatable | schema | requires |      comment       
---------+---------+-----------+---------+-------------+--------+----------+--------------------
 new_ext | 1.0     | f         | f       | f           |        | {pg_tle} | Test TLE Functions
(1 row)

-- check avaialble version update paths -- should be none
SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target | path 
--------+--------+------
(0 rows)

-- add the update back in
SELECT pgtle.install_update_path
(
 'new_ext',
 '1.0',
 '1.1',
$_pgtle_$
  CREATE OR REPLACE FUNCTION fun()
  RETURNS INT AS $$ SELECT 2; $$ LANGUAGE SQL;
$_pgtle_$
);
 install_update_path 
---------------------
 t
(1 row)

-- check avaialble version update paths -- should be 1.0=>1.1
SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 
(2 rows)

-- try to add a duplicate update path
-- fail
SELECT pgtle.install_update_path
(
 'new_ext',
 '1.0',
 '1.1',
$_pgtle_$
  CREATE OR REPLACE FUNCTION fun()
  RETURNS INT AS $$ SELECT 2; $$ LANGUAGE SQL;
$_pgtle_$
);
ERROR:  extension "new_ext" update path "1.0-1.1" already installed
HINT:  To update this specific install path, first use "pgtle.uninstall_update_path".
-- add a downgrade path
SELECT pgtle.install_update_path
(
 'new_ext',
 '1.1',
 '1.0',
$_pgtle_$
  CREATE OR REPLACE FUNCTION fun()
  RETURNS INT AS $$ SELECT 1; $$ LANGUAGE SQL;
$_pgtle_$
);
 install_update_path 
---------------------
 t
(1 row)

-- check avaiable update paths
SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 1.1--1.0
(2 rows)

-- only uninstall the downgrade path
SELECT pgtle.uninstall_update_path('new_ext', '1.1', '1.0');
 uninstall_update_path 
-----------------------
 t
(1 row)

-- check avaiable update paths
SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 
(2 rows)

-- try uninstalling again
-- fail
SELECT pgtle.uninstall_update_path('new_ext', '1.1', '1.0');
ERROR:  Extension new_ext does not exist
CONTEXT:  PL/pgSQL function uninstall_update_path(text,text,text) line 16 at RAISE
-- try ininstall with if exists, see false
SELECT pgtle.uninstall_update_path_if_exists('new_ext', '1.1', '1.0');
 uninstall_update_path_if_exists 
---------------------------------
 f
(1 row)

-- check avaiable update paths
SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 
(2 rows)

-- install again and uninstall with "if exists", see true
SELECT pgtle.install_update_path
(
 'new_ext',
 '1.1',
 '1.0',
$_pgtle_$
  CREATE OR REPLACE FUNCTION fun()
  RETURNS INT AS $$ SELECT 1; $$ LANGUAGE SQL;
$_pgtle_$
);
 install_update_path 
---------------------
 t
(1 row)

SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 1.1--1.0
(2 rows)

SELECT pgtle.uninstall_update_path_if_exists('new_ext', '1.1', '1.0');
 uninstall_update_path_if_exists 
---------------------------------
 t
(1 row)

-- check avaiable update paths
SELECT *
FROM pgtle.extension_update_paths('new_ext') x
ORDER BY x.source;
 source | target |   path   
--------+--------+----------
 1.0    | 1.1    | 1.0--1.1
 1.1    | 1.0    | 
(2 rows)

-- ok, install it again
SELECT pgtle.install_update_path
(
 'new_ext',
 '1.1',
 '1.0',
$_pgtle_$
  CREATE OR REPLACE FUNCTION fun()
  RETURNS INT AS $$ SELECT 1; $$ LANGUAGE SQL;
$_pgtle_$
);
 install_update_path 
---------------------
 t
(1 row)

-- test the default version, should be 1.0
SELECT * FROM pgtle.available_extensions() x WHERE x.name = 'new_ext';
  name   | default_version |      comment       
---------+-----------------+--------------------
 new_ext | 1.0             | Test TLE Functions
(1 row)

-- set the new default
SELECT pgtle.set_default_version('new_ext', '1.1');
 set_default_version 
---------------------
 t
(1 row)

-- test the default version, should be 1.1
SELECT * FROM pgtle.available_extensions() x WHERE x.name = 'new_ext';
  name   | default_version |      comment       
---------+-----------------+--------------------
 new_ext | 1.1             | Test TLE Functions
(1 row)

-- try setting a default version that does not exist
-- fail
SELECT pgtle.set_default_version('new_ext', '1.2');
ERROR:  extension and version do not exist
HINT:  Try installing the extension with "pgtle.install_extension".
-- try setting a default version on an extension that does not exist
-- fail
SELECT pgtle.set_default_version('bogus', '1.2');
ERROR:  extension and version do not exist
HINT:  Try installing the extension with "pgtle.install_extension".
-- uninstall
SELECT pgtle.uninstall_extension('new_ext');
 uninstall_extension 
---------------------
 t
(1 row)

-- OK let's try to install an extension with a control file that has errors
SELECT pgtle.install_extension
(
 'broken_ext',
 '0.1',
 $$Distance functions for two points'
 directory = '/tmp/$$,
$_pg_tle_$
    CREATE FUNCTION dist(x1 numeric, y1 numeric, x2 numeric, y2 numeric, l numeric)
    RETURNS numeric
    AS $$
      SELECT ((x2 ^ l - x1 ^ l) ^ (1 / l)) + ((y2 ^ l - y1 ^ l) ^ (1 / l));
    $$ LANGUAGE SQL;
$_pg_tle_$
);
 install_extension 
-------------------
 t
(1 row)

-- this should lead to a sytnax error that we catch
-- fail
SELECT * FROM pgtle.available_extensions();
ERROR:  syntax error in extension control function for "broken_ext"
DETAIL:  Could not parse extension control function "pgtle"."broken_ext.control".
HINT:  You may need to reinstall the extension to correct this error.
-- shoo shoo and uninstall
SELECT pgtle.uninstall_extension('broken_ext');
 uninstall_extension 
---------------------
 t
(1 row)

-- uninstall with a non-existent extension
-- error
SELECT pgtle.uninstall_extension('bogus');
ERROR:  Extension bogus does not exist
CONTEXT:  PL/pgSQL function uninstall_extension(text) line 18 at RAISE
-- uninstall_if_exists with a non-existent extension
-- returns false, no error
SELECT pgtle.uninstall_extension_if_exists('bogus');
 uninstall_extension_if_exists 
-------------------------------
 f
(1 row)

-- uninstall_if_exists with an extension that exists
SELECT pgtle.install_extension
(
 'test42',
 '1.0',
 'Test TLE Functions',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test42_func()
  RETURNS INT AS $$
  (
    SELECT 42
  )$$ LANGUAGE sql;
$_pgtle_$
);
 install_extension 
-------------------
 t
(1 row)

SELECT pgtle.uninstall_extension_if_exists('test42');
 uninstall_extension_if_exists 
-------------------------------
 t
(1 row)

-- back to our regular program: these should work
-- removal of artifacts requires semi-privileged role
SET SESSION AUTHORIZATION dbadmin;
SELECT CURRENT_USER;
 current_user 
--------------
 dbadmin
(1 row)

SELECT pgtle.uninstall_extension('test123');
 uninstall_extension 
---------------------
 t
(1 row)

SELECT pgtle.uninstall_extension('test_no_switch_to_superuser_when_trusted');
 uninstall_extension 
---------------------
 t
(1 row)

RESET SESSION AUTHORIZATION;
SELECT pgtle.install_extension
(
 'test42',
 '1.0',
 'Test TLE Functions',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test42_func()
  RETURNS INT AS $$
  (
    SELECT 42
  )$$ LANGUAGE sql;
$_pgtle_$
);
 install_extension 
-------------------
 t
(1 row)

SELECT pgtle.available_extension_versions();
           available_extension_versions            
---------------------------------------------------
 (test42,1.0,f,f,f,,{pg_tle},"Test TLE Functions")
(1 row)

SELECT pgtle.install_extension_version_sql
(
 'test42',
 '2.0',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test42_func()
  RETURNS INT AS $$
  (
    SELECT 4242
  )$$ LANGUAGE sql;
$_pgtle_$
);
 install_extension_version_sql 
-------------------------------
 t
(1 row)

SELECT pgtle.available_extension_versions();
           available_extension_versions            
---------------------------------------------------
 (test42,1.0,f,f,f,,{pg_tle},"Test TLE Functions")
 (test42,2.0,f,f,f,,{pg_tle},"Test TLE Functions")
(2 rows)

-- install an already installed version
-- fails
SELECT pgtle.install_extension_version_sql
(
 'test42',
 '2.0',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test42_func()
  RETURNS INT AS $$
  (
    SELECT 4242
  )$$ LANGUAGE sql;
$_pgtle_$
);
ERROR:  version "2.0" of extension "test42" already installed
SELECT pgtle.available_extension_versions();
           available_extension_versions            
---------------------------------------------------
 (test42,1.0,f,f,f,,{pg_tle},"Test TLE Functions")
 (test42,2.0,f,f,f,,{pg_tle},"Test TLE Functions")
(2 rows)

-- uninstall default version
-- fails
SELECT pgtle.uninstall_extension('test42', '1.0');
ERROR:  Can not uninstall default version of extension test42, use set_default_version to update the default to another available version and retry
CONTEXT:  PL/pgSQL function uninstall_extension(text,text) line 28 at RAISE
SELECT pgtle.available_extension_versions();
           available_extension_versions            
---------------------------------------------------
 (test42,1.0,f,f,f,,{pg_tle},"Test TLE Functions")
 (test42,2.0,f,f,f,,{pg_tle},"Test TLE Functions")
(2 rows)

-- uninstall a version that is not the default version
-- succeeds
SELECT pgtle.uninstall_extension('test42', '2.0');
 uninstall_extension 
---------------------
 t
(1 row)

SELECT pgtle.available_extension_versions();
           available_extension_versions            
---------------------------------------------------
 (test42,1.0,f,f,f,,{pg_tle},"Test TLE Functions")
(1 row)

-- uninstall non-existent version
-- fails
SELECT pgtle.uninstall_extension('test42', '3.0');
ERROR:  Version test42 of extension 3.0 is not installed and therefore can not be uninstalled
CONTEXT:  PL/pgSQL function uninstall_extension(text,text) line 50 at RAISE
SELECT pgtle.available_extension_versions();
           available_extension_versions            
---------------------------------------------------
 (test42,1.0,f,f,f,,{pg_tle},"Test TLE Functions")
(1 row)

-- uninstall the only installed and  default version
SELECT pgtle.uninstall_extension('test42', '1.0');
 uninstall_extension 
---------------------
 t
(1 row)

SELECT pgtle.available_extension_versions();
 available_extension_versions 
------------------------------
(0 rows)

SELECT pgtle.install_extension
(
 'test42',
 '1.0',
 'Test TLE Functions',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test42_func()
  RETURNS INT AS $$
  (
    SELECT 42
  )$$ LANGUAGE sql;
$_pgtle_$
);
 install_extension 
-------------------
 t
(1 row)

SELECT pgtle.install_extension_version_sql
(
 'test42',
 '2.0',
$_pgtle_$
  CREATE OR REPLACE FUNCTION test42_func()
  RETURNS INT AS $$
  (
    SELECT 4242
  )$$ LANGUAGE sql;
$_pgtle_$
);
 install_extension_version_sql 
-------------------------------
 t
(1 row)

SELECT pgtle.available_extension_versions();
           available_extension_versions            
---------------------------------------------------
 (test42,1.0,f,f,f,,{pg_tle},"Test TLE Functions")
 (test42,2.0,f,f,f,,{pg_tle},"Test TLE Functions")
(2 rows)

-- uninstall extension with multiple versions
SELECT pgtle.uninstall_extension('test42');
 uninstall_extension 
---------------------
 t
(1 row)

SELECT pgtle.available_extension_versions();
 available_extension_versions 
------------------------------
(0 rows)

-- Skip TransactionStmts
BEGIN;
SELECT pgtle.available_extension_versions();
 available_extension_versions 
------------------------------
(0 rows)

SELECT 1/0;
ERROR:  division by zero
SELECT pgtle.available_extension_versions();
ERROR:  current transaction is aborted, commands ignored until end of transaction block
ROLLBACK;
-- clean up
RESET SESSION AUTHORIZATION;
DROP FUNCTION superuser_only();
REVOKE CREATE, USAGE ON SCHEMA PUBLIC FROM dbstaff;
REVOKE CREATE, USAGE ON SCHEMA PUBLIC FROM pgtle_admin;
DROP ROLE dbadmin;
DO
$$
  DECLARE
    objname text;
    sql text;
  BEGIN
    SELECT current_database() INTO objname;
    EXECUTE format('REVOKE ALL ON DATABASE %I FROM dbstaff;', objname);
  END;
$$ LANGUAGE plpgsql;
DROP ROLE dbstaff;
DROP ROLE dbstaff2;
DROP ROLE dbguest;
DROP EXTENSION pg_tle;
DROP SCHEMA pgtle;
DROP ROLE pgtle_admin;