优化 Postgres 中的行级安全表达式

Question

我有一个包含数百万行的表，分为几个类别（比如数字 1-5）。访问数据库的应用程序使用一个单一的数据库帐户。但是，该应用程序有自己的用户帐户，每个应用程序用户只能访问某些类别。因此，允许的类别列表通过会话变量传递给数据库：

SET SESSION mydb.allowed_categories = '1,3,5';

我使用 RLS 根据会话变量过滤行：

CREATE POLICY table_select_policy ON big_table
FOR SELECT
USING (ARRAY[category] && string_to_array(current_setting('mydb.allowed_categories'),',')::int[]));

问题在于，使用这种方法，RLS 行过滤需要花费大量时间。另一方面，当我通过以下方式进行实验时：

CREATE POLICY table_select_policy ON big_table
FOR SELECT
USING (category = 1 OR category = 3 OR category = 5);

RLS 过滤几乎快了 10 倍。当然，这种硬编码是行不通的，因为我想在应用程序中动态更改允许的类别列表。

这 category列有一个 btree 索引，但是由于类别的数量相当少，查询计划器总是倾向于对 RLS 过滤器进行顺序扫描。

所以我的问题是 - 有没有办法优化 RLS 表达式，使其至少更接近硬编码方法？你会建议一个不同的解决方案吗？该应用程序有很多用户，所以我不想为每个用户都创建一个数据库帐户。

Answer 1

乍一看，由于current_setting和string_to_array函数分别是稳定和不可变的，并且category列上有一个索引，因此以下条件可以解决问题。

CREATE POLICY table_select_policy ON big_table
FOR SELECT
USING (category = ANY(string_to_array(current_setting('mydb.allowed_categories'),',')::int[])));

然而，真正的问题与条件的形式无关。如果您查看 PostgreSQL 系统目录并找到条件中使用的那两个函数，您可能会注意到这两个函数的cost参数都设置为1。这使得优化器假设在执行位图堆扫描时调用这些函数来重新检查条件是廉价的。

为了说明这一点，请考虑下big_table表。

CREATE TABLE big_table AS
SELECT c AS category
FROM generate_series(1, 1000000) g,
     generate_series(1, 10) c;

CREATE INDEX big_table_category_idx ON big_table (category);

ANALYZE big_table;

使用下面的查询查询表会产生以下计划（查看索引重新检查删除的行数）。

EXPLAIN ANALYZE SELECT * FROM big_table WHERE category = ANY(string_to_array(current_setting('mydb.allowed_categories'), ',')::int[]);

"Bitmap Heap Scan on big_table  (cost=52478.56..195418.17 rows=3036665 width=4) (actual time=166.613..9273.010 rows=3000000 loops=1)"
"  Recheck Cond: (category = ANY ((string_to_array(current_setting('mydb.allowed_categories'::text), ','::text))::integer[]))"
"  Rows Removed by Index Recheck: 5209365"
"  ->  Bitmap Index Scan on big_table_category_idx  (cost=0.00..51719.39 rows=3036665 width=0) (actual time=164.782..164.782 rows=3000000 loops=1)"
"        Index Cond: (category = ANY ((string_to_array(current_setting('mydb.allowed_categories'::text), ','::text))::integer[]))"
"Total runtime: 9341.568 ms"

为避免这种情况，您有两种选择。第一个是增加work_mem服务器配置中的参数。这将允许服务器将完整的位图存储在内存中，并在重新检查条件时节省大量时间。下面的计划是在work_mem参数设置为时获得的1000M。

EXPLAIN ANALYZE SELECT * FROM big_table WHERE category = ANY(string_to_array(current_setting('mydb.allowed_categories'), ',')::int[]);


"Bitmap Heap Scan on big_table  (cost=52478.56..195418.17 rows=3036665 width=4) (actual time=193.613..449.385 rows=3000000 loops=1)"
"  Recheck Cond: (category = ANY ((string_to_array(current_setting('mydb.allowed_categories'::text), ','::text))::integer[]))"
"  ->  Bitmap Index Scan on big_table_category_idx  (cost=0.00..51719.39 rows=3036665 width=0) (actual time=184.858..184.858 rows=3000000 loops=1)"
"        Index Cond: (category = ANY ((string_to_array(current_setting('mydb.allowed_categories'::text), ','::text))::integer[]))"
"Total runtime: 513.528 ms"

第二种选择是将查询条件包装在一个“昂贵”的函数中。

CREATE OR REPLACE FUNCTION my_categories()
RETURNS int[] AS $$
BEGIN
  RETURN string_to_array(current_setting('mydb.allowed_categories'), ',')::int[];
END;
$$ LANGUAGE plpgsql STABLE COST 100000;

EXPLAIN ANALYZE SELECT * FROM big_table WHERE category = ANY(my_categories());

"Index Only Scan using big_table_category_idx on big_table  (cost=250.44..9363945.19 rows=3036665 width=4) (actual time=0.035..577.094 rows=3000000 loops=1)"
"  Index Cond: (category = ANY (my_categories()))"
"  Heap Fetches: 3000000"
"Total runtime: 642.522 ms"

此选项还将使优化器最终执行索引扫描而不是位图扫描，这可能会变慢一点。