PostgreSQL 事务使用查询结果插入和查询另一个表然后返回原始查询结果

Question

我正在编写一个存储文件样本和 YARA 签名数据的应用程序。本质上，在单个事务中，我需要执行查询，在插入和另一个查询中引用这些结果，然后返回原始结果。我有三个与此讨论相关的表:

1. samples - 这是存储需要使用相关 YARA 签名扫描的文件信息的表格。
2. yararules - 存储 YARA 规则信息的表。
3. yaratracker - 一个跟踪到目前为止已处理的样本/规则对的表。

在单个事务中，应用程序需要:

1. 获取一批尚未处理的唯一样本/规则对。优选地，此查询将获取与单个样本关联的所有未处理规则(即，如果我要对样本运行 YARA 规则，我想运行该样本上尚未处理的所有 YARA 规则，以便我只需将示例加载到内存中一次)。
2. 从步骤 1 中找到的批处理中获取唯一的 id,sha256 列表。
3. 将第 1 步中的批处理插入到 yaraqueue 中，matchcount 列等于 0，complete 列设置为 false。

我可以使用下面的查询完成第 1 步，但我不知道如何引用这些结果来完成第 2 步。我尝试查看变量，但显然没有一个可以容纳多行。我研究过使用游标，但我似乎无法在后续命令中使用游标然后返回游标。

SELECT s.id,r.id
FROM sample s CROSS JOIN yararules r
WHERE r.status = 'Disabled' AND NOT EXISTS(
    SELECT 1 FROM yaratracker q
    WHERE q.sample_id = s.id AND q.rule_id = r.id
) 
ORDER BY s.id
LIMIT 1000;

相关的数据库模式如下所示。

CREATE TYPE samplelist AS ENUM ('Whitelist', 'Blacklist', 'Greylist', 'Unknown');

CREATE TABLE samples (
    id SERIAL PRIMARY KEY,
    md5 CHAR(32) NOT NULL,
    sha1 CHAR(40) NOT NULL,
    sha256 CHAR(64) NOT NULL,
    total INT NOT NULL,
    positives INT NOT NULL,
    list SAMPLELIST NOT NULL,
    filetype VARCHAR(16) NOT NULL,
    submitted TIMESTAMP WITH TIME ZONE NOT NULL,
    user_id SERIAL REFERENCES users;
);

CREATE UNIQUE INDEX md5_idx ON {0} (md5);
CREATE UNIQUE INDEX sha1_idx ON {0} (sha1);
CREATE UNIQUE INDEX sha256_idx ON {0} (sha256);

CREATE TYPE rulestatus AS ENUM ('Enabled', 'Disabled');

CREATE TABLE yararules (
    id SERIAL PRIMARY KEY,
    name VARCHAR(32) NOT NULL UNIQUE,
    description TEXT NOT NULL,
    rules TEXT NOT NULL,
    lastmodified TIMESTAMP WITH TIME ZONE NOT NULL,
    status rulestatus NOT NULL,
    user_id SERIAL REFERENCES users ON DELETE CASCADE
);

CREATE TABLE yaratracker (
    id SERIAL PRIMARY KEY,
    rule_id SERIAL REFERENCES yararules ON DELETE CASCADE,
    sample_id SERIAL REFERENCES sample ON DELETE CASCADE,
    matchcount INT NOT NULL,
    complete BOOL NOT NULL
);

CREATE INDEX composite_idx ON yaratracker (rule_id, sample_id);
CREATE INDEX complete_idx ON yaratracker (complete);

Answer 1

INSERT INTO target_table(a,b,c,...)
SELECT sid, rid, sha, ...
FROM (
        SELECT s.id AS sid
          ,r.id AS rid
          , s.sha256 AS sha
          , ...
          , ROW_NUMBER() OVER (PARTITION BY s.id) AS rn -- <<<--- HERE
        FROM sample s CROSS JOIN yararules r
        WHERE r.status = 'Disabled' AND NOT EXISTS(
         SELECT 1 FROM yaratracker q
          WHERE q.sample_id = s.id
            AND q.rule_id = r.id
        )
        ORDER BY s.id
        LIMIT 1000;
) src
WHERE src.rn = 1; -- <<<--- HERE

---

WHERE src.rn = 1 将限制交叉连接为每个 sample.id(id 和sha256 在 sample 表中是唯一的，因此选择一个唯一的 id 与选择一个唯一的 sha256 具有相同的效果)

永远不会生成完整的交叉连接结果；优化器足够聪明，可以将 WHERE rn=1 条件下推到子查询中。

注意:LIMIT 1000 可能应该被删除(或拉高到更高的级别)

---

如果您真的需要保存来自 CROSS JOIN 的结果，您可以使用 CTE 链(预期性能下降......)

---

WITH big AS (
        SELECT s.id AS sample_id
                ,r.id AS rule_id
                , s.sha256
                -- , ...
        , ROW_NUMBER() OVER (PARTITION BY s.id) AS rn -- <<<--- HERE
        FROM sample s 
        CROSS JOIN yararules r
        WHERE r.status = 'Disabled' AND NOT EXISTS(
                SELECT 1 FROM yaratracker q
                WHERE q.sample_id = s.id AND q.rule_id = r.id
                )
        )
, ins AS (
        INSERT INTO target_table(a,b,c,...)
        SELECT b.sample_id, b.rule_id, b.sha256 , ...
        FROM big b
        WHERE b.rn = 1; -- <<<--- HERE
        RETURNING *
        )
INSERT INTO yaratracker (rule_id, sample_id, matchcount, complete )
SELECT  b.sample_id, b.rule_id, 0, False
FROM big b
-- LEFT JOIN ins i ON i.a = b.sample_id AND i.b= b.rule_id
        ;

---

注意:yaratracker(rule_id,sample_id) 应该不是连续的，而只是普通整数，引用 yararules(id) 和 sample(id)