PostgreSQL：能否基于查询函数实现多表插入与测试数据生成？

我有一个用于聚合和查询关联表数据的PostgreSQL plpgsql函数aggregate_sales_data，代码如下：

CREATE OR REPLACE FUNCTION public.aggregate_sales_data(
    sales_rep_id integer)
    RETURNS TABLE(product_id integer, customer_id integer, total_sales_amount numeric, last_sale_date timestamp with time zone, region_id smallint) 
    LANGUAGE 'plpgsql'
    COST 100
    VOLATILE SECURITY DEFINER PARALLEL UNSAFE
    ROWS 1000

AS $BODY$
DECLARE
BEGIN
RETURN QUERY
SELECT DISTINCT ON (s.product_id) s.product_id, c.customer_id, SUM(s.amount) as total_sales_amount, MAX(s.sale_date) as last_sale_date, MAX(s.region_id) as region_id
FROM sales s
LEFT JOIN customers c ON s.customer_id = c.customer_id
WHERE s.sale_date > CURRENT_TIMESTAMP - INTERVAL '1 month' AND s.sales_rep_id = sales_rep_id
GROUP BY s.product_id, c.customer_id
ORDER BY s.product_id, c.customer_id ASC
LIMIT 500;
END;
$BODY$;

需求说明 #

• 需求1：希望该函数能返回额外的指定行（比如(2, 2, 3, NOW(), 1)），但不想手动向customers和sales表插入数据，目前想到用硬编码UNION的方式，想找更合理的方案。
• 需求2：想要实现函数insert_customers_and_products(123, 456, 100.00, '2024-03-06 10:00:00', 1)，执行后自动向customers和products表插入对应数据，且aggregate_sales_data能查询到该行数据，无需手动编写每个表的INSERT语句，能否借助初始的aggregate_sales_data函数实现？

针对需求1的优化方案 #

不需要用UNION硬拼结果，可以直接在原有查询后追加一行构造好的硬编码数据，代码更清晰，还能避免UNION可能引发的类型匹配问题。修改后的函数如下：

CREATE OR REPLACE FUNCTION public.aggregate_sales_data(
    sales_rep_id integer)
    RETURNS TABLE(product_id integer, customer_id integer, total_sales_amount numeric, last_sale_date timestamp with time zone, region_id smallint) 
    LANGUAGE 'plpgsql'
    COST 100
    VOLATILE SECURITY DEFINER PARALLEL UNSAFE
    ROWS 1000

AS $BODY$
DECLARE
BEGIN
-- 原有查询逻辑
RETURN QUERY
SELECT DISTINCT ON (s.product_id) s.product_id, c.customer_id, SUM(s.amount) as total_sales_amount, MAX(s.sale_date) as last_sale_date, MAX(s.region_id) as region_id
FROM sales s
LEFT JOIN customers c ON s.customer_id = c.customer_id
WHERE s.sale_date > CURRENT_TIMESTAMP - INTERVAL '1 month' AND s.sales_rep_id = sales_rep_id
GROUP BY s.product_id, c.customer_id
ORDER BY s.product_id, c.customer_id ASC
LIMIT 500;

-- 追加指定行数据，显式指定类型避免隐式转换问题
RETURN QUERY
SELECT 2::integer, 2::integer, 3::numeric, NOW()::timestamp with time zone, 1::smallint;
END;
$BODY$;

如果需要根据条件动态决定是否追加这行数据，可添加判断逻辑：

-- 示例：仅当销售代表ID为1时追加指定行
IF sales_rep_id = 1 THEN
    RETURN QUERY
    SELECT 2::integer, 2::integer, 3::numeric, NOW()::timestamp with time zone, 1::smallint;
END IF;

针对需求2的实现方案 #

aggregate_sales_data是查询类函数，无法直接用来实现插入逻辑，但可以基于它的表关联规则，编写一个新的插入函数，自动处理多表插入逻辑（注：需求中提到的products表未在原函数中关联，这里默认是笔误，若确实需要插入products，可参照customers的插入逻辑扩展）。

实现思路：

1. 先判断customers表中是否存在目标customer_id，不存在则插入
2. 向sales表插入销售数据（若需校验product_id在products表中存在，可添加类似判断）
3. 插入完成后调用aggregate_sales_data返回最新聚合结果

函数代码示例：

CREATE OR REPLACE FUNCTION public.insert_customers_and_sales(
    p_product_id integer,
    p_customer_id integer,
    p_amount numeric,
    p_sale_date timestamp with time zone,
    p_region_id smallint,
    p_sales_rep_id integer)
    RETURNS TABLE(product_id integer, customer_id integer, total_sales_amount numeric, last_sale_date timestamp with time zone, region_id smallint)
    LANGUAGE 'plpgsql'
    VOLATILE SECURITY DEFINER
AS $BODY$
BEGIN
    -- 插入customers（主键冲突则跳过）
    INSERT INTO customers(customer_id)
    VALUES(p_customer_id)
    ON CONFLICT (customer_id) DO NOTHING;

    -- 插入sales数据
    INSERT INTO sales(product_id, customer_id, amount, sale_date, region_id, sales_rep_id)
    VALUES(p_product_id, p_customer_id, p_amount, p_sale_date, p_region_id, p_sales_rep_id);

    -- 调用原聚合函数返回最新结果
    RETURN QUERY
    SELECT * FROM public.aggregate_sales_data(p_sales_rep_id);
END;
$BODY$;

调用示例：

SELECT * FROM insert_customers_and_sales(123, 456, 100.00, '2024-03-06 10:00:00', 1, 你的销售代表ID);

若确实需要插入products表，添加以下逻辑即可：

-- 插入products（主键冲突则跳过）
INSERT INTO products(product_id)
VALUES(p_product_id)
ON CONFLICT (product_id) DO NOTHING;

内容的提问来源于stack exchange，提问作者Tony