douyin-archive/lib/search-query.ts

// 搜索的 SQL 部分：搜索接口（分页取结果）和播放页（取完整的有序 ID 列表算上一条/下一条）
// 共用同一段匹配与排序逻辑，保证两边顺序完全一致。
import { Prisma } from "@prisma/client";
import { prisma } from "@/lib/prisma";
import { termToPgPattern } from "@/lib/search";

export type SearchInput = {
  /** 已转小写的关键词（parseSearchQuery 的结果） */
  terms: string[];
  excludes: string[];
  contentTag: string | null;
};

export type SearchRow = {
  awemeId: string;
  type: "video" | "image";
  score: number;
  inDesc: boolean;
  inTranscript: boolean;
  rawDesc: string;
  rawTranscript: string;
  contentTags: string[];
  total: number;
};

// 子串匹配：每个关键词都必须出现在「描述 + 语音转写」里（不区分大小写），
// 不依赖分词，所以输入任意片段都能搜到。数据量小，全表扫描足够快。
// 打分：描述命中权重更高，按出现次数累加（设上限防止长转写刷分）；同分按发布时间倒序。
// 没有关键词、只按标签筛选时，score 都是 0，等价于按时间排序。
function matchedCte({ terms, excludes, contentTag }: SearchInput) {
  const patterns = terms.map(termToPgPattern);
  const excludePatterns = excludes.map(termToPgPattern);
  return Prisma.sql`
    WITH params AS (
      SELECT ${patterns}::text[] AS terms, ${excludePatterns}::text[] AS excludes
    ),
    docs AS (
      SELECT
        v.aweme_id AS "awemeId",
        'video'::text AS type,
        v."desc" AS "rawDesc",
        COALESCE(array_to_string(vt.transcript, E'\n'), '') AS "rawTranscript",
        v.created_at AS created,
        v.manual_content_tags AS "contentTags"
      FROM "Video" v
      LEFT JOIN "VideoTranscript" vt ON vt."videoId" = v.aweme_id
      WHERE ${contentTag}::text IS NULL
         OR v.manual_content_tags @> ARRAY[${contentTag}]::text[]

      UNION ALL

      SELECT
        ip.aweme_id,
        'image'::text,
        ip."desc",
        '',
        ip.created_at,
        ARRAY[]::text[]
      FROM "ImagePost" ip
      WHERE ${contentTag}::text IS NULL
    ),
    lowered AS (
      SELECT docs.*, lower(docs."rawDesc") AS d, lower(docs."rawTranscript") AS t
      FROM docs
    ),
    matched AS (
      SELECT
        l.*,
        COALESCE((
          SELECT SUM(
            4 * LEAST(regexp_count(l.d, term), 3)
            + LEAST(regexp_count(l.t, term), 10)
          )
          FROM unnest(p.terms) term
        ), 0)::float8 AS score,
        EXISTS (SELECT 1 FROM unnest(p.terms) term WHERE l.d ~ term) AS "inDesc",
        EXISTS (SELECT 1 FROM unnest(p.terms) term WHERE l.t ~ term) AS "inTranscript"
      FROM lowered l, params p
      WHERE NOT EXISTS (
          SELECT 1 FROM unnest(p.terms) term
          WHERE l.d !~ term AND l.t !~ term
        )
        AND NOT EXISTS (
          SELECT 1 FROM unnest(p.excludes) x
          WHERE l.d ~ x OR l.t ~ x
        )
    )
  `;
}

const ORDER_BY = Prisma.sql`ORDER BY score DESC, created DESC, "awemeId"`;

export function searchRows(input: SearchInput, offset: number, limit: number) {
  return prisma.$queryRaw<SearchRow[]>`
    ${matchedCte(input)}
    SELECT
      "awemeId", type, score, "inDesc", "inTranscript",
      "rawDesc", "rawTranscript", "contentTags",
      (COUNT(*) OVER ())::int AS total
    FROM matched
    ${ORDER_BY}
    OFFSET ${offset}
    LIMIT ${limit};
  `;
}

/** 按搜索结果顺序返回全部命中作品的 ID。 */
export async function searchIds(input: SearchInput): Promise<string[]> {
  const rows = await prisma.$queryRaw<{ awemeId: string }[]>`
    ${matchedCte(input)}
    SELECT "awemeId" FROM matched
    ${ORDER_BY};
  `;
  return rows.map((r) => r.awemeId);
}