--complain if script is sourced in psql, rather than via CREATE EXTENSION
\echo Use "CREATE EXTENSION pgsqlmock" to load this file. \quit

create function _routineresult ( _routine_signature text )
returns text as $$
begin
	/*
	We need this function because without handling errors we will run into an error like
	SQL Error [42883]: ERROR: function
	"pg_toast.mock_func(text,text,text,text,anyelement)" dose not exist
	It happens when we write filter 'routine_return.returns is not null' in tap_funky view.
	Handling 'undefined_object' does not work
	*/
    return pg_catalog.pg_get_function_result((_routine_signature::regprocedure)::oid);
exception when others then
    return null;
end;
$$ language plpgsql stable;

create function _routineargsdefs ( _routine_signature text )
returns text as $$
begin
	/*
	We need this function because without handling errors we will run into an error like
	SQL Error [42883]: ERROR: function
	"pg_toast.mock_func(text,text,text,text,anyelement)" dose not exist
	It happens when we write filter 'routine_args.with_defs is not null' in tap_funky view.
	Handling 'undefined_object' does not work
	*/
    return pg_catalog.pg_get_function_arguments((_routine_signature::regprocedure)::oid);
exception when others then
	return null;
end;
$$ language plpgsql stable;

create view tap_funky_ext as
select
	tf.oid,
	tf.schema,
	tf.name,
	tf.owner,
	tf.args,
	lower(coalesce(
		routine_info."returns",
		tf."returns"))                       AS "returns",
	tf.langoid,
	tf.is_strict,
	tf.kind,
	tf.is_definer,
	tf.returns_set,
	tf.volatility,
	tf.is_visible,
	l.lanname                                AS langname,
	format('(%s)',
		routine_info.args_with_defs
	)                                        AS args_with_defs
from tap_funky tf
join pg_catalog.pg_proc p on p.oid = tf.oid
join pg_catalog.pg_namespace n ON p.pronamespace = n.oid
left join pg_catalog.pg_language l ON l.oid = p.prolang
left join lateral (
	select
		format('%1I.%2I', tf.schema, tf.name) as qualified
) routine_name on true
left join lateral (
	select
		case
			when lower(tf.schema) != 'pg_catalog'
				then _routineresult(concat(routine_name.qualified, '(', tf.args, ')'))
			else null
		end  collate "default" as "returns",
		case
			when lower(tf.schema) != 'pg_catalog'
				then _routineargsdefs(concat(routine_name.qualified, '(', tf.args, ')'))
			else null
		end  collate "default" as "args_with_defs"
) as routine_info on true;


create function get_routine_signature(
	_routine_schema name
	, _routine_name name)
 returns table (routine_schema text, routine_name text, routine_params text)
 language sql stable
as $function$
	select
		"schema", "name", args_with_defs
	from tap_funky_ext
	where
		"schema" = _routine_schema and
		"name" = _routine_name;
$function$;

create function get_routine_signature(
	_routine_name name)
 returns table (routine_schema text, routine_name text, routine_params text)
 language sql stable
as $function$
	select
		"schema", "name", args_with_defs
	from tap_funky_ext
	where
		"name" = _routine_name;
$function$;

--this function creates a mock in place of a real function
create function mock_func(
    _func_schema text
    , _func_name text
    , _func_args text
    , _return_set_value text default null
    , _return_scalar_value anyelement default null::text
)
returns void
--creates a mock in place of a real function
 LANGUAGE plpgsql
AS $function$
declare
    _mock_ddl text;
    _func_result_type text;
    _func_qualified_name text;
    _func_language text;
	_returns_set bool;
	_variants text;
	_ex_msg text;
begin
	--First of all, we have to identify which function we must mock. If there is no such function, throw an error.
	begin
	    select "returns", langname, returns_set
	    into strict _func_result_type, _func_language, _returns_set
	    from tap_funky_ext
	    where "schema" = _func_schema
	        and "name" = _func_name
	        and args_with_defs = _func_args;
		exception when NO_DATA_FOUND or TOO_MANY_ROWS then
			select string_agg(E'\t - ' || format('%I.%I %s', "schema", "name", args_with_defs), E'\n')::text
			into _variants
			from tap_funky_ext
			where "name" = _func_name;
			_ex_msg = format('Routine %I.%I %s does not exist.',
				_func_schema, _func_name, _func_args) || E'\n' || 'Possible variants are:' || E'\n' ||
				coalesce(_variants, 'There is no such function in any schema');
            raise exception '%', coalesce(_ex_msg, 'Нет описания');
	end;
	--This is the case when we need to mock a function written in SQL.
	--But in order to be able to execute the mocking functionality, we need to have a function written in plpgsql.
	--That is why we create a hidden function which name starts with "__".
	if _func_language = 'sql' and _returns_set then
		_mock_ddl = format('
	        create function %1$I.__%2$I(_name text)
	             returns %3$s
	             language plpgsql
	        AS %4$sfunction%4$s
			begin
	            return query execute _query(_name);
			end;
	        %4$sfunction%4$s;',
			_func_schema/*1*/, _func_name/*2*/, _func_result_type/*3*/, '$'/*4*/);
	    execute _mock_ddl;
		_mock_ddl = format('
	        create or replace function %1$I.%2$I %3$s
	             returns %4$s
	             language %5$s
	        AS %7$sfunction%7$s
	            select * from %1$I.__%2$I ( ''%6$s'' );
	        %7$sfunction%7$s;',
			_func_schema/*1*/, _func_name/*2*/, _func_args/*3*/, _func_result_type/*4*/,
			_func_language/*5*/, _return_set_value/*6*/, '$'/*7*/);
	    execute _mock_ddl;
	end if;

	if _func_language = 'plpgsql' and _returns_set then
		_mock_ddl = format('
	        create or replace function %1$I.%2$I %3$s
	             returns %4$s
	             language plpgsql
	        AS %6$sfunction%6$s
			begin
	            return query execute _query( ''%5$s'' );
			end;
	        %6$sfunction%6$s;',
			_func_schema/*1*/, _func_name/*2*/, _func_args/*3*/, _func_result_type/*4*/,
			_return_set_value/*5*/, '$'/*6*/);
	    execute _mock_ddl;
	end if;

	if not _returns_set then
		_mock_ddl = format('
	        create or replace function %1$I.%2$I %3$s
	             RETURNS %4$s
	             LANGUAGE %5$s
	        AS %8$sfunction%8$s
	            select %6$L::%7$s;
	        %8$sfunction%8$s;',
			_func_schema/*1*/,  _func_name/*2*/, _func_args/*3*/, _func_result_type/*4*/,
			_func_language/*5*/, _return_scalar_value/*6*/, pg_typeof(_return_scalar_value)/*7*/, '$'/*8*/);
	    execute _mock_ddl;
	end if;
end $function$;

--This function creates a mock in place of a real view
create function mock_view(
    _view_schema text
    , _view_name text
    , _return_set_sql text default null
)
returns void
--Create a mock in place of a real view
language plpgsql
as $function$
declare
    _mock_ddl text;
    _func_result_type text;
    _func_qualified_name text;
    _func_language text;
    _returns_set bool;
begin
    _mock_ddl = format('drop view %I.%I', _view_schema, _view_name);
    execute _mock_ddl;
    _mock_ddl = format('create view %I.%I as %s', _view_schema, _view_name, _return_set_sql);
    execute _mock_ddl;
end; $function$;

create function fake_table(
    _table_ident text[],
    _make_table_empty boolean default false,
    _leave_primary_key boolean default false,
    _drop_not_null boolean DEFAULT false,
    _drop_collation boolean DEFAULT false,
    _drop_partitions boolean DEFAULT false
)
returns void
--It frees a table from any constraint (we call such a table as a fake)
--faked table is a full copy of _table_name, but has no any constraint
--without foreign and primary things you can do whatever you want in testing context
 LANGUAGE plpgsql
AS $function$
declare
    _table record;
    _fk_table record;
    _part_table record;
    _fake_ddl text;
    _not_null_ddl text;
begin
    for _table in
        select
            quote_ident(coalesce((parse_ident(table_ident))[1], '')) table_schema,
            quote_ident(coalesce((parse_ident(table_ident))[2], '')) table_name,
            coalesce((parse_ident(table_ident))[1], '') table_schema_l,
            coalesce((parse_ident(table_ident))[2], '') table_name_l
        from
            unnest(_table_ident) as t(table_ident)
        loop
            for _fk_table in
                -- collect all table's relations including primary key and unique constraint
                select distinct *
                from (
                    select
                        fk_schema_name table_schema, fk_table_name table_name
						, fk_constraint_name constraint_name, false as is_pk, 1 as ord
                    from
                        pg_all_foreign_keys
                    where
                        fk_schema_name = _table.table_schema_l and fk_table_name = _table.table_name_l
                    union all
                    select
                        fk_schema_name table_schema, fk_table_name table_name
						, fk_constraint_name constraint_name, false as is_pk, 1 as ord
                    from
                        pg_all_foreign_keys
                    where
                        pk_schema_name = _table.table_schema_l and pk_table_name = _table.table_name_l
                    union all
                    select
                        table_schema, table_name
						, constraint_name
						, case when constraint_type = 'PRIMARY KEY' then true else false end as is_pk, 2 as ord
                    from
                        information_schema.table_constraints
                    where
                        table_schema = _table.table_schema_l
                        and table_name = _table.table_name_l
                        and constraint_type in ('PRIMARY KEY', 'UNIQUE')
                ) as t
                order by ord
            loop
				if not(_leave_primary_key and _fk_table.is_pk) then
	                _fake_ddl = format('alter table %1$I.%2$I drop constraint %3$I;',
	                	_fk_table.table_schema/*1*/, _fk_table.table_name/*2*/, _fk_table.constraint_name/*3*/
	                );
	                execute _fake_ddl;
				end if;
            end loop;

            if _make_table_empty then
                _fake_ddl = format('truncate table %1$s.%2$s;', _table.table_schema, _table.table_name);
                execute _fake_ddl;
            end if;

            --Free table from not null constraints
            _fake_ddl = format('alter table %1$s.%2$s ', _table.table_schema, _table.table_name);
            if _drop_not_null then
                select
                    string_agg(format('alter column %1$I drop not null', t.attname), ', ')
                into
                    _not_null_ddl
                from
                    pg_catalog.pg_attribute t
                where t.attrelid = (_table.table_schema || '.' || _table.table_name)::regclass
                    and t.attnum > 0 and attnotnull and attidentity = '' /*'d' or 'a' means generated as identity*/
					--We must be sure that the current column is not part of PK
					and not exists(
						select * from (
							select unnest(_keys(_table.table_schema, _table.table_name, 'p')) as col_name
						) as tt where col_name = t.attname::text);

                _fake_ddl = _fake_ddl || _not_null_ddl || ';';
            else
                _fake_ddl = null;
            end if;

            if _fake_ddl is not null then
                execute _fake_ddl;
            end if;

            if _drop_partitions then
                if pg_version_num() < 100000 then
                    raise exception 'Sorry, but declarative partitioning was introduced only starting with PostgreSQL version 10.';
                end if;
                for _part_table in select _parts from _parts(_table.table_schema, _table.table_name) loop
                    _fake_ddl = 'drop table if exists ' || _part_table._parts || ';';
                    execute _fake_ddl;
                end loop;
            end if;
        end loop;
end $function$;

create function call_count(
	_call_count int
	, _func_schema name
	, _func_name name
	, _func_args name[])
 RETURNS text
 LANGUAGE plpgsql
AS $function$
declare
    _actual_call_count int;
	_track_functions_setting text;
begin
	select current_setting('track_functions') into _track_functions_setting;

	if _track_functions_setting != 'all' then
	    return fail('track_functions setting is not set. Must be all');
	end if;

	select calls into _actual_call_count
	from pg_stat_xact_user_functions
	where funcid = _get_func_oid(_func_schema, _func_name, _func_args);

    return ok(
        _actual_call_count = _call_count
        , format('routine %I.%I must has been called %L times, actual call count is %L'
        	, _func_schema, _func_name, _call_count, _actual_call_count)
    );
end $function$;

create function drop_prepared_statement(_statements text[])
returns setof bool as $$
declare
    _statement record;
begin
	for _statement in select _name from unnest(_statements) as t(_name) loop
	    if exists(select * from pg_prepared_statements where "name" = _statement._name) then
	        EXECUTE format('deallocate %I;', _statement._name);
	        return next true;
		else
			return next false;
	    end if;
	end loop;
end
$$
language plpgsql;


create function print_table_as_json(in _table_schema text, in _table_name text)
returns void
 language plpgsql
AS $function$
declare
    _ddl text;
    _json text;
    _columns text;
--returns a query which you can execute and see your table as normal dataset
--you can find the returned query in the output window in DBeaver, where we see raise notice command output
--note! the returned dataset is limited to 1000 records. that's why you didn't get any jdbc error in dbeaver in case of huge amount of rows
begin
    _ddl = format('
        select json_agg(
            array(select %1$I from %2$I.%1$I limit 1000
        )) as j;', _table_name, _table_schema);
    execute _ddl into _json;
    _json = '[' || ltrim(rtrim(_json::text, ']'), '[') || ']';

    select string_agg(concat(quote_ident(c.column_name), ' ', case when lower(c.data_type) = 'array' then e.data_type || '[]' else c.data_type end), ', ')
    into _columns
    from information_schema."columns" c
    left join information_schema.element_types e
     on ((c.table_catalog, c.table_schema, c.table_name, 'TABLE', c.dtd_identifier)
       = (e.object_catalog, e.object_schema, e.object_name, e.object_type, e.collection_type_identifier))
    where c.table_schema = _table_schema
        and c.table_name = _table_name;

    _json = format('select * from /*%1$I.%2$I*/ json_to_recordset(%3$L) as t(%4$s)',
		_table_schema/*1*/, _table_name/*2*/, _json/*3*/, _columns/*4*/);
    raise notice '%', _json;
end $function$;

create function print_query_as_json(in _prepared_statement_name text)
returns void
 language plpgsql
as $function$
declare
    _ddl text;
	_table_name text;
--returns a query which you can execute and see your table as normal dataset
--you can find the returned query in the output window in DBeaver, where we see raise notice command output
--note! the returned dataset is limited to 1000 records. that's why you didn't get any jdbc error in dbeaver in case of huge amount of rows
begin
	_table_name = _prepared_statement_name || '_' || gen_random_uuid();
    _ddl = format('create table public.%1$I as execute %2$s', _table_name, _prepared_statement_name);
	execute _ddl;
	perform print_table_as_json('public', _table_name::text);
end;
$function$;

create function _get_func_oid(name, name, name[])
 returns oid
 language sql
as $function$
    select oid
      from tap_funky_ext
     where "schema" = $1
	   and "name" = $2
       and args = _funkargs($3)
       and is_visible
$function$;