Mini-challenge — TaskFlow advanced search: FTS + JSONB + pgvector

-- =====================================================
-- Extension setup (chay mot lan)
-- =====================================================
CREATE EXTENSION IF NOT EXISTS unaccent;
CREATE EXTENSION IF NOT EXISTS pg_trgm;
CREATE EXTENSION IF NOT EXISTS vector;

-- Vietnamese FTS config: copy simple, add unaccent mapping (see lesson 05)
-- CREATE TEXT SEARCH CONFIGURATION has no IF NOT EXISTS -> guard by DO block
DO $$
BEGIN
  IF NOT EXISTS (SELECT 1 FROM pg_ts_config WHERE cfgname = 'vietnamese') THEN
    CREATE TEXT SEARCH CONFIGURATION vietnamese (COPY = simple);
    ALTER TEXT SEARCH CONFIGURATION vietnamese
      ALTER MAPPING FOR hword, hword_part, word
      WITH unaccent, simple;
  END IF;
END $$;

-- =====================================================
-- Extend tasks table (chay mot lan, idempotent)
-- =====================================================
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS body          TEXT;
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS task_metadata JSONB NOT NULL DEFAULT '{}';

-- Generated column: tu dong cap nhat khi title/body thay doi
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS body_tsv tsvector
  GENERATED ALWAYS AS (
    to_tsvector('vietnamese',
      coalesce(title, '') || ' ' || coalesce(body, ''))
  ) STORED;

-- pgvector: 1536 dims cho text-embedding-3-small
-- NULL ban dau -- backfill tu external script (Step 5)
ALTER TABLE tasks ADD COLUMN IF NOT EXISTS embedding vector(1536);

-- =====================================================
-- Seed: 100k task voi title VN + metadata
-- Mat khoang 30-60 giay
-- =====================================================
INSERT INTO tasks (
  project_id, assignee_id, title, body,
  status, due_at, task_metadata, created_at, updated_at
)
SELECT
  (random() * 10 + 1)::int,               -- project_id 1-10
  (random() * 20 + 1)::int,               -- assignee_id 1-20

  -- Title tieng Viet diverse
  (ARRAY[
    'Trien khai feature thanh toan ' || gs,
    'Sua loi dang nhap nguoi dung ' || gs,
    'Toi uu toc do truy van ' || gs,
    'Viet unit test cho module gio hang ' || gs,
    'Review code pull request ' || gs,
    'Cap nhat tai lieu API ' || gs,
    'Fix payment gateway timeout ' || gs,
    'Deploy len moi truong staging ' || gs,
    'Kiem tra bao mat endpoint ' || gs,
    'Khac phuc loi xuat hoa don ' || gs,
    'Tich hop webhook thanh toan ' || gs,
    'Refactor service xu ly don hang ' || gs,
    'Monitor latency sau deploy ' || gs,
    'Cau hinh alert cho loi 5xx ' || gs,
    'Migrate du lieu sang schema moi ' || gs
  ])[floor(random() * 15) + 1],

  -- Body mo ta chi tiet hon
  (ARRAY[
    'Can kiem tra flow thanh toan end-to-end truoc khi len prod.',
    'Nguoi dung bao cao loi khi dang nhap bang Google OAuth.',
    'Query dashboard chay 3 giay, can them index phu hop.',
    'Stripe webhook tra ve 500 trong 2 phan tram request.',
    'Sau deploy 14h, p99 latency tang tu 120ms len 450ms.'
  ])[floor(random() * 5) + 1],

  -- Status: mot random() duy nhat (r) cho ca CASE, neu goi random() moi nhanh
  -- thi phan bo lech (~45/41/13/1 thay vi 45/30/20/5)
  (ARRAY['todo','doing','done','archived'])[
    CASE
      WHEN r < 0.45 THEN 1   -- 45% todo
      WHEN r < 0.75 THEN 2   -- 30% doing
      WHEN r < 0.95 THEN 3   -- 20% done
      ELSE 4                 -- 5% archived
    END
  ],

  -- due_at: -8..+32 ngay -> 20% da qua han, 80% trong ~30 ngay toi
  now() + ((random() * 40 - 8) * interval '1 day'),

  -- task_metadata JSONB: labels + priority + source
  jsonb_build_object(
    'labels', (
      CASE floor(random() * 8)::int
        WHEN 0 THEN '["urgent","backend"]'::jsonb
        WHEN 1 THEN '["urgent","frontend"]'::jsonb
        WHEN 2 THEN '["backend","performance"]'::jsonb
        WHEN 3 THEN '["frontend","ui"]'::jsonb
        WHEN 4 THEN '["urgent","security"]'::jsonb
        WHEN 5 THEN '["backend","migration"]'::jsonb
        WHEN 6 THEN '["devops","deployment"]'::jsonb
        ELSE        '["testing"]'::jsonb
      END
    ),
    'priority', (ARRAY['low','medium','high','critical'])[floor(random() * 4) + 1],
    'source',   (ARRAY['jira','linear','internal'])[floor(random() * 3) + 1]
  ),

  now() - (random() * interval '60 days'),
  now() - (random() * interval '30 days')
FROM generate_series(1, 100000) gs
CROSS JOIN LATERAL (SELECT random() AS r) x;   -- r moi cho moi row

-- Verify
SELECT status, count(*) FROM tasks GROUP BY status ORDER BY count(*) DESC;
-- Expected: todo ~45k, doing ~30k, done ~20k, archived ~5k

SELECT count(*) FROM tasks WHERE task_metadata @> '{"labels":["urgent"]}';
-- Expected: ~37.5k (3 trong 8 nhanh label deu co 'urgent' -> 37.5% cua 100k)

ANALYZE tasks;

-- =====================================================
-- Index 1: GIN cho FTS (body_tsv generated column)
-- =====================================================
CREATE INDEX IF NOT EXISTS idx_tasks_body_tsv
  ON tasks USING gin(body_tsv);

-- =====================================================
-- Index 2: GIN jsonb_path_ops cho JSONB containment
-- jsonb_path_ops chi ho tro @> (containment) nhung nho hon
-- va nhanh hon gin_default_ops cho pattern nay
-- =====================================================
CREATE INDEX IF NOT EXISTS idx_tasks_metadata
  ON tasks USING gin(task_metadata jsonb_path_ops);

-- =====================================================
-- Index 3: B-tree cho due_at range filter
-- =====================================================
CREATE INDEX IF NOT EXISTS idx_tasks_due_at
  ON tasks(due_at);

-- =====================================================
-- Index 4: HNSW cho pgvector cosine similarity
-- m=16, ef_construction=64 la default production-safe
-- =====================================================
CREATE INDEX IF NOT EXISTS idx_tasks_embedding
  ON tasks USING hnsw (embedding vector_cosine_ops)
  WITH (m = 16, ef_construction = 64);

SELECT indexname, indexdef
FROM pg_indexes
WHERE tablename = 'tasks'
  AND indexname LIKE 'idx_tasks_%'
ORDER BY indexname;

EXPLAIN (ANALYZE, BUFFERS)
SELECT
  id, title, due_at,
  ts_rank(body_tsv,
    websearch_to_tsquery('vietnamese', 'deploy')) AS rank
FROM tasks,
     websearch_to_tsquery('vietnamese', 'deploy') query
WHERE task_metadata @> '{"labels":["urgent"]}'
  AND body_tsv @@ query
  AND due_at BETWEEN now() AND now() + INTERVAL '7 days'
ORDER BY rank DESC, due_at ASC
LIMIT 20;

Limit  (cost=2341.89..2341.94 rows=20 width=96)
       (actual time=8.341..8.347 rows=20 loops=1)
  ->  Sort  (cost=2341.89..2342.01 rows=48 width=96)
            (actual time=8.339..8.343 rows=20 loops=1)
        Sort Key: (ts_rank(...)) DESC, due_at
        Sort Method: quicksort  Memory: 28kB
        ->  Bitmap Heap Scan on tasks
              (cost=312.45..2340.34 rows=48 width=96)
              (actual time=3.812..8.201 rows=47 loops=1)
              Recheck Cond: ((task_metadata @> '{"labels":["urgent"]}')
                             AND (body_tsv @@ query)
                             AND (due_at >= now())
                             AND (due_at <= (now() + '7 days'::interval)))
              Heap Blocks: exact=43
              Buffers: shared hit=312 read=28
              ->  BitmapAnd
                    (cost=312.45..312.45 rows=48 width=0)
                    (actual time=3.641..3.641 rows=0 loops=1)
                    ->  Bitmap Index Scan on idx_tasks_metadata
                          (cost=0.00..89.12 rows=2431 width=0)
                          (actual time=0.912..0.912 rows=2413 loops=1)
                          Index Cond: (task_metadata @>
                                       '{"labels":["urgent"]}'::jsonb)
                    ->  Bitmap Index Scan on idx_tasks_body_tsv
                          (cost=0.00..143.21 rows=1892 width=0)
                          (actual time=1.234..1.234 rows=1876 loops=1)
                          Index Cond: (body_tsv @@ query)
                    ->  Bitmap Index Scan on idx_tasks_due_at
                          (cost=0.00..78.34 rows=7124 width=0)
                          (actual time=0.891..0.891 rows=7089 loops=1)
                          Index Cond: ((due_at >= now()) AND
                                       (due_at <= (now() + '7 days')))
Planning Time: 1.423 ms
Execution Time: 8.891 ms

import openai

def get_embedding(text: str) -> list[float]:
    resp = openai.embeddings.create(
        input=text,
        model="text-embedding-3-small"   # 1536 dims
    )
    return resp.data[0].embedding

# Backfill: tinh embedding cho tasks chua co
def backfill_embeddings(db, batch_size: int = 100):
    while True:
        rows = db.query(
            "SELECT id, title, body FROM tasks "
            "WHERE embedding IS NULL ORDER BY id LIMIT %s",
            (batch_size,)
        )
        if not rows:
            break
        texts = [f"{r.title}. {r.body or ''}" for r in rows]
        resp = openai.embeddings.create(
            input=texts,
            model="text-embedding-3-small"
        )
        for row, emb in zip(rows, resp.data):
            db.execute(
                "UPDATE tasks SET embedding = %s WHERE id = %s",
                (emb.embedding, row.id)
            )

EXPLAIN ANALYZE
SELECT id, title, embedding <=> $1 AS distance
FROM tasks
WHERE embedding IS NOT NULL
ORDER BY embedding <=> $1
LIMIT 5;

Limit  (cost=0.01..0.26 rows=5 width=40)
       (actual time=12.341..12.398 rows=5 loops=1)
  ->  Index Scan using idx_tasks_embedding on tasks
        (cost=0.01..4123.45 rows=79845 width=40)
        (actual time=12.339..12.394 rows=5 loops=1)
        Order By: (embedding <=> $1)
        Filter: (embedding IS NOT NULL)
Execution Time: 12.543 ms
-- Expected: 10-30ms cho 100k vector co embedding
-- ANN (approximate): HNSW co the bo sot mot it result o bien
-- ef_search mac dinh 40, tang len 100 cho recall tot hon

-- $1 = embedding vector cua query string (tinh o app layer)
-- $2 = keyword string cho FTS (vi du: 'payment bug')
-- Weight: vector 60%, FTS 40% (dieu chinh theo use case)

WITH fts AS (
  SELECT
    id,
    ts_rank(body_tsv,
      websearch_to_tsquery('vietnamese', $2)) AS fts_score
  FROM tasks
  WHERE body_tsv @@ websearch_to_tsquery('vietnamese', $2)
  ORDER BY fts_score DESC
  LIMIT 50
),

vec AS (
  SELECT
    id,
    1.0 - (embedding <=> $1)   AS vec_score  -- cosine similarity: 1=identical
  FROM tasks
  WHERE embedding IS NOT NULL
  ORDER BY embedding <=> $1
  LIMIT 50
)

SELECT
  t.id,
  t.title,
  t.status,
  t.due_at,
  COALESCE(fts.fts_score,  0) * 0.4
    + COALESCE(vec.vec_score, 0) * 0.6       AS hybrid_score,
  COALESCE(fts.fts_score,  0)                AS fts_score,
  COALESCE(vec.vec_score,  0)                AS vec_score
FROM tasks t
LEFT JOIN fts USING (id)
LEFT JOIN vec USING (id)
WHERE fts.id IS NOT NULL OR vec.id IS NOT NULL
ORDER BY hybrid_score DESC
LIMIT 10;

-- Kiem tra extension da duoc cai
SELECT name, installed_version FROM pg_available_extensions
WHERE name IN ('unaccent','pg_trgm','vector');
-- Phai thay ca 3 co installed_version khac NULL

-- Kiem tra 4 index da ton tai
SELECT indexname FROM pg_indexes
WHERE tablename = 'tasks' AND indexname LIKE 'idx_tasks_%'
ORDER BY indexname;
-- Expected: idx_tasks_body_tsv, idx_tasks_due_at,
--           idx_tasks_embedding, idx_tasks_metadata

-- Kiem tra composite query co BitmapAnd
EXPLAIN (ANALYZE, BUFFERS)
SELECT id FROM tasks,
  websearch_to_tsquery('vietnamese', 'deploy') q
WHERE task_metadata @> '{"labels":["urgent"]}'
  AND body_tsv @@ q
  AND due_at BETWEEN now() AND now() + INTERVAL '7 days'
LIMIT 20;
-- Phai thay: BitmapAnd -> 3x Bitmap Index Scan