From 76c0c1903442ef41bd4a6e15f5b1672d4464c90c Mon Sep 17 00:00:00 2001
From: amit <amitlangote09@gmail.com>
Date: Tue, 24 Jan 2017 13:22:17 +0900
Subject: [PATCH] Add relkind checks to certain contrib modules

Contrib functions such a pg_visibility, pg_relpages are only applicable
to certain relkinds; error out if otherwise.

Also, RELKIND_PARTITIONED_TABLE was not added to the pageinspect's and
pgstattuple's list of unsupported relkinds.

Add tests for the newly added checks.
---
 contrib/pageinspect/expected/page.out            |   5 ++
 contrib/pageinspect/rawpage.c                    |   5 ++
 contrib/pageinspect/sql/page.sql                 |   5 ++
 contrib/pg_visibility/.gitignore                 |   4 +
 contrib/pg_visibility/Makefile                   |   2 +
 contrib/pg_visibility/expected/pg_visibility.out | 103 +++++++++++++++++++++++
 contrib/pg_visibility/pg_visibility.c            |  44 +++++++---
 contrib/pg_visibility/sql/pg_visibility.sql      |  66 +++++++++++++++
 contrib/pgstattuple/expected/pgstattuple.out     |  29 +++++++
 contrib/pgstattuple/pgstatindex.c                |  30 +++++++
 contrib/pgstattuple/pgstattuple.c                |   3 +
 contrib/pgstattuple/sql/pgstattuple.sql          |  24 ++++++
 12 files changed, 306 insertions(+), 14 deletions(-)
 create mode 100644 contrib/pg_visibility/.gitignore
 create mode 100644 contrib/pg_visibility/expected/pg_visibility.out
 create mode 100644 contrib/pg_visibility/sql/pg_visibility.sql

diff --git a/contrib/pageinspect/expected/page.out b/contrib/pageinspect/expected/page.out
index 13964cd878..85b183db4d 100644
--- a/contrib/pageinspect/expected/page.out
+++ b/contrib/pageinspect/expected/page.out
@@ -71,3 +71,8 @@ SELECT * FROM fsm_page_contents(get_raw_page('test1', 'fsm', 0));
 (1 row)
 
 DROP TABLE test1;
+-- check that using any of these functions with a partitioned table would fail
+create table test_partitioned (a int) partition by range (a);
+select get_raw_page('test_partitioned', 0);
+ERROR:  cannot get raw page from partitioned table "test_partitioned"
+drop table test_partitioned;
diff --git a/contrib/pageinspect/rawpage.c b/contrib/pageinspect/rawpage.c
index 102f360c6f..1ccc3ff320 100644
--- a/contrib/pageinspect/rawpage.c
+++ b/contrib/pageinspect/rawpage.c
@@ -123,6 +123,11 @@ get_raw_page_internal(text *relname, ForkNumber forknum, BlockNumber blkno)
 				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
 				 errmsg("cannot get raw page from foreign table \"%s\"",
 						RelationGetRelationName(rel))));
+	if (rel->rd_rel->relkind == RELKIND_PARTITIONED_TABLE)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("cannot get raw page from partitioned table \"%s\"",
+						RelationGetRelationName(rel))));
 
 	/*
 	 * Reject attempts to read non-local temporary relations; we would be
diff --git a/contrib/pageinspect/sql/page.sql b/contrib/pageinspect/sql/page.sql
index 97eef9829a..fa3533f2af 100644
--- a/contrib/pageinspect/sql/page.sql
+++ b/contrib/pageinspect/sql/page.sql
@@ -28,3 +28,8 @@ SELECT tuple_data_split('test1'::regclass, t_data, t_infomask, t_infomask2, t_bi
 SELECT * FROM fsm_page_contents(get_raw_page('test1', 'fsm', 0));
 
 DROP TABLE test1;
+
+-- check that using any of these functions with a partitioned table would fail
+create table test_partitioned (a int) partition by range (a);
+select get_raw_page('test_partitioned', 0);
+drop table test_partitioned;
diff --git a/contrib/pg_visibility/.gitignore b/contrib/pg_visibility/.gitignore
new file mode 100644
index 0000000000..5dcb3ff972
--- /dev/null
+++ b/contrib/pg_visibility/.gitignore
@@ -0,0 +1,4 @@
+# Generated subdirectories
+/log/
+/results/
+/tmp_check/
diff --git a/contrib/pg_visibility/Makefile b/contrib/pg_visibility/Makefile
index 379591a098..bc42944426 100644
--- a/contrib/pg_visibility/Makefile
+++ b/contrib/pg_visibility/Makefile
@@ -7,6 +7,8 @@ EXTENSION = pg_visibility
 DATA = pg_visibility--1.1.sql pg_visibility--1.0--1.1.sql
 PGFILEDESC = "pg_visibility - page visibility information"
 
+REGRESS = pg_visibility
+
 ifdef USE_PGXS
 PG_CONFIG = pg_config
 PGXS := $(shell $(PG_CONFIG) --pgxs)
diff --git a/contrib/pg_visibility/expected/pg_visibility.out b/contrib/pg_visibility/expected/pg_visibility.out
new file mode 100644
index 0000000000..4496fe36f8
--- /dev/null
+++ b/contrib/pg_visibility/expected/pg_visibility.out
@@ -0,0 +1,103 @@
+CREATE EXTENSION pg_visibility;
+--
+-- check that using the module's functions with unsupported relations will fail
+--
+create table test_partitioned (a int) partition by list (a);
+select pg_visibility('test_partitioned', 0);
+ERROR:  "test_partitioned" is not a table, materialized view, or TOAST table
+select pg_visibility_map('test_partitioned');
+ERROR:  "test_partitioned" is not a table, materialized view, or TOAST table
+select pg_visibility_map_summary('test_partitioned');
+ERROR:  "test_partitioned" is not a table, materialized view, or TOAST table
+select pg_check_frozen('test_partitioned');
+ERROR:  "test_partitioned" is not a table, materialized view, or TOAST table
+select pg_truncate_visibility_map('test_partitioned');
+ERROR:  "test_partitioned" is not a table, materialized view, or TOAST table
+create table test_partition partition of test_partitioned for values in (1);
+create index test_index on test_partition (a);
+select pg_visibility('test_index', 0);
+ERROR:  "test_index" is not a table, materialized view, or TOAST table
+select pg_visibility_map('test_index');
+ERROR:  "test_index" is not a table, materialized view, or TOAST table
+select pg_visibility_map_summary('test_index');
+ERROR:  "test_index" is not a table, materialized view, or TOAST table
+select pg_check_frozen('test_index');
+ERROR:  "test_index" is not a table, materialized view, or TOAST table
+select pg_truncate_visibility_map('test_index');
+ERROR:  "test_index" is not a table, materialized view, or TOAST table
+create view test_view as select 1;
+select pg_visibility('test_view', 0);
+ERROR:  "test_view" is not a table, materialized view, or TOAST table
+select pg_visibility_map('test_view');
+ERROR:  "test_view" is not a table, materialized view, or TOAST table
+select pg_visibility_map_summary('test_view');
+ERROR:  "test_view" is not a table, materialized view, or TOAST table
+select pg_check_frozen('test_view');
+ERROR:  "test_view" is not a table, materialized view, or TOAST table
+select pg_truncate_visibility_map('test_view');
+ERROR:  "test_view" is not a table, materialized view, or TOAST table
+create sequence test_sequence;
+select pg_visibility('test_sequence', 0);
+ERROR:  "test_sequence" is not a table, materialized view, or TOAST table
+select pg_visibility_map('test_sequence');
+ERROR:  "test_sequence" is not a table, materialized view, or TOAST table
+select pg_visibility_map_summary('test_sequence');
+ERROR:  "test_sequence" is not a table, materialized view, or TOAST table
+select pg_check_frozen('test_sequence');
+ERROR:  "test_sequence" is not a table, materialized view, or TOAST table
+select pg_truncate_visibility_map('test_sequence');
+ERROR:  "test_sequence" is not a table, materialized view, or TOAST table
+create foreign data wrapper dummy;
+create server dummy_server foreign data wrapper dummy;
+create foreign table test_foreign_table () server dummy_server;
+select pg_visibility('test_foreign_table', 0);
+ERROR:  "test_foreign_table" is not a table, materialized view, or TOAST table
+select pg_visibility_map('test_foreign_table');
+ERROR:  "test_foreign_table" is not a table, materialized view, or TOAST table
+select pg_visibility_map_summary('test_foreign_table');
+ERROR:  "test_foreign_table" is not a table, materialized view, or TOAST table
+select pg_check_frozen('test_foreign_table');
+ERROR:  "test_foreign_table" is not a table, materialized view, or TOAST table
+select pg_truncate_visibility_map('test_foreign_table');
+ERROR:  "test_foreign_table" is not a table, materialized view, or TOAST table
+/* check some of the allowed relkinds */
+create table regular_table (a int);
+insert into regular_table values (1), (2);
+vacuum regular_table;
+select count(*) > 0 from pg_visibility('regular_table');
+ ?column? 
+----------
+ t
+(1 row)
+
+truncate regular_table;
+select count(*) > 0 from pg_visibility('regular_table');
+ ?column? 
+----------
+ f
+(1 row)
+
+create materialized view matview_visibility_test as select * from regular_table;
+vacuum matview_visibility_test;
+select count(*) > 0 from pg_visibility('matview_visibility_test');
+ ?column? 
+----------
+ f
+(1 row)
+
+insert into regular_table values (1), (2);
+refresh materialized view matview_visibility_test;
+select count(*) > 0 from pg_visibility('matview_visibility_test');
+ ?column? 
+----------
+ t
+(1 row)
+
+drop table test_partitioned;
+drop view test_view;
+drop sequence test_sequence;
+drop foreign table test_foreign_table;
+drop server dummy_server;
+drop foreign data wrapper dummy;
+drop materialized view matview_visibility_test;
+drop table regular_table;
diff --git a/contrib/pg_visibility/pg_visibility.c b/contrib/pg_visibility/pg_visibility.c
index 5580637870..2ecfa3c7d0 100644
--- a/contrib/pg_visibility/pg_visibility.c
+++ b/contrib/pg_visibility/pg_visibility.c
@@ -95,6 +95,22 @@ pg_visibility_map(PG_FUNCTION_ARGS)
 }
 
 /*
+ * check_relation_relkind - convenience routine to check that relation
+ * is of the relkind supported by the callers
+ */
+static void
+check_relation_relkind(Relation rel)
+{
+	if (rel->rd_rel->relkind != RELKIND_RELATION &&
+		rel->rd_rel->relkind != RELKIND_MATVIEW &&
+		rel->rd_rel->relkind != RELKIND_TOASTVALUE)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("\"%s\" is not a table, materialized view, or TOAST table",
+						RelationGetRelationName(rel))));
+}
+
+/*
  * Visibility map information for a single block of a relation, plus the
  * page-level information for the same block.
  */
@@ -114,6 +130,9 @@ pg_visibility(PG_FUNCTION_ARGS)
 
 	rel = relation_open(relid, AccessShareLock);
 
+	/* Only some relkind have the visibility info */
+	check_relation_relkind(rel);
+
 	if (blkno < 0 || blkno > MaxBlockNumber)
 		ereport(ERROR,
 				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
@@ -257,6 +276,10 @@ pg_visibility_map_summary(PG_FUNCTION_ARGS)
 	bool		nulls[2];
 
 	rel = relation_open(relid, AccessShareLock);
+
+	/* Only some relkind have the visibility info */
+	check_relation_relkind(rel);
+
 	nblocks = RelationGetNumberOfBlocks(rel);
 
 	for (blkno = 0; blkno < nblocks; ++blkno)
@@ -369,13 +392,8 @@ pg_truncate_visibility_map(PG_FUNCTION_ARGS)
 
 	rel = relation_open(relid, AccessExclusiveLock);
 
-	if (rel->rd_rel->relkind != RELKIND_RELATION &&
-		rel->rd_rel->relkind != RELKIND_MATVIEW &&
-		rel->rd_rel->relkind != RELKIND_TOASTVALUE)
-		ereport(ERROR,
-				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
-		   errmsg("\"%s\" is not a table, materialized view, or TOAST table",
-				  RelationGetRelationName(rel))));
+	/* Only some relkind have the visibility info */
+	check_relation_relkind(rel);
 
 	RelationOpenSmgr(rel);
 	rel->rd_smgr->smgr_vm_nblocks = InvalidBlockNumber;
@@ -464,6 +482,9 @@ collect_visibility_data(Oid relid, bool include_pd)
 
 	rel = relation_open(relid, AccessShareLock);
 
+	/* Only some relkind have the visibility info */
+	check_relation_relkind(rel);
+
 	nblocks = RelationGetNumberOfBlocks(rel);
 	info = palloc0(offsetof(vbits, bits) +nblocks);
 	info->next = 0;
@@ -543,13 +564,8 @@ collect_corrupt_items(Oid relid, bool all_visible, bool all_frozen)
 
 	rel = relation_open(relid, AccessShareLock);
 
-	if (rel->rd_rel->relkind != RELKIND_RELATION &&
-		rel->rd_rel->relkind != RELKIND_MATVIEW &&
-		rel->rd_rel->relkind != RELKIND_TOASTVALUE)
-		ereport(ERROR,
-				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
-		   errmsg("\"%s\" is not a table, materialized view, or TOAST table",
-				  RelationGetRelationName(rel))));
+	/* Only some relkind have the visibility info */
+	check_relation_relkind(rel);
 
 	nblocks = RelationGetNumberOfBlocks(rel);
 
diff --git a/contrib/pg_visibility/sql/pg_visibility.sql b/contrib/pg_visibility/sql/pg_visibility.sql
new file mode 100644
index 0000000000..9f481d23f7
--- /dev/null
+++ b/contrib/pg_visibility/sql/pg_visibility.sql
@@ -0,0 +1,66 @@
+CREATE EXTENSION pg_visibility;
+
+--
+-- check that using the module's functions with unsupported relations will fail
+--
+create table test_partitioned (a int) partition by list (a);
+select pg_visibility('test_partitioned', 0);
+select pg_visibility_map('test_partitioned');
+select pg_visibility_map_summary('test_partitioned');
+select pg_check_frozen('test_partitioned');
+select pg_truncate_visibility_map('test_partitioned');
+
+create table test_partition partition of test_partitioned for values in (1);
+create index test_index on test_partition (a);
+select pg_visibility('test_index', 0);
+select pg_visibility_map('test_index');
+select pg_visibility_map_summary('test_index');
+select pg_check_frozen('test_index');
+select pg_truncate_visibility_map('test_index');
+
+create view test_view as select 1;
+select pg_visibility('test_view', 0);
+select pg_visibility_map('test_view');
+select pg_visibility_map_summary('test_view');
+select pg_check_frozen('test_view');
+select pg_truncate_visibility_map('test_view');
+
+create sequence test_sequence;
+select pg_visibility('test_sequence', 0);
+select pg_visibility_map('test_sequence');
+select pg_visibility_map_summary('test_sequence');
+select pg_check_frozen('test_sequence');
+select pg_truncate_visibility_map('test_sequence');
+
+create foreign data wrapper dummy;
+create server dummy_server foreign data wrapper dummy;
+create foreign table test_foreign_table () server dummy_server;
+select pg_visibility('test_foreign_table', 0);
+select pg_visibility_map('test_foreign_table');
+select pg_visibility_map_summary('test_foreign_table');
+select pg_check_frozen('test_foreign_table');
+select pg_truncate_visibility_map('test_foreign_table');
+
+/* check some of the allowed relkinds */
+create table regular_table (a int);
+insert into regular_table values (1), (2);
+vacuum regular_table;
+select count(*) > 0 from pg_visibility('regular_table');
+truncate regular_table;
+select count(*) > 0 from pg_visibility('regular_table');
+
+create materialized view matview_visibility_test as select * from regular_table;
+vacuum matview_visibility_test;
+select count(*) > 0 from pg_visibility('matview_visibility_test');
+insert into regular_table values (1), (2);
+refresh materialized view matview_visibility_test;
+select count(*) > 0 from pg_visibility('matview_visibility_test');
+
+drop table test_partitioned;
+drop view test_view;
+drop sequence test_sequence;
+drop foreign table test_foreign_table;
+drop server dummy_server;
+drop foreign data wrapper dummy;
+drop materialized view matview_visibility_test;
+drop table regular_table;
diff --git a/contrib/pgstattuple/expected/pgstattuple.out b/contrib/pgstattuple/expected/pgstattuple.out
index 169d1932b2..dfd15b0c63 100644
--- a/contrib/pgstattuple/expected/pgstattuple.out
+++ b/contrib/pgstattuple/expected/pgstattuple.out
@@ -138,3 +138,32 @@ select * from pgstathashindex('test_hashidx');
        2 |            4 |              0 |            1 |          0 |          0 |          0 |          100
 (1 row)
 
+-- check that using any of these functions with unsupported relations will fail
+create table test_partitioned (a int) partition by range (a);
+select pgstattuple('test_partitioned');
+ERROR:  "test_partitioned" (partitioned table) is not supported
+select pgstattuple_approx('test_partitioned');
+ERROR:  "test_partitioned" is not a table or materialized view
+select pg_relpages('test_partitioned');
+ERROR:  "test_partitioned" is not a table, index, materialized view, sequence, or TOAST table
+create view test_view as select 1;
+select pgstattuple('test_view');
+ERROR:  "test_view" (view) is not supported
+select pgstattuple_approx('test_view');
+ERROR:  "test_view" is not a table or materialized view
+select pg_relpages('test_view');
+ERROR:  "test_view" is not a table, index, materialized view, sequence, or TOAST table
+create foreign data wrapper dummy;
+create server dummy_server foreign data wrapper dummy;
+create foreign table test_foreign_table () server dummy_server;
+select pgstattuple('test_foreign_table');
+ERROR:  "test_foreign_table" (foreign table) is not supported
+select pgstattuple_approx('test_foreign_table');
+ERROR:  "test_foreign_table" is not a table or materialized view
+select pg_relpages('test_foreign_table');
+ERROR:  "test_foreign_table" is not a table, index, materialized view, sequence, or TOAST table
+drop table test_partitioned;
+drop view test_view;
+drop foreign table test_foreign_table;
+drop server dummy_server;
+drop foreign data wrapper dummy;
diff --git a/contrib/pgstattuple/pgstatindex.c b/contrib/pgstattuple/pgstatindex.c
index 17a53e3bb7..04b20915f8 100644
--- a/contrib/pgstattuple/pgstatindex.c
+++ b/contrib/pgstattuple/pgstatindex.c
@@ -361,6 +361,24 @@ pgstatindex_impl(Relation rel, FunctionCallInfo fcinfo)
 	return result;
 }
 
+/*
+ * check_relation_relkind - convenience routine to check that relation
+ * is of the relkind supported by the callers
+ */
+static void
+check_relation_relkind(Relation rel)
+{
+	if (rel->rd_rel->relkind != RELKIND_RELATION &&
+		rel->rd_rel->relkind != RELKIND_INDEX &&
+		rel->rd_rel->relkind != RELKIND_MATVIEW &&
+		rel->rd_rel->relkind != RELKIND_SEQUENCE &&
+		rel->rd_rel->relkind != RELKIND_TOASTVALUE)
+		ereport(ERROR,
+				(errcode(ERRCODE_WRONG_OBJECT_TYPE),
+				 errmsg("\"%s\" is not a table, index, materialized view, sequence, or TOAST table",
+						RelationGetRelationName(rel))));
+}
+
 /* --------------------------------------------------------
  * pg_relpages()
  *
@@ -388,6 +406,9 @@ pg_relpages(PG_FUNCTION_ARGS)
 	relrv = makeRangeVarFromNameList(textToQualifiedNameList(relname));
 	rel = relation_openrv(relrv, AccessShareLock);
 
+	/* only some relkinds have storage */
+	check_relation_relkind(rel);
+
 	/* note: this will work OK on non-local temp tables */
 
 	relpages = RelationGetNumberOfBlocks(rel);
@@ -409,6 +430,9 @@ pg_relpages_v1_5(PG_FUNCTION_ARGS)
 	relrv = makeRangeVarFromNameList(textToQualifiedNameList(relname));
 	rel = relation_openrv(relrv, AccessShareLock);
 
+	/* only some relkinds have storage */
+	check_relation_relkind(rel);
+
 	/* note: this will work OK on non-local temp tables */
 
 	relpages = RelationGetNumberOfBlocks(rel);
@@ -433,6 +457,9 @@ pg_relpagesbyid(PG_FUNCTION_ARGS)
 
 	rel = relation_open(relid, AccessShareLock);
 
+	/* only some relkinds have storage */
+	check_relation_relkind(rel);
+
 	/* note: this will work OK on non-local temp tables */
 
 	relpages = RelationGetNumberOfBlocks(rel);
@@ -452,6 +479,9 @@ pg_relpagesbyid_v1_5(PG_FUNCTION_ARGS)
 
 	rel = relation_open(relid, AccessShareLock);
 
+	/* only some relkinds have storage */
+	check_relation_relkind(rel);
+
 	/* note: this will work OK on non-local temp tables */
 
 	relpages = RelationGetNumberOfBlocks(rel);
diff --git a/contrib/pgstattuple/pgstattuple.c b/contrib/pgstattuple/pgstattuple.c
index 06a1992bb1..b2432f43ed 100644
--- a/contrib/pgstattuple/pgstattuple.c
+++ b/contrib/pgstattuple/pgstattuple.c
@@ -293,6 +293,9 @@ pgstat_relation(Relation rel, FunctionCallInfo fcinfo)
 		case RELKIND_FOREIGN_TABLE:
 			err = "foreign table";
 			break;
+		case RELKIND_PARTITIONED_TABLE:
+			err = "partitioned table";
+			break;
 		default:
 			err = "unknown";
 			break;
diff --git a/contrib/pgstattuple/sql/pgstattuple.sql b/contrib/pgstattuple/sql/pgstattuple.sql
index 81fd5d693b..e42ac5db95 100644
--- a/contrib/pgstattuple/sql/pgstattuple.sql
+++ b/contrib/pgstattuple/sql/pgstattuple.sql
@@ -51,3 +51,27 @@ select * from pgstatginindex('test_ginidx');
 create index test_hashidx on test using hash (b);
 
 select * from pgstathashindex('test_hashidx');
+
+-- check that using any of these functions with unsupported relations will fail
+create table test_partitioned (a int) partition by range (a);
+select pgstattuple('test_partitioned');
+select pgstattuple_approx('test_partitioned');
+select pg_relpages('test_partitioned');
+
+create view test_view as select 1;
+select pgstattuple('test_view');
+select pgstattuple_approx('test_view');
+select pg_relpages('test_view');
+
+create foreign data wrapper dummy;
+create server dummy_server foreign data wrapper dummy;
+create foreign table test_foreign_table () server dummy_server;
+select pgstattuple('test_foreign_table');
+select pgstattuple_approx('test_foreign_table');
+select pg_relpages('test_foreign_table');
+
+drop table test_partitioned;
+drop view test_view;
+drop foreign table test_foreign_table;
+drop server dummy_server;
+drop foreign data wrapper dummy;
-- 
2.11.0

