From 9074310e3ba56bc06c258c446784ccd259d65327 Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Sat, 4 May 2024 10:39:24 -0700
Subject: [PATCH] Added option for materialized views

---
 app/controllers/pg_hero/home_controller.rb | 2 +-
 lib/pghero/methods/space.rb                | 4 ++--
 test/space_test.rb                         | 7 +++++++
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/app/controllers/pg_hero/home_controller.rb b/app/controllers/pg_hero/home_controller.rb
index e7c56860..9c8a709d 100644
--- a/app/controllers/pg_hero/home_controller.rb
+++ b/app/controllers/pg_hero/home_controller.rb
@@ -87,7 +87,7 @@ def space
       @days = (params[:days] || 7).to_i
       @database_size = @database.database_size
       @only_tables = params[:tables].present?
-      @relation_sizes, @sizes_timeout = rescue_timeout([]) { @only_tables ? @database.table_sizes : @database.relation_sizes }
+      @relation_sizes, @sizes_timeout = rescue_timeout([]) { @only_tables ? @database.table_sizes : @database.relation_sizes(matviews: true) }
       @space_stats_enabled = @database.space_stats_enabled? && !@only_tables
       if @space_stats_enabled
         space_growth = @database.space_growth(days: @days, relation_sizes: @relation_sizes)
diff --git a/lib/pghero/methods/space.rb b/lib/pghero/methods/space.rb
index f54be0e6..9df6d3dc 100644
--- a/lib/pghero/methods/space.rb
+++ b/lib/pghero/methods/space.rb
@@ -5,7 +5,7 @@ def database_size
         PgHero.pretty_size select_one("SELECT pg_database_size(current_database())")
       end
 
-      def relation_sizes
+      def relation_sizes(matviews: false)
         select_all_size <<~SQL
           SELECT
             n.nspname AS schema,
@@ -19,7 +19,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', 'i'#{matviews ? ", 'm'" : ""})
           ORDER BY
             pg_table_size(c.oid) DESC,
             2 ASC
diff --git a/test/space_test.rb b/test/space_test.rb
index 551b4fdf..1cb2ce84 100644
--- a/test/space_test.rb
+++ b/test/space_test.rb
@@ -9,6 +9,13 @@ def test_relation_sizes
     relation_sizes = database.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" }
+    refute relation_sizes.find { |r| r[:relation] == "all_users" && r[:type] == "matview" }
+  end
+
+  def test_relation_sizes_matviews
+    relation_sizes = database.relation_sizes(matviews: true)
+    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" }
   end