Refactor columnar tests to use helper functions for storage ID retrieval and stripe counting

Author: m3hm3t

src/test/regress/expected/columnar_create.out

@@ -214,8 +214,9 @@ SELECT COUNT(*) FROM columnar_temp WHERE i < 5;
      4
 (1 row)
 
-SELECT columnar.get_storage_id(oid) AS columnar_temp_storage_id
-FROM pg_class WHERE relname='columnar_temp' \gset
+SELECT COALESCE(columnar_test_helpers.get_storage_id_if_visible('columnar_temp'::regclass), 0)
+       AS columnar_temp_storage_id
+\gset
 BEGIN;
   DROP TABLE columnar_temp;
   -- show that we drop stripes properly

src/test/regress/expected/columnar_drop.out

@@ -12,28 +12,35 @@
 -- 'postgres' directory is excluded from comparison to have the same result.
 -- store postgres database oid
 SELECT oid postgres_oid FROM pg_database WHERE datname = 'postgres' \gset
-SELECT count(distinct storage_id) AS columnar_stripes_before_drop FROM columnar.stripe \gset
+SELECT columnar_test_helpers.safe_visible_stripe_count() AS columnar_stripes_before_drop \gset
 -- DROP columnar tables
 DROP TABLE contestant;
 DROP TABLE contestant_compressed;
 -- make sure DROP deletes metadata
-SELECT :columnar_stripes_before_drop - count(distinct storage_id) FROM columnar.stripe;
- ?column?
+SELECT :columnar_stripes_before_drop - columnar_test_helpers.safe_visible_stripe_count() AS dropped;
+ dropped
 ---------------------------------------------------------------------
-        2
+       2
 (1 row)
 
 -- Create a columnar table under a schema and drop it.
 CREATE SCHEMA test_schema;
 CREATE TABLE test_schema.test_table(data int) USING columnar;
 INSERT INTO test_schema.test_table VALUES (1);
-SELECT count(*) AS columnar_stripes_before_drop FROM columnar.stripe \gset
+-- capture schema OID (via catalog), then save it to a var
+SELECT oid AS test_schema_oid FROM pg_catalog.pg_namespace WHERE nspname = 'test_schema' \gset
+-- baseline, scoped to that schema OID
+SELECT columnar_test_helpers.safe_visible_stripe_count_by_oid(:test_schema_oid) AS columnar_stripes_before_drop \gset
+-- drop the schema
 DROP SCHEMA test_schema CASCADE;
 NOTICE:  drop cascades to table test_schema.test_table
-SELECT :columnar_stripes_before_drop - count(distinct storage_id) FROM columnar.stripe;
- ?column?
+-- delta using the saved OID (works even though the schema name no longer exists)
+SELECT :columnar_stripes_before_drop
+     - columnar_test_helpers.safe_visible_stripe_count_by_oid(:test_schema_oid)
+  AS dropped;
+ dropped
 ---------------------------------------------------------------------
-        1
+       1
 (1 row)
 
 SELECT current_database() datname \gset

src/test/regress/expected/columnar_first_row_number.out

@@ -1,5 +1,5 @@
 CREATE SCHEMA columnar_first_row_number;
-SET search_path tO columnar_first_row_number;
+SET search_path to columnar_first_row_number;
 CREATE TABLE col_table_1 (a int) USING columnar;
 INSERT INTO col_table_1 SELECT i FROM generate_series(1, 10) i;
 BEGIN;
@@ -9,9 +9,13 @@ ROLLBACK;
 INSERT INTO col_table_1 SELECT i FROM generate_series(1, 12) i;
 ALTER TABLE col_table_1 SET (columnar.stripe_row_limit = 1000);
 INSERT INTO col_table_1 SELECT i FROM generate_series(1, 2350) i;
-SELECT row_count, first_row_number FROM columnar.stripe a
-WHERE a.storage_id = columnar.get_storage_id('col_table_1'::regclass)
-ORDER BY stripe_num;
+WITH s AS (
+  SELECT columnar.get_storage_id('col_table_1'::regclass) AS sid
+)
+SELECT a.row_count, a.first_row_number
+FROM columnar_internal.stripe a, s
+WHERE a.storage_id = s.sid
+ORDER BY a.stripe_num;
  row_count | first_row_number
 ---------------------------------------------------------------------
         10 |                1
@@ -39,9 +43,13 @@ BEGIN;
   INSERT INTO col_table_1 SELECT i FROM generate_series(1, 16) i;
 COMMIT;
 -- show that we start with first_row_number=1 after TRUNCATE
-SELECT row_count, first_row_number FROM columnar.stripe a
-WHERE a.storage_id = columnar.get_storage_id('col_table_1'::regclass)
-ORDER BY stripe_num;
+WITH s AS (
+  SELECT columnar.get_storage_id('col_table_1'::regclass) AS sid
+)
+SELECT a.row_count, a.first_row_number
+FROM columnar_internal.stripe a, s
+WHERE a.storage_id = s.sid
+ORDER BY a.stripe_num;
  row_count | first_row_number
 ---------------------------------------------------------------------
         32 |                1

src/test/regress/expected/columnar_indexes.out

@@ -599,8 +599,10 @@ begin;
   -- i) abort before flushing
 rollback;
 insert into uniq select generate_series(1,100);
-SELECT COUNT(*)=1 FROM columnar.stripe cs
-WHERE cs.storage_id = columnar.get_storage_id('columnar_indexes.uniq'::regclass);
+SELECT COUNT(*) = 1
+FROM columnar_internal.stripe AS st
+WHERE st.storage_id =
+      columnar.get_storage_id('columnar_indexes.uniq'::regclass);
  ?column?
 ---------------------------------------------------------------------
  t
@@ -618,8 +620,10 @@ begin;
 
 rollback;
 insert into uniq select generate_series(1,100);
-SELECT COUNT(*)=1 FROM columnar.stripe cs
-WHERE cs.storage_id = columnar.get_storage_id('columnar_indexes.uniq'::regclass);
+SELECT COUNT(*) = 1
+FROM columnar_internal.stripe AS st
+WHERE st.storage_id =
+      columnar.get_storage_id('columnar_indexes.uniq'::regclass);
  ?column?
 ---------------------------------------------------------------------
  t
@@ -632,8 +636,10 @@ begin;
   -- i) abort before flushing
 rollback;
 insert into uniq select generate_series(1,100);
-SELECT COUNT(*)=1 FROM columnar.stripe cs
-WHERE cs.storage_id = columnar.get_storage_id('columnar_indexes.uniq'::regclass);
+SELECT COUNT(*) = 1
+FROM columnar_internal.stripe AS st
+WHERE st.storage_id =
+      columnar.get_storage_id('columnar_indexes.uniq'::regclass);
  ?column?
 ---------------------------------------------------------------------
  t
@@ -651,8 +657,10 @@ begin;
 
 rollback;
 insert into uniq select generate_series(1,100);
-SELECT COUNT(*)=1 FROM columnar.stripe cs
-WHERE cs.storage_id = columnar.get_storage_id('columnar_indexes.uniq'::regclass);
+SELECT COUNT(*) = 1
+FROM columnar_internal.stripe AS st
+WHERE st.storage_id =
+      columnar.get_storage_id('columnar_indexes.uniq'::regclass);
  ?column?
 ---------------------------------------------------------------------
  t
@@ -668,17 +676,23 @@ begin;
   -- values. So the following insert should just work.
   insert into uniq select generate_series(1,100);
   -- didn't flush anything yet, but should see the in progress stripe-write
-  SELECT stripe_num, first_row_number, row_count FROM columnar.stripe cs
-  WHERE cs.storage_id = columnar.get_storage_id('columnar_indexes.uniq'::regclass);
+  SELECT st.stripe_num, st.first_row_number, st.row_count
+  FROM columnar_internal.stripe AS st
+  WHERE st.storage_id =
+        columnar.get_storage_id('columnar_indexes.uniq'::regclass)
+  ORDER BY st.stripe_num;
  stripe_num | first_row_number | row_count
 ---------------------------------------------------------------------
           2 |           150001 |         0
 (1 row)
 
 commit;
 -- should have completed the stripe reservation
-SELECT stripe_num, first_row_number, row_count FROM columnar.stripe cs
-WHERE cs.storage_id = columnar.get_storage_id('columnar_indexes.uniq'::regclass);
+SELECT st.stripe_num, st.first_row_number, st.row_count
+FROM columnar_internal.stripe AS st
+WHERE st.storage_id =
+      columnar.get_storage_id('columnar_indexes.uniq'::regclass)
+ORDER BY st.stripe_num;
  stripe_num | first_row_number | row_count
 ---------------------------------------------------------------------
           2 |           150001 |       100

src/test/regress/expected/columnar_insert.out

@@ -194,8 +194,14 @@ select
              2 |             0 |                  6 |              750001
 (1 row)
 
-SELECT relname, stripe_num, chunk_group_count, row_count FROM columnar.stripe a, pg_class b
-WHERE columnar.get_storage_id(b.oid)=a.storage_id AND relname = 'zero_col'
+SELECT c.relname, st.stripe_num, st.chunk_group_count, st.row_count
+FROM pg_catalog.pg_class      AS c
+JOIN pg_catalog.pg_namespace  AS n ON n.oid = c.relnamespace
+JOIN pg_catalog.pg_am         AS am ON am.oid = c.relam AND am.amname = 'columnar'
+JOIN LATERAL (SELECT columnar.get_storage_id(c.oid) AS sid) AS s ON true
+JOIN columnar_internal.stripe AS st ON st.storage_id = s.sid
+WHERE c.relname = 'zero_col'
+  AND (c.relpersistence <> 't' OR c.relnamespace = pg_catalog.pg_my_temp_schema())
 ORDER BY 1,2,3,4;
  relname  | stripe_num | chunk_group_count | row_count
 ---------------------------------------------------------------------
@@ -206,15 +212,27 @@ ORDER BY 1,2,3,4;
  zero_col |          5 |                 1 |        64
 (5 rows)
 
-SELECT relname, stripe_num, value_count FROM columnar.chunk a, pg_class b
-WHERE columnar.get_storage_id(b.oid)=a.storage_id AND relname = 'zero_col'
+SELECT c.relname, ch.stripe_num, ch.value_count
+FROM pg_catalog.pg_class       AS c
+JOIN pg_catalog.pg_namespace   AS n ON n.oid = c.relnamespace
+JOIN pg_catalog.pg_am          AS am ON am.oid = c.relam AND am.amname = 'columnar'
+JOIN LATERAL (SELECT columnar.get_storage_id(c.oid) AS sid) AS s ON true
+JOIN columnar_internal.chunk   AS ch ON ch.storage_id = s.sid
+WHERE c.relname = 'zero_col'
+  AND (c.relpersistence <> 't' OR c.relnamespace = pg_catalog.pg_my_temp_schema())
 ORDER BY 1,2,3;
  relname | stripe_num | value_count
 ---------------------------------------------------------------------
 (0 rows)
 
-SELECT relname, stripe_num, chunk_group_num, row_count FROM columnar.chunk_group a, pg_class b
-WHERE columnar.get_storage_id(b.oid)=a.storage_id AND relname = 'zero_col'
+SELECT c.relname, cg.stripe_num, cg.chunk_group_num, cg.row_count
+FROM pg_catalog.pg_class            AS c
+JOIN pg_catalog.pg_namespace        AS n ON n.oid = c.relnamespace
+JOIN pg_catalog.pg_am               AS am ON am.oid = c.relam AND am.amname = 'columnar'
+JOIN LATERAL (SELECT columnar.get_storage_id(c.oid) AS sid) AS s ON true
+JOIN columnar_internal.chunk_group  AS cg ON cg.storage_id = s.sid
+WHERE c.relname = 'zero_col'
+  AND (c.relpersistence <> 't' OR c.relnamespace = pg_catalog.pg_my_temp_schema())
 ORDER BY 1,2,3,4;
  relname  | stripe_num | chunk_group_num | row_count
 ---------------------------------------------------------------------

src/test/regress/expected/columnar_matview.out

@@ -63,13 +63,14 @@ SELECT * FROM t_view a ORDER BY a;
 -- verify that we have created metadata entries for the materialized view
 SELECT columnar.get_storage_id(oid) AS storageid
 FROM pg_class WHERE relname='t_view' \gset
-SELECT count(*) FROM columnar.stripe WHERE storage_id=:storageid;
+SELECT columnar.get_storage_id('t_view'::regclass) AS storageid \gset
+SELECT count(*) FROM columnar_internal.stripe WHERE storage_id = :storageid;
  count
 ---------------------------------------------------------------------
      1
 (1 row)
 
-SELECT count(*) FROM columnar.chunk WHERE storage_id=:storageid;
+SELECT count(*) FROM columnar_internal.chunk  WHERE storage_id = :storageid;
  count
 ---------------------------------------------------------------------
      3
@@ -78,13 +79,13 @@ SELECT count(*) FROM columnar.chunk WHERE storage_id=:storageid;
 DROP TABLE t CASCADE;
 NOTICE:  drop cascades to materialized view t_view
 -- dropping must remove metadata
-SELECT count(*) FROM columnar.stripe WHERE storage_id=:storageid;
+SELECT count(*) FROM columnar_internal.stripe WHERE storage_id = :storageid;
  count
 ---------------------------------------------------------------------
      0
 (1 row)
 
-SELECT count(*) FROM columnar.chunk WHERE storage_id=:storageid;
+SELECT count(*) FROM columnar_internal.chunk  WHERE storage_id = :storageid;
  count
 ---------------------------------------------------------------------
      0

src/test/regress/expected/columnar_recursive.out

@@ -11,10 +11,19 @@ $$ LANGUAGE SQL;
 INSERT INTO t2 SELECT i, f(i) FROM generate_series(1, 5) i;
 -- there are no subtransactions, so above statement should batch
 -- INSERTs inside the UDF and create on stripe per table.
-SELECT relname, count(*) FROM columnar.stripe a, pg_class b
-WHERE columnar.get_storage_id(b.oid)=a.storage_id AND relname IN ('t1', 't2')
-GROUP BY relname
-ORDER BY relname;
+WITH rels(rel) AS (
+  VALUES ('t1'::regclass), ('t2'::regclass)
+),
+sids AS (
+  SELECT rel, columnar.get_storage_id(rel) AS sid
+  FROM rels
+)
+SELECT c.relname, COUNT(*) AS count
+FROM columnar_internal.stripe st
+JOIN sids s ON st.storage_id = s.sid
+JOIN pg_catalog.pg_class c ON c.oid = s.rel
+GROUP BY c.relname
+ORDER BY c.relname;
  relname | count
 ---------------------------------------------------------------------
  t1      |     1

src/test/regress/expected/columnar_rollback.out

@@ -3,8 +3,15 @@
 --
 CREATE TABLE t(a int, b int) USING columnar;
 CREATE VIEW t_stripes AS
-SELECT * FROM columnar.stripe a, pg_class b
-WHERE a.storage_id = columnar.get_storage_id(b.oid) AND b.relname = 't';
+SELECT st.stripe_num, st.first_row_number, st.row_count, st.chunk_group_count
+FROM columnar_internal.stripe AS st
+JOIN LATERAL (
+  SELECT columnar.get_storage_id(
+           (quote_ident(current_schema())||'.t')::regclass
+         ) AS sid
+) s ON true
+WHERE st.storage_id = s.sid
+ORDER BY st.stripe_num;
 BEGIN;
 INSERT INTO t SELECT i, i+1 FROM generate_series(1, 10) i;
 ROLLBACK;

src/test/regress/expected/columnar_test_helpers.out

@@ -146,3 +146,76 @@ BEGIN
     RETURN NEXT;
   END LOOP;
 END; $$ language plpgsql;
+-- Count only stripes that belong to *visible* columnar tables:
+--  • permanent/unlogged tables from any schema
+--  • or temp tables from *my* session (exclude other sessions' temps)
+CREATE OR REPLACE FUNCTION safe_visible_stripe_count(ns regnamespace DEFAULT NULL)
+RETURNS bigint
+LANGUAGE sql
+STABLE
+AS $$
+WITH rels AS (
+  SELECT c.oid
+  FROM pg_catalog.pg_class c
+  JOIN pg_catalog.pg_am am ON am.oid = c.relam
+  WHERE am.amname = 'columnar'
+    AND (c.relpersistence <> 't' OR c.relnamespace = pg_catalog.pg_my_temp_schema())
+    AND ($1 IS NULL OR c.relnamespace = $1)
+),
+sids AS (
+  -- Safe to call: we filtered out other sessions' temp rels above
+  SELECT columnar.get_storage_id(oid) AS sid
+  FROM rels
+)
+SELECT COUNT(DISTINCT st.storage_id)
+FROM columnar_internal.stripe AS st
+JOIN sids ON st.storage_id = sids.sid;
+$$;
+CREATE OR REPLACE FUNCTION safe_visible_stripe_count_by_oid(ns_oid oid DEFAULT NULL)
+RETURNS bigint
+LANGUAGE sql
+STABLE
+AS $$
+WITH rels AS (
+  SELECT c.oid
+  FROM pg_catalog.pg_class c
+  JOIN pg_catalog.pg_am am ON am.oid = c.relam
+  WHERE am.amname = 'columnar'
+    AND (c.relpersistence <> 't' OR c.relnamespace = pg_catalog.pg_my_temp_schema())
+    AND ($1 IS NULL OR c.relnamespace = $1)
+),
+sids AS (
+  SELECT columnar.get_storage_id(oid) AS sid
+  FROM rels
+)
+SELECT COUNT(DISTINCT st.storage_id)
+FROM columnar_internal.stripe AS st
+JOIN sids ON st.storage_id = sids.sid;
+$$;
+CREATE OR REPLACE FUNCTION get_storage_id_if_visible(rel regclass)
+RETURNS bigint
+LANGUAGE sql STABLE AS $$
+  SELECT CASE
+           WHEN c.relpersistence = 't'
+                AND c.relnamespace <> pg_catalog.pg_my_temp_schema()
+             THEN NULL  -- other session’s temp → don’t touch
+           ELSE columnar.get_storage_id(c.oid)
+         END
+  FROM pg_catalog.pg_class c
+  WHERE c.oid = $1::oid
+$$;
+CREATE OR REPLACE FUNCTION columnar_test_helpers.chunks_for(rel regclass)
+RETURNS TABLE (
+  stripe_num int,
+  attr_num int,
+  chunk_group_num int,
+  minimum_value text,
+  maximum_value text
+)
+LANGUAGE sql STABLE AS $$
+  SELECT ch.stripe_num, ch.attr_num, ch.chunk_group_num,
+         ch.minimum_value, ch.maximum_value
+  FROM columnar_internal.chunk AS ch
+  WHERE ch.storage_id = columnar.get_storage_id($1)
+  ORDER BY ch.stripe_num, ch.attr_num, ch.chunk_group_num
+$$;

src/test/regress/expected/columnar_truncate.out

@@ -7,7 +7,8 @@ CREATE TABLE columnar_truncate_test_second (a int, b int) USING columnar;
 -- COMPRESSED
 CREATE TABLE columnar_truncate_test_compressed (a int, b int) USING columnar;
 CREATE TABLE columnar_truncate_test_regular (a int, b int);
-SELECT count(distinct storage_id) AS columnar_data_files_before_truncate FROM columnar.stripe \gset
+SELECT columnar_test_helpers.safe_visible_stripe_count()
+       AS columnar_data_files_before_truncate \gset
 INSERT INTO columnar_truncate_test select a, a from generate_series(1, 10) a;
 set columnar.compression = 'pglz';
 INSERT INTO columnar_truncate_test_compressed select a, a from generate_series(1, 10) a;