数据库内核月报 － 2018 / 02

create extension dblink;
create extension oss_fdw;

-- 异步数据装载的准备工作
CREATE OR REPLACE FUNCTION rds_oss_fdw_load_data_prepare(t_from text, t_to text)
  RETURNS bool AS
$BODY$
DECLARE
	t_exist  int;
	curs1 refcursor;
	r	record;
	filepath text;
	fileindex int8;
	s1 text;
	s2 text;
	s3 text;
	c int = 0;
	s4 text;
	s5 text;
	ss4 text;
	ss5 text;
	sql text;
BEGIN
	create table if not exists oss_fdw_load_status(id BIGSERIAL primary key, filename text, size int8, rows int8 default 0, status int default 0);

	select count(*) into t_exist from oss_fdw_load_status;

	if t_exist != 0 then
		RAISE NOTICE 'oss_fdw_load_status not empty';
		return false;
	end if;

	-- 通过 oss_fdw_list_file 函数，把外部表 t_from 匹配的 OSS 中的文件列到表中
	insert into oss_fdw_load_status (filename, size) select name,size from oss_fdw_list_file(t_from) order by size desc;

	select count(*) into t_exist from oss_fdw_load_status;
	if t_exist = 0 then
		RAISE NOTICE 'oss_fdw_load_status empty,not task found';
		return false;
	end if;

	return true;
END;
$BODY$
	LANGUAGE plpgsql;

-- 数据装载的工作函数
CREATE OR REPLACE FUNCTION rds_oss_fdw_load_data_execute(t_from text, t_to text, num_work int, pass text)
  RETURNS bool AS
$BODY$
DECLARE
	t_exist  int;
	curs1 refcursor;
	r	record;
	filepath text;
	fileindex int8;
	s1 text;
	s2 text;
	s3 text;
	c int = 0;
	s4 text;
	s5 text;
	ss4 text;
	ss5 text;
	sql text;
	db text;
	user text;
BEGIN
	select count(*) into t_exist from oss_fdw_load_status;
	if t_exist = 0 then
		RAISE NOTICE 'oss_fdw_load_status empty';
		return false;
	end if;

	s4 = 'oss_loader';
	s5 = 'idle';
	ss4 = '''' || s4 ||'''';
	ss5 = '''' || s5 ||'''';
	sql = 'select count(*) from pg_stat_activity where application_name = ' || ss4 || ' and state != ' || ss5;

	select current_database() into db;
	select current_user into user;

	-- 通过游标，不断获取单个任务
	OPEN curs1 FOR SELECT id, filename FROM oss_fdw_load_status order by id;
	loop
		fetch curs1 into r;
		if not found then
			exit;
		end if;
		fileindex = r.id;
		filepath = r.filename;

		s1 = '''' || t_from ||'''';
		s2 = '''' || t_to ||'''';
		s3 = '''' || filepath ||'''';

		LOOP
			-- 查看当前正在工作的任务数，过达到并发数就在这里等待
			select a into c from dblink('dbname='||db ||' user='||user || ' password='||pass ,sql)as t(a int);
			IF c < num_work THEN
				EXIT;
			END IF;
			RAISE NOTICE 'current runing % loader', c;
			perform pg_sleep(1);
		END LOOP;

		-- 通过 DBLINK 创建异步任务
		perform dis_conn('oss_loader_'||fileindex);
		perform dblink_connect('oss_loader_'||fileindex, 'dbname='||db ||' user='||user || ' application_name=oss_loader' || ' password='||pass);
		perform dblink_send_query('oss_loader_'||fileindex, format('
			begin;
			select rds_oss_fdw_load_single_file(%s,%s,%s,%s);
			end;'
			, fileindex, s1, s2, s3)
		);
		RAISE NOTICE 'runing loader task % filename %',fileindex, filepath;
	end loop;
	close curs1;

	-- 任务分配完成，等待所有任务完成
	LOOP
		select a into c from dblink('dbname='||db ||' user='||user || ' password='||pass ,sql)as t(a int);
		IF c = 0 THEN
			EXIT;
		END IF;
		RAISE NOTICE 'current runing % loader', c;
		perform pg_sleep(1);
	END LOOP;

	return true;
END;
$BODY$
	LANGUAGE plpgsql;

-- 单个文件的数据装在函数
CREATE OR REPLACE FUNCTION rds_oss_fdw_load_single_file(taskid int8, t_from text, t_to text, filepath text)
  RETURNS void AS
$BODY$
DECLARE
	rowscount int8 = 0;
	current text;
	sql text;
BEGIN
	-- 配置 GUC 参数，指定要导入的 OSS 上的文件
	perform set_config('oss_fdw.rds_read_one_file',filepath,true);
	select current_setting('oss_fdw.rds_read_one_file') into current;
	RAISE NOTICE 'begin load %', current;

	-- 通过动态 SQL 导入数据
	EXECUTE 'insert into '|| t_to || ' select * from ' || t_from;
	GET DIAGNOSTICS rowscount = ROW_COUNT;

	-- 导入完成后，把结果保存到状态表中
	RAISE NOTICE 'end load id % % to % % rows', taskid, filepath, t_to, rowscount;
	update oss_fdw_load_status set rows = rowscount,status = 1 where id = taskid;
	return;

EXCEPTION
	when others then
	RAISE 'run rds_oss_fdw_load_single_file with error';
END;
$BODY$
	LANGUAGE plpgsql;

-- 关闭连接不报错
create or replace function dis_conn(name) returns void as $$  
declare  
begin  
  perform dblink_disconnect($1);  
  return;  
exception when others then  
  return;  
end;  
$$ language plpgsql strict;  

select rds_oss_fdw_load_data_prepare('oss_table','lineitem');

 select rds_oss_fdw_load_data_execute('oss_table','lineitem',10,'mypassword');

 select application_name, state, pid,query, now() - xact_start as xact  from pg_stat_activity where state != 'idle' and application_name='oss_loader' order by xact desc;

select pg_terminate_backend(pid),application_name, state ,query from pg_stat_activity where state != 'idle' and pid != pg_backend_pid() and application_name='oss_loader';

select
(
select sum(size)/1024/1024 as complete from oss_fdw_load_status where status = 1
)a,
(
select sum(size)/1024/1024 as full from oss_fdw_load_status
)b;

select rds_oss_fdw_load_data_prepare('t_oss2','lineitem');

select rds_oss_fdw_load_data_execute('t_oss2','lineitem',10,'123456Zwj');

select sum(size)/1024/1024 from oss_fdw_load_status;
      ?column?      
--------------------
 22561.919849395752
(1 row)

select pg_size_pretty(pg_relation_size(oid)) from pg_class where relname = 'lineitem';
 pg_size_pretty 
----------------
 101 GB
(1 row)