From 546bce4d7c0e418fdb086cebfb77b135ac6eb4d3 Mon Sep 17 00:00:00 2001
From: fatkodima <fatkodima123@gmail.com>
Date: Tue, 20 Feb 2024 18:59:12 +0200
Subject: [PATCH] Show space usage for partitioned tables

---
 CHANGELOG.md                |  1 +
 lib/pghero/methods/space.rb |  9 +++++++--
 test/space_test.rb          |  3 +++
 test/test_helper.rb         | 11 +++++++++++
 4 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ddb32765..60cff291 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,6 @@
 ## 3.4.2 (unreleased)
 
+- Added partitioned tables to space page
 - Added materialized views to space page
 - Fixed error with `slow_queries` method
 
diff --git a/lib/pghero/methods/space.rb b/lib/pghero/methods/space.rb
index f54be0e6..cc030220 100644
--- a/lib/pghero/methods/space.rb
+++ b/lib/pghero/methods/space.rb
@@ -10,7 +10,12 @@ def relation_sizes
           SELECT
             n.nspname AS schema,
             c.relname AS relation,
-            CASE c.relkind WHEN 'r' THEN 'table' WHEN 'm' then 'matview' ELSE 'index' END AS type,
+            CASE c.relkind
+              WHEN 'r' THEN 'table'
+              WHEN 'p' THEN 'partitioned table'
+              WHEN 'm' then 'matview'
+              ELSE 'index'
+            END AS type,
             pg_table_size(c.oid) AS size_bytes
           FROM
             pg_class c
@@ -19,7 +24,7 @@ def relation_sizes
           WHERE
             n.nspname NOT IN ('pg_catalog', 'information_schema')
             AND n.nspname !~ '^pg_toast'
-            AND c.relkind IN ('r', 'm', 'i')
+            AND c.relkind IN ('r', 'p', 'm', 'i')
           ORDER BY
             pg_table_size(c.oid) DESC,
             2 ASC
diff --git a/test/space_test.rb b/test/space_test.rb
index 551b4fdf..3f9a3d89 100644
--- a/test/space_test.rb
+++ b/test/space_test.rb
@@ -10,6 +10,9 @@ def test_relation_sizes
     assert relation_sizes.find { |r| r[:relation] == "users" && r[:type] == "table" }
     assert relation_sizes.find { |r| r[:relation] == "users_pkey" && r[:type] == "index" }
     assert relation_sizes.find { |r| r[:relation] == "all_users" && r[:type] == "matview" }
+    assert relation_sizes.find { |r| r[:relation] == "partitioned_users" && r[:type] == "partitioned table" }
+    assert relation_sizes.find { |r| r[:relation] == "partitioned_users_1" && r[:type] == "table" }
+    assert relation_sizes.find { |r| r[:relation] == "partitioned_users_2" && r[:type] == "table" }
   end
 
   def test_table_sizes
diff --git a/test/test_helper.rb b/test/test_helper.rb
index 5ae96916..2c45ddc5 100644
--- a/test/test_helper.rb
+++ b/test/test_helper.rb
@@ -74,3 +74,14 @@ class User < ActiveRecord::Base
 User.insert_all!(users)
 ActiveRecord::Base.connection.execute("ANALYZE users")
 ActiveRecord::Base.connection.execute("CREATE MATERIALIZED VIEW all_users AS SELECT * FROM users")
+ActiveRecord::Base.connection.execute(<<~SQL)
+  CREATE TABLE partitioned_users (
+    city_id integer NOT NULL
+  ) PARTITION BY RANGE (city_id);
+
+  CREATE TABLE partitioned_users_1 PARTITION OF partitioned_users
+    FOR VALUES FROM (1) TO (1000);
+
+  CREATE TABLE partitioned_users_2 PARTITION OF partitioned_users
+    FOR VALUES FROM (1001) TO (2000);
+SQL