Postgres: Sync table row count (#39494)

resources/migrations/001_update_migrations.yaml

@@ -6055,6 +6055,19 @@ databaseChangeLog:
             constraintName: idx_cache_config_unique_model
             columnNames: model, model_id
 
+  - changeSet:
+      id: v50.2024-03-21T17:41:00
+      author: qnkhuat
+      comment: Add metabase_table.estimated_row_count
+      changes:
+        - addColumn:
+            tableName: metabase_table
+            columns:
+              - column:
+                  name: estimated_row_count
+                  type: bigint
+                  remarks: The estimated row count
+
   # >>>>>>>>>> DO NOT ADD NEW MIGRATIONS BELOW THIS LINE! ADD THEM ABOVE <<<<<<<<<<
 
 ########################################################################################################################

src/metabase/api/notify.clj

@@ -58,8 +58,8 @@
       (let [driver (driver.u/database->driver database)
             {db-tables :tables} (driver/describe-database driver database)]
         (if-let [table (some (fn [table-in-db]
-                               (when (= (dissoc table-in-db :description)
-                                        {:schema schema_name :name table_name})
+                               (when (and (= schema_name (:schema table-in-db))
+                                          (= table_name (:name table-in-db)))
                                  table-in-db))
                              db-tables)]
           (let [created (sync-tables/create-or-reactivate-table! database table)]

src/metabase/driver/postgres.clj

@@ -227,10 +227,12 @@
                            [:inline "m"] [:inline "MATERIALIZED VIEW"]
                            :else nil]
                           :type]
-                         [:d.description :description]]
-             :from      [[:pg_catalog.pg_namespace :n]
-                         [:pg_catalog.pg_class :c]]
-             :left-join [[:pg_catalog.pg_description :d] [:and [:= :c.oid :d.objoid] [:= :d.objsubid 0] [:= :d.classoid [:raw "'pg_class'::regclass"]]]]
+                         [:d.description :description]
+                         [:stat.n_live_tup :estimated_row_count]]
+             :from      [[:pg_catalog.pg_class :c]]
+             :join      [[:pg_catalog.pg_namespace :n]   [:= :c.relnamespace :n.oid]]
+             :left-join [[:pg_catalog.pg_description :d] [:and [:= :c.oid :d.objoid] [:= :d.objsubid 0] [:= :d.classoid [:raw "'pg_class'::regclass"]]]
+                         [:pg_stat_user_tables :stat]    [:and [:= :n.nspname :stat.schemaname] [:= :c.relname :stat.relname]]]
              :where     [:and [:= :c.relnamespace :n.oid]
                          ;; filter out system tables
                          [(keyword "!~") :n.nspname "^pg_"] [:<> :n.nspname "information_schema"]

src/metabase/driver/sql_jdbc.clj

@@ -74,8 +74,8 @@
     ((get-method driver/db-default-timezone :metabase.driver/driver) driver database)))
 
 (defmethod driver/execute-reducible-query :sql-jdbc
-  [driver query chans respond]
-  (sql-jdbc.execute/execute-reducible-query driver query chans respond))
+  [driver query context respond]
+  (sql-jdbc.execute/execute-reducible-query driver query context respond))
 
 (defmethod driver/notify-database-updated :sql-jdbc
   [_ database]

src/metabase/sync/interface.clj

@@ -8,12 +8,15 @@
    [metabase.util.malli.schema :as ms]))
 
 (mr/def ::DatabaseMetadataTable
-  [:map
-   [:name           ::lib.schema.common/non-blank-string]
-   [:schema         [:maybe ::lib.schema.common/non-blank-string]]
-   [:require-filter {:optional true} :boolean]
+  [:map {:closed true}
+   [:name                                     ::lib.schema.common/non-blank-string]
+   [:schema                                   [:maybe ::lib.schema.common/non-blank-string]]
+   ;; for databases that store an estimated row count in system tables (e.g: postgres)
+   [:estimated_row_count     {:optional true} [:maybe :int]]
+   ;; for databases that support forcing query to include a filter (e.g: partitioned table on bigquery)
+   [:database_require_filter {:optional true} [:maybe :boolean]]
    ;; `:description` in this case should be a column/remark on the Table, if there is one.
-   [:description    {:optional true} [:maybe :string]]])
+   [:description             {:optional true} [:maybe :string]]])
 
 (def DatabaseMetadataTable
   "Schema for the expected output of `describe-database` for a Table."

src/metabase/sync/sync_metadata.clj

@@ -57,9 +57,9 @@
    (sync-util/create-sync-step "sync-fks" sync-fks/sync-fks! sync-fks-summary)
    ;; Sync index info if the database supports it
    (sync-util/create-sync-step "sync-indexes" sync-indexes/maybe-sync-indexes! sync-indexes-summary)
-   ;; finally, sync the metadata metadata table if it exists.
+   ;; Sync the metadata metadata table if it exists.
    (sync-util/create-sync-step "sync-metabase-metadata" #(metabase-metadata/sync-metabase-metadata! % db-metadata))
-   ;; Now sync the table privileges
+   ;; Now sync the table privileges if the database enables it
    (sync-util/create-sync-step "sync-table-privileges" sync-table-privileges/sync-table-privileges!)])
 
 (mu/defn sync-db-metadata!

src/metabase/sync/sync_metadata/tables.clj

@@ -156,7 +156,7 @@
   [table-metadata :- i/DatabaseMetadataTable
    metabase-table :- (ms/InstanceOf :model/Table)]
   (log/infof "Updating table metadata for %s" (sync-util/name-for-logging metabase-table))
-  (let [to-update-keys [:description :database_require_filter]
+  (let [to-update-keys [:description :database_require_filter :estimated_row_count]
         old-table      (select-keys metabase-table to-update-keys)
         new-table      (select-keys (merge
                                      (zipmap to-update-keys (repeat nil))
@@ -193,10 +193,10 @@
         (remove metabase-metadata/is-metabase-metadata-table?)
         (:tables db-metadata)))
 
-(mu/defn ^:private db->our-metadata :- [:set i/DatabaseMetadataTable]
+(mu/defn ^:private db->our-metadata :- [:set (ms/InstanceOf :model/Table)]
   "Return information about what Tables we have for this DB in the Metabase application DB."
   [database :- i/DatabaseInstance]
-  (set (t2/select [:model/Table :id :name :schema :description :database_require_filter]
+  (set (t2/select [:model/Table :id :name :schema :description :database_require_filter :estimated_row_count]
                   :db_id  (u/the-id database)
                   :active true)))
 

test/metabase/driver/postgres_test.clj

@@ -251,8 +251,23 @@
                  (mt/run-mbql-query people
                    {:fields [$name $bird_id->birds.name]}))))))))
 
-(defn- default-table-result [table-name]
-  {:name table-name, :schema "public", :description nil})
+(defn- default-table-result
+  ([table-name]
+   (default-table-result table-name {}))
+  ([table-name opts]
+   (merge
+    {:name table-name :schema "public" :description nil
+     ;; estimated-row-count is estimated, so the value can't be known for sure "during" test without
+     ;; VACUUM-ing. So for tests that don't concern the exact value of estimated-row-count, use schema instead
+     :estimated_row_count (mt/malli=? [:maybe :int])}
+    opts)))
+
+(defn- describe-database->tables
+  "Returns a seq of tables sorted by name from calling [[driver/describe-database]]."
+  [& args]
+  (->> (apply driver/describe-database args)
+      :tables
+      (sort-by :name)))
 
 (deftest materialized-views-test
   (mt/test-driver :postgres
@@ -264,9 +279,9 @@
                        ["DROP MATERIALIZED VIEW IF EXISTS test_mview;
                        CREATE MATERIALIZED VIEW test_mview AS
                        SELECT 'Toucans are the coolest type of bird.' AS true_facts;"])
-        (t2.with-temp/with-temp [Database database {:engine :postgres, :details (assoc details :dbname "materialized_views_test")}]
-          (is (= {:tables #{(default-table-result "test_mview")}}
-                 (driver/describe-database :postgres database))))))))
+        (mt/with-temp [:model/Database database {:engine :postgres, :details (assoc details :dbname "materialized_views_test")}]
+          (is (=? [(default-table-result "test_mview" {:estimated_row_count (mt/malli=? int?)})]
+                  (describe-database->tables :postgres database))))))))
 
 (deftest foreign-tables-test
   (mt/test-driver :postgres
@@ -276,7 +291,8 @@
         ;; You need to set `MB_POSTGRESQL_TEST_USER` in order for this to work apparently.
         ;;
         ;; make sure that the details include optional stuff like `:user`. Otherwise the test is going to FAIL. You can
-        ;; set it at run time from the REPL using [[mt/db-test-env-var!]].
+        ;; set it at run time from the REPL using [[mt/db-test-env-var!]]
+        ;; (mt/db-test-env-var! :postgresql :user "postgres").
         (is (mc/coerce [:map
                         [:port :int]
                         [:host :string]
@@ -297,8 +313,9 @@
                                 OPTIONS (user '" (:user details) "');
                               GRANT ALL ON public.local_table to PUBLIC;")])
         (t2.with-temp/with-temp [Database database {:engine :postgres, :details (assoc details :dbname "fdw_test")}]
-          (is (= {:tables (set (map default-table-result ["foreign_table" "local_table"]))}
-                 (driver/describe-database :postgres database))))))))
+          (is (=? [(default-table-result "foreign_table")
+                   (default-table-result "local_table" {:estimated_row_count (mt/malli=? int?)})]
+                  (describe-database->tables :postgres database))))))))
 
 (deftest recreated-views-test
   (mt/test-driver :postgres
@@ -363,8 +380,8 @@
                 ;; now sync the DB
                 (sync!)
                 ;; all three of these tables should appear in the metadata (including, importantly, the "main" table)
-                (is (= {:tables (set (map default-table-result ["part_vals" "part_vals_0" "part_vals_1"]))}
-                       (driver/describe-database :postgres database)))))
+                (is (=? (map default-table-result ["part_vals" "part_vals_0" "part_vals_1"])
+                        (describe-database->tables :postgres database)))))
             (log/warn
              (u/format-color
               'yellow
@@ -1378,7 +1395,9 @@
                            (is (= (into #{} (#'sql-jdbc.describe-database/jdbc-get-tables
                                              :postgres (.getMetaData conn) nil schema-pattern table-pattern
                                              ["TABLE" "PARTITIONED TABLE" "VIEW" "FOREIGN TABLE" "MATERIALIZED VIEW"]))
-                                  (into #{} (#'postgres/get-tables (mt/db) schemas tables)))))]
+                                  (into #{} (map #(dissoc % :estimated_row_count))
+                                        (#'postgres/get-tables (mt/db) schemas tables)))))]
+
              (doseq [stmt ["CREATE TABLE public.table (id INTEGER, type TEXT);"
                            "CREATE UNIQUE INDEX idx_table_type ON public.table(type);"
                            "CREATE TABLE public.partition_table (id INTEGER) PARTITION BY RANGE (id);"

test/metabase/sync/sync_metadata/tables_test.clj

@@ -3,7 +3,9 @@
   (:require
    [clojure.java.jdbc :as jdbc]
    [clojure.test :refer :all]
+   [metabase.driver :as driver]
    [metabase.driver.sql-jdbc.connection :as sql-jdbc.conn]
+   [metabase.driver.sql-jdbc.execute :as sql-jdbc.execute]
    [metabase.models :refer [Database Table]]
    [metabase.sync :as sync]
    [metabase.sync.sync-metadata.tables :as sync-tables]
@@ -11,6 +13,7 @@
    [metabase.test.data.interface :as tx]
    [metabase.test.data.sql :as sql.tx]
    [metabase.util :as u]
+   [next.jdbc :as next.jdbc]
    [toucan2.core :as t2]))
 
 (tx/defdataset db-with-some-cruft
@@ -60,3 +63,18 @@
         (is (=? {:active true
                  :description "added comment"}
                 (t2/select-one :model/Table (mt/id :user))))))))
+
+(deftest sync-estimated-row-count-test
+  (mt/test-driver :postgres
+    (testing "Can sync row count"
+      (mt/dataset test-data
+        ;; row count is estimated so we VACUUM so the statistic table is updated before syncing
+        (sql-jdbc.execute/do-with-connection-with-options
+         driver/*driver*
+         (mt/db)
+         nil
+         (fn [conn]
+           (next.jdbc/execute! conn ["VACUUM;"])))
+        (sync/sync-database! (mt/db) {:scan :schema})
+        (is (= 100
+               (t2/select-one-fn :estimated_row_count :model/Table (mt/id :venues))))))))