捕获受动态 sql 影响的行数？

Question

我正在尝试从 plpgsql 函数中的 QUERY EXEUTE 获取返回，以便能够检查动态更新查询影响了多少行。我的用例是在插入或更新到动态设置的表时将事件（带有自定义有效负载）添加到单独的表中。因为我的事件有一个自定义有效负载，所以我无法使用数据库触发器（例如，插入前触发器）。作为一个简化的例子，假设我有这个表：

CREATE TABLE users (user_id text primary key, name text)

这是我的简化事件表：

CREATE TABLE events(event_id text primary key, payload json)

这是我的简化函数：

CREATE OR REPLACE FUNCTION my_function(_rowtype anyelement, q text, payload jsonb)
    RETURNS SETOF anyelement AS
$func$
DECLARE
    event_id text;
BEGIN
    SELECT jsonb_object_field_text (payload, 'id')::text INTO STRICT event_id;
    execute format('insert into event(event_id, payload) values ($1, $2)') using event_id, payload;
    RETURN QUERY EXECUTE format('%s', q);
END
$func$ LANGUAGE plpgsql;

我们的目标是让这项工作与有人在交易中创建这些工作完全相同。在插入的伪代码中：

BEGIN
insert into events(id, payload) values($1, $2)
insert into users(columns) values(<any values>)
COMMIT

同样的更新：

BEGIN
insert into events(id, payload) values($1, $2)
result, error := query(`update users set name = 'hello' where id = 'Not Exists Thus No Rows Modified'`);
if result.rowsAffected() == 0 {
   ROLLBACK
}
COMMIT

函数my_function 几乎可以工作，除了一个边缘情况：当更新实际上不影响任何行时。 例如，这有效：

select * from my_function(NULL::users, 
       'insert into users(id,name) values('u1', ''a2'') returning *',
       payload => '{"id": "e1", "custom": "s1", "field": "2019-10-12T07:20:50.52Z"}')      

正如预期的那样，完成此操作后，用户表中的一行和事件表都被创建了。 失败的情况如下：

select * from my_function(NULL::users, 
       'update users set name = ''hello'' where user_id = ''NotExists'' returning *',
       payload => '{"id": "e2", "custom": "s3", "field": "2019-10-12T07:20:50.52Z"}')     

在这里，在 events 表中创建了一行（我的目标是不应该创建它）。 我知道这种方法并不优雅，而且我知道这很容易受到 SQL 注入的影响。我很想提出更好的方法来解决这个问题（包括取消我们现在正在做的事情）。但要直接回答这个问题，我希望存储 QUERY EXECUTE 的结果，检查是否有任何行受到影响，并引发错误，这样就不会出现在事件表中创建行的情况用户表中没有真正的相应变化。用户表只是一个例子，一般来说，它可以是任何动态设置的表。

Answer 1

返回查询不需要走到函数的末尾，它只说：“这个查询的结果是结果集的一部分”。

因此，您可以使用 RETURN QUERY，请求 FOUND 并采取相应措施。这是您的函数修改后以这种方式工作：

CREATE OR REPLACE FUNCTION public.my_function(_rowtype anyelement, q text, payload jsonb)
 RETURNS SETOF anyelement
 LANGUAGE plpgsql
AS $function$
DECLARE
    event_id text;
BEGIN
    SELECT jsonb_object_field_text (payload, 'id')::text INTO STRICT event_id;

    RETURN QUERY EXECUTE format('%s', q);
    IF FOUND THEN
        execute format('insert into events(event_id, payload) values ($1, $2)') using event_id, payload;
    END IF;

    RETURN;
END
$function$

PD：也许您也可以使用转换表 OLD 和 NEW（从 v10 开始提供，https://www.postgresql.org/docs/10/sql-createtrigger.html）使用 FOR EACH STATEMENT 触发器来解决您的问题

Answer 2

CREATE OR REPLACE FUNCTION my_function(_rowtype anyelement, q text, payload jsonb)
  RETURNS SETOF anyelement
  LANGUAGE plpgsql AS
$func$
BEGIN
   RETURN QUERY EXECUTE q;

   IF NOT FOUND THEN 
      RETURN;  -- nothing happened yet, we can exit silently.
      -- Or you WANT an error for this case. Then do this instead:
      -- RAISE EXCEPTION 'Query passed in parameter "q" did not affect any rows. Doing nothing!';
   END IF;
   
   INSERT INTO event(event_id, payload)
   VALUES (payload->>'id', payload);
END
$func$;

如前所述，RETURN QUERY 不会从函数返回。 The manual:

RETURN NEXT 和 RETURN QUERY 实际上并没有从 函数——它们只是将零或多行附加到函数的 结果集。然后继续执行中的下一条语句 PL/pgSQL 函数。作为连续的RETURN NEXT 或RETURN QUERY 命令被执行，结果集被建立。最后一个RETURN， 它应该没有参数，导致控制退出函数（或 你可以让控制到达函数的末尾）。

在手册那一章的底部有一个code example 用于您的案例。实际上，来自我。来自这里：

• FUNCTION syntax error

建议使用GET DIAGNOSTICS 而不是更简单的FOUND。确实EXECUTE 没有设置FOUND 的状态。但是RETURN QUERY 可以。所以继续使用更简单的FOUND。相关：

• Dynamic SQL (EXECUTE) as condition for IF statement

您的原始文件中有两次 format()。虽然这通常对动态 SQL 非常有用，但在您的情况下它是无用的。 EXECUTE format('%s', q) 与 EXECUTE q 完全相同，但增加了成本。在传递用户输入时，两者都为 SQL 注入打开了大门。

虽然事务很可能会回滚，但请从关键步骤开始，然后再执行其余步骤。避免浪费工作。所以我将执行q 移到顶部。 假设它不依赖于“有效负载”行，现在稍后插入。

另外，INSERT INTO events 可以是纯 SQL。那里没有什么动态的。不需要format() 或EXECUTE。

最后，假设您的jsonb_object_field_text (payload, 'id')::text 只是payload->>'id' 的一种奇特方式。不需要额外的变量和另一个SELECT INTO。

SQL 注入警告

将用户输入（示例中的参数q）转换为动态执行的代码是最直接的SQL 注入漏洞。我不想那样做被夹在我的内衣里。