Add a pg_catalog.citus_nodes view

src/backend/distributed/metadata/metadata_utility.c

@@ -67,6 +67,7 @@
 
 
 /* Local functions forward declarations */
+static uint64 NodeDatabaseSize(char *nodeName, int nodePort, char *databaseName);
 static uint64 * AllocateUint64(uint64 value);
 static void RecordDistributedRelationDependencies(Oid distributedRelationId);
 static GroupShardPlacement * TupleToGroupShardPlacement(TupleDesc tupleDesc,
@@ -81,11 +82,67 @@ static void ErrorIfNotSuitableToGetSize(Oid relationId);
 
 
 /* exports for SQL callable functions */
+PG_FUNCTION_INFO_V1(citus_database_size);
+PG_FUNCTION_INFO_V1(citus_node_database_size);
 PG_FUNCTION_INFO_V1(citus_table_size);
 PG_FUNCTION_INFO_V1(citus_total_relation_size);
 PG_FUNCTION_INFO_V1(citus_relation_size);
 
 
+/*
+ * citus_database_size is a UDF that returns the sum of the database sizes
+ * across all nodes.
+ */
+Datum
+citus_database_size(PG_FUNCTION_ARGS)
+{
+	Datum databaseNameDatum = PG_GETARG_DATUM(0);
+	Name databaseName = DatumGetName(databaseNameDatum);
+	char *databaseNameString = NameStr(*databaseName);
+
+	CheckCitusVersion(ERROR);
+
+	Datum databaseSizeDatum = DirectFunctionCall1(pg_database_size_name,
+												  databaseNameDatum);
+
+	uint64 databaseSize = Int64GetDatum(databaseSizeDatum);
+
+	/* we calculated the coordinator database size above, so only query workers */
+	List *workerNodeList = ActiveReadableNonCoordinatorNodeList();
+	WorkerNode *workerNode = NULL;
+
+	foreach_ptr(workerNode, workerNodeList)
+	{
+		databaseSize += NodeDatabaseSize(workerNode->workerName, workerNode->workerPort,
+										 databaseNameString);
+	}
+
+	PG_RETURN_UINT64(databaseSize);
+}
+
+
+/*
+ * citus_node_database_size is a UDF that returns the size of the database
+ * on a given node.
+ */
+Datum
+citus_node_database_size(PG_FUNCTION_ARGS)
+{
+	text *nodeNameText = PG_GETARG_TEXT_P(0);
+	char *nodeNameString = text_to_cstring(nodeNameText);
+	int32 nodePort = PG_GETARG_INT32(1);
+	Name databaseName = PG_GETARG_NAME(2);
+	char *databaseNameString = NameStr(*databaseName);
+
+	CheckCitusVersion(ERROR);
+
+	uint64 nodeDatabaseSize = NodeDatabaseSize(nodeNameString, nodePort,
+											   databaseNameString);
+
+	PG_RETURN_UINT64(nodeDatabaseSize);
+}
+
+
 /*
  * citus_total_relation_size accepts a table name and returns a distributed table
  * and its indexes' total relation size.
@@ -177,6 +234,41 @@ citus_relation_size(PG_FUNCTION_ARGS)
 }
 
 
+/*
+ * NodeDatabaseSize returns the size of a database on a given node.
+ */
+static uint64
+NodeDatabaseSize(char *nodeName, int nodePort, char *databaseName)
+{
+	uint32 connectionFlags = 0;
+	PGresult *result = NULL;
+	StringInfo sizeQuery = makeStringInfo();
+	bool raiseErrors = true;
+
+	appendStringInfo(sizeQuery, "SELECT pg_catalog.pg_database_size(%s)",
+					 quote_literal_cstr(quote_identifier(databaseName)));
+
+	MultiConnection *connection = GetNodeConnection(connectionFlags, nodeName, nodePort);
+
+	int queryResult = ExecuteOptionalRemoteCommand(connection, sizeQuery->data,
+												   &result);
+	if (queryResult != 0)
+	{
+		ereport(ERROR, (errcode(ERRCODE_CONNECTION_FAILURE),
+						errmsg("cannot get the size because of a connection error")));
+	}
+
+	List *sizeList = ReadFirstColumnAsText(result);
+	StringInfo databaseSizeStringInfo = (StringInfo) linitial(sizeList);
+	uint64 databaseSize = SafeStringToUint64(databaseSizeStringInfo->data);
+
+	PQclear(result);
+	ClearResults(connection, raiseErrors);
+
+	return databaseSize;
+}
+
+
 /*
  * DistributedTableSize is helper function for each kind of citus size functions.
  * It first checks whether the table is distributed and size query can be run on

src/backend/distributed/sql/citus--9.5-1--10.0-1.sql

@@ -4,6 +4,9 @@ DROP FUNCTION IF EXISTS pg_catalog.citus_total_relation_size(regclass);
 
 #include "udfs/citus_total_relation_size/10.0-1.sql"
 #include "udfs/citus_tables/10.0-1.sql"
+#include "udfs/citus_database_size/10.0-1.sql"
+#include "udfs/citus_node_database_size/10.0-1.sql"
+#include "udfs/citus_nodes/10.0-1.sql"
 #include "udfs/citus_finish_pg_upgrade/10.0-1.sql"
 
 #include "../../columnar/sql/columnar--9.5-1--10.0-1.sql"

src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql

@@ -9,3 +9,8 @@ DROP VIEW public.citus_tables;
 DROP FUNCTION pg_catalog.citus_total_relation_size(regclass,boolean);
 
 #include "../udfs/citus_total_relation_size/7.0-1.sql"
+
+DROP VIEW pg_catalog.citus_nodes;
+
+DROP FUNCTION pg_catalog.citus_database_size(name);
+DROP FUNCTION pg_catalog.citus_node_database_size(text,int,name);

src/backend/distributed/sql/udfs/citus_database_size/10.0-1.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_database_size(
+	dbname name default current_database())
+RETURNS bigint
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_database_size$$;
+COMMENT ON FUNCTION pg_catalog.citus_database_size(name)
+IS 'returns the size of the database across all nodes';

src/backend/distributed/sql/udfs/citus_database_size/latest.sql

@@ -0,0 +1,7 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_database_size(
+	dbname name default current_database())
+RETURNS bigint
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_database_size$$;
+COMMENT ON FUNCTION pg_catalog.citus_database_size(name)
+IS 'returns the size of the database across all nodes';

src/backend/distributed/sql/udfs/citus_node_database_size/10.0-1.sql

@@ -0,0 +1,9 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_node_database_size(
+	nodename text,
+	nodeport int default 5432,
+	dbname name default current_database())
+RETURNS bigint
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_node_database_size$$;
+COMMENT ON FUNCTION pg_catalog.citus_node_database_size(text, int, name)
+IS 'returns the size of the database on the given node';

src/backend/distributed/sql/udfs/citus_node_database_size/latest.sql

@@ -0,0 +1,9 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_node_database_size(
+	nodename text,
+	nodeport int default 5432,
+	dbname name default current_database())
+RETURNS bigint
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_node_database_size$$;
+COMMENT ON FUNCTION pg_catalog.citus_node_database_size(text, int, name)
+IS 'returns the size of the database on the given node';

src/backend/distributed/sql/udfs/citus_nodes/10.0-1.sql

@@ -0,0 +1,41 @@
+CREATE VIEW citus.citus_nodes AS
+SELECT
+  nodename AS "Host",
+  nodeport AS "Port",
+  CASE WHEN groupid = 0 THEN 'coordinator' ELSE 'worker' END AS "Node Type",
+  pg_size_pretty(citus_node_database_size(nodename, nodeport)) AS "Database Size",
+  (SELECT
+     count(*)
+   FROM
+     pg_dist_partition t
+   JOIN
+     pg_dist_shard s
+   USING
+     (logicalrelid)
+   JOIN
+     pg_dist_placement p
+   USING
+     (shardid)
+   WHERE p.groupid = n.groupid AND t.partmethod <> 'n'
+  ) AS "Distributed Table Shards",
+  (SELECT
+     count(*)
+   FROM
+     pg_dist_partition t
+   JOIN
+     pg_dist_shard s
+   USING
+     (logicalrelid)
+   JOIN
+     pg_dist_placement p
+   USING
+     (shardid)
+   WHERE p.groupid = n.groupid AND t.partmethod = 'n' AND t.repmodel = 't'
+  ) AS "Reference Tables"
+FROM
+  pg_dist_node n
+ORDER BY
+  groupid;
+
+ALTER VIEW citus.citus_nodes SET SCHEMA pg_catalog;
+GRANT SELECT ON pg_catalog.citus_nodes TO public;

src/backend/distributed/sql/udfs/citus_nodes/latest.sql

@@ -0,0 +1,41 @@
+CREATE VIEW citus.citus_nodes AS
+SELECT
+  nodename AS "Host",
+  nodeport AS "Port",
+  CASE WHEN groupid = 0 THEN 'coordinator' ELSE 'worker' END AS "Node Type",
+  pg_size_pretty(citus_node_database_size(nodename, nodeport)) AS "Database Size",
+  (SELECT
+     count(*)
+   FROM
+     pg_dist_partition t
+   JOIN
+     pg_dist_shard s
+   USING
+     (logicalrelid)
+   JOIN
+     pg_dist_placement p
+   USING
+     (shardid)
+   WHERE p.groupid = n.groupid AND t.partmethod <> 'n'
+  ) AS "Distributed Table Shards",
+  (SELECT
+     count(*)
+   FROM
+     pg_dist_partition t
+   JOIN
+     pg_dist_shard s
+   USING
+     (logicalrelid)
+   JOIN
+     pg_dist_placement p
+   USING
+     (shardid)
+   WHERE p.groupid = n.groupid AND t.partmethod = 'n' AND t.repmodel = 't'
+  ) AS "Reference Tables"
+FROM
+  pg_dist_node n
+ORDER BY
+  groupid;
+
+ALTER VIEW citus.citus_nodes SET SCHEMA pg_catalog;
+GRANT SELECT ON pg_catalog.citus_nodes TO public;

src/test/regress/expected/multi_extension.out

@@ -482,7 +482,9 @@ SELECT * FROM print_extension_changes();
                                               | access method columnar
                                               | function alter_columnar_table_reset(regclass,boolean,boolean,boolean,boolean)
                                               | function alter_columnar_table_set(regclass,integer,integer,name,integer)
+                                              | function citus_database_size(name)
                                               | function citus_internal.columnar_ensure_objects_exist()
+                                              | function citus_node_database_size(text,integer,name)
                                               | function citus_total_relation_size(regclass,boolean)
                                               | function columnar.columnar_handler(internal)
                                               | schema columnar
@@ -490,8 +492,9 @@ SELECT * FROM print_extension_changes();
                                               | table columnar.columnar_skipnodes
                                               | table columnar.columnar_stripes
                                               | table columnar.options
+                                              | view citus_nodes
                                               | view citus_tables
-(13 rows)
+(16 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/multi_extension_0.out

@@ -479,15 +479,18 @@ SELECT * FROM print_extension_changes();
                previous_object                |                     current_object
 ---------------------------------------------------------------------
  function citus_total_relation_size(regclass) |
+                                              | function citus_database_size(name)
                                               | function citus_internal.columnar_ensure_objects_exist()
+                                              | function citus_node_database_size(text,integer,name)
                                               | function citus_total_relation_size(regclass,boolean)
                                               | schema columnar
                                               | sequence columnar.storageid_seq
                                               | table columnar.columnar_skipnodes
                                               | table columnar.columnar_stripes
                                               | table columnar.options
+                                              | view citus_nodes
                                               | view citus_tables
-(9 rows)
+(12 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -31,6 +31,7 @@ ORDER BY 1;
  function citus_add_rebalance_strategy(name,regproc,regproc,regproc,real,real)
  function citus_blocking_pids(integer)
  function citus_create_restore_point(text)
+ function citus_database_size(name)
  function citus_dist_stat_activity()
  function citus_drop_trigger()
  function citus_executor_name(integer)
@@ -51,6 +52,7 @@ ORDER BY 1;
  function citus_jsonb_concatenate(jsonb,jsonb)
  function citus_jsonb_concatenate_final(jsonb)
  function citus_node_capacity_1(integer)
+ function citus_node_database_size(text,integer,name)
  function citus_prepare_pg_upgrade()
  function citus_query_stats()
  function citus_relation_size(regclass)
@@ -212,11 +214,12 @@ ORDER BY 1;
  type noderole
  view citus_dist_stat_activity
  view citus_lock_waits
+ view citus_nodes
  view citus_shard_indexes_on_worker
  view citus_shards_on_worker
  view citus_stat_statements
  view citus_tables
  view citus_worker_stat_activity
  view pg_dist_shard_placement
-(202 rows)
+(205 rows)
 

src/test/regress/expected/upgrade_list_citus_objects_0.out

@@ -28,6 +28,7 @@ ORDER BY 1;
  function citus_add_rebalance_strategy(name,regproc,regproc,regproc,real,real)
  function citus_blocking_pids(integer)
  function citus_create_restore_point(text)
+ function citus_database_size(name)
  function citus_dist_stat_activity()
  function citus_drop_trigger()
  function citus_executor_name(integer)
@@ -48,6 +49,7 @@ ORDER BY 1;
  function citus_jsonb_concatenate(jsonb,jsonb)
  function citus_jsonb_concatenate_final(jsonb)
  function citus_node_capacity_1(integer)
+ function citus_node_database_size(text,integer,name)
  function citus_prepare_pg_upgrade()
  function citus_query_stats()
  function citus_relation_size(regclass)
@@ -208,11 +210,12 @@ ORDER BY 1;
  type noderole
  view citus_dist_stat_activity
  view citus_lock_waits
+ view citus_nodes
  view citus_shard_indexes_on_worker
  view citus_shards_on_worker
  view citus_stat_statements
  view citus_tables
  view citus_worker_stat_activity
  view pg_dist_shard_placement
-(198 rows)
+(201 rows)