From 4a08bb604d5deddddcf1f45963dd10438b28027c Mon Sep 17 00:00:00 2001
From: Alvaro Herrera <alvherre@alvh.no-ip.org>
Date: Fri, 28 Aug 2015 15:10:20 -0300
Subject: [PATCH 13/24] add colstore function to dbsize

---
 src/backend/utils/adt/dbsize.c | 64 ++++++++++++++++++++++++++++++++++++++++++
 src/include/catalog/pg_proc.h  |  2 ++
 src/include/utils/builtins.h   |  1 +
 3 files changed, 67 insertions(+)

diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
index 82311b4..6dd8155 100644
--- a/src/backend/utils/adt/dbsize.c
+++ b/src/backend/utils/adt/dbsize.c
@@ -443,6 +443,46 @@ calculate_indexes_size(Relation rel)
 	return size;
 }
 
+/*
+ * Calculate total on-disk size of all column stores attached to the given table.
+ *
+ * Can be applied safely to a column store, but you'll just get zero.
+ */
+static int64
+calculate_column_stores_size(Relation rel)
+{
+	int64		size = 0;
+
+	/*
+	 * Aggregate all indexes on the given relation
+	 */
+	if (rel->rd_rel->relhascstore)
+	{
+		List	   *cstore_oids = RelationGetColStoreList(rel);
+		ListCell   *cell;
+
+		foreach(cell, cstore_oids)
+		{
+			Oid			cstoreOid = lfirst_oid(cell);
+			Relation	cstoreRel;
+			ForkNumber	forkNum;
+
+			cstoreRel = relation_open(cstoreOid, AccessShareLock);
+
+			for (forkNum = 0; forkNum <= MAX_FORKNUM; forkNum++)
+				size += calculate_relation_size(&(cstoreRel->rd_node),
+												cstoreRel->rd_backend,
+												forkNum);
+
+			relation_close(cstoreRel, AccessShareLock);
+		}
+
+		list_free(cstore_oids);
+	}
+
+	return size;
+}
+
 Datum
 pg_table_size(PG_FUNCTION_ARGS)
 {
@@ -481,6 +521,25 @@ pg_indexes_size(PG_FUNCTION_ARGS)
 	PG_RETURN_INT64(size);
 }
 
+Datum
+pg_column_stores_size(PG_FUNCTION_ARGS)
+{
+	Oid			relOid = PG_GETARG_OID(0);
+	Relation	rel;
+	int64		size;
+
+	rel = try_relation_open(relOid, AccessShareLock);
+
+	if (rel == NULL)
+		PG_RETURN_NULL();
+
+	size = calculate_column_stores_size(rel);
+
+	relation_close(rel, AccessShareLock);
+
+	PG_RETURN_INT64(size);
+}
+
 /*
  *	Compute the on-disk size of all files for the relation,
  *	including heap data, index data, toast data, FSM, VM.
@@ -501,6 +560,11 @@ calculate_total_relation_size(Relation rel)
 	 */
 	size += calculate_indexes_size(rel);
 
+	/*
+	 * Add size of all attached column stores as well
+	 */
+	size += calculate_column_stores_size(rel);
+
 	return size;
 }
 
diff --git a/src/include/catalog/pg_proc.h b/src/include/catalog/pg_proc.h
index ddf7c67..f0481bd 100644
--- a/src/include/catalog/pg_proc.h
+++ b/src/include/catalog/pg_proc.h
@@ -3664,6 +3664,8 @@ DATA(insert OID = 2997 ( pg_table_size			PGNSP PGUID 12 1 0 0 0 f f f f t f v 1
 DESCR("disk space usage for the specified table, including TOAST, free space and visibility map");
 DATA(insert OID = 2998 ( pg_indexes_size		PGNSP PGUID 12 1 0 0 0 f f f f t f v 1 0 20 "2205" _null_ _null_ _null_ _null_ _null_ pg_indexes_size _null_ _null_ _null_ ));
 DESCR("disk space usage for all indexes attached to the specified table");
+DATA(insert OID = 3350 ( pg_column_stores_size	PGNSP PGUID 12 1 0 0 0 f f f f t f v 1 0 20 "2205" _null_ _null_ _null_ _null_ _null_ pg_column_stores_size _null_ _null_ _null_ ));
+DESCR("disk space usage for all column stores attached to the specified table");
 DATA(insert OID = 2999 ( pg_relation_filenode	PGNSP PGUID 12 1 0 0 0 f f f f t f s 1 0 26 "2205" _null_ _null_ _null_ _null_ _null_ pg_relation_filenode _null_ _null_ _null_ ));
 DESCR("filenode identifier of relation");
 DATA(insert OID = 3454 ( pg_filenode_relation PGNSP PGUID 12 1 0 0 0 f f f f t f s 2 0 2205 "26 26" _null_ _null_ _null_ _null_ _null_ pg_filenode_relation _null_ _null_ _null_ ));
diff --git a/src/include/utils/builtins.h b/src/include/utils/builtins.h
index fc1679e..aa79730 100644
--- a/src/include/utils/builtins.h
+++ b/src/include/utils/builtins.h
@@ -464,6 +464,7 @@ extern Datum pg_size_pretty(PG_FUNCTION_ARGS);
 extern Datum pg_size_pretty_numeric(PG_FUNCTION_ARGS);
 extern Datum pg_table_size(PG_FUNCTION_ARGS);
 extern Datum pg_indexes_size(PG_FUNCTION_ARGS);
+extern Datum pg_column_stores_size(PG_FUNCTION_ARGS);
 extern Datum pg_relation_filenode(PG_FUNCTION_ARGS);
 extern Datum pg_filenode_relation(PG_FUNCTION_ARGS);
 extern Datum pg_relation_filepath(PG_FUNCTION_ARGS);
-- 
2.1.4