ABAP-Indexing/index/db.py

"""Stage 4 — 인덱스 DB.

기본 백엔드는 SQLite(FTS5) — 개발/폐쇄망에서 무의존 동작.
PostgreSQL(+pgvector) 전환은 ASSUMPTIONS.md 와 schema_postgres.sql 참고.
계획서 §6.1 스키마의 SQLite 대응판이다 (embedding 컬럼은 임베딩 모델 확정 전까지 보류).

검색의 1차 단위는 logic_chunk(LLM 이 골라낸 로직 조각) 이고 unit 은 컨테이너·얇은 색인이다
(docs/logic-chunk-design.md).
"""
from __future__ import annotations

import json
import re
import sqlite3
from pathlib import Path

from config.settings import settings

SCHEMA = """
CREATE TABLE IF NOT EXISTS package(
  devclass TEXT PRIMARY KEY, text_ko TEXT, summary_json TEXT
);
CREATE TABLE IF NOT EXISTS program(
  name TEXT PRIMARY KEY, devclass TEXT, title_ko TEXT,
  created_on TEXT, changed_on TEXT, source_hash TEXT,
  summary_json TEXT, summary_status TEXT DEFAULT 'none', has_source INTEGER DEFAULT 0,
  text_symbols_json TEXT
);
CREATE TABLE IF NOT EXISTS include(
  program TEXT, include TEXT, code_hash TEXT, line_count INTEGER, code TEXT,
  PRIMARY KEY(program, include)
);
CREATE TABLE IF NOT EXISTS unit(
  unit_id TEXT PRIMARY KEY, program TEXT, include TEXT, unit_type TEXT, name TEXT,
  line_start INTEGER, line_end INTEGER, code_hash TEXT, signature TEXT,
  header_comment TEXT, loc INTEGER, refs_json TEXT, sub_chunks_json TEXT,
  summary_json TEXT, summary_status TEXT DEFAULT 'none', prompt_version INTEGER DEFAULT 0,
  summary_error TEXT, chunk_count INTEGER DEFAULT 0
);
CREATE INDEX IF NOT EXISTS idx_unit_program ON unit(program);
CREATE TABLE IF NOT EXISTS symbol(
  program TEXT, name TEXT, scope TEXT, unit_id TEXT, decl_include TEXT,
  decl_line INTEGER, type_text TEXT, ddic_ref TEXT, kind TEXT
);
CREATE INDEX IF NOT EXISTS idx_symbol ON symbol(program, name);
-- 정의부 — 조각을 복사해 붙여넣을 때 함께 가야 하는 **선언 원문**.
-- symbol 과 행이 겹쳐 보이지만 목적이 다르다: symbol 은 추적(어디서 채워지나), declaration 은
-- 붙여넣기(무엇을 함께 가져가야 컴파일되나). 그래서 code/줄범위/의존을 따로 담는다.
-- → parser/declarations.py, index/decls.py
CREATE TABLE IF NOT EXISTS declaration(
  program TEXT, name TEXT, kind TEXT, scope TEXT, unit_id TEXT, include TEXT,
  line_start INTEGER, line_end INTEGER, code TEXT, type_text TEXT, depends_json TEXT
);
CREATE INDEX IF NOT EXISTS idx_decl ON declaration(program, name);
CREATE INDEX IF NOT EXISTS idx_decl_unit ON declaration(unit_id);
CREATE TABLE IF NOT EXISTS symbol_write(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  program TEXT, symbol TEXT, unit_id TEXT, include TEXT, line INTEGER, kind TEXT,
  source_symbols TEXT, source_tables TEXT, stmt_text TEXT, callee TEXT
);
CREATE INDEX IF NOT EXISTS idx_write ON symbol_write(program, symbol);
CREATE TABLE IF NOT EXISTS symbol_read(
  program TEXT, symbol TEXT, unit_id TEXT, include TEXT, line INTEGER
);
CREATE INDEX IF NOT EXISTS idx_read ON symbol_read(program, symbol);
CREATE TABLE IF NOT EXISTS table_ref(
  program TEXT, unit_id TEXT, table_name TEXT, mode TEXT, line INTEGER
);
CREATE INDEX IF NOT EXISTS idx_table_ref ON table_ref(table_name);
CREATE TABLE IF NOT EXISTS call_edge(
  program TEXT, from_unit TEXT, to_unit TEXT, external_name TEXT, call_type TEXT, line INTEGER
);
CREATE INDEX IF NOT EXISTS idx_edge_to ON call_edge(to_unit);
CREATE INDEX IF NOT EXISTS idx_edge_ext ON call_edge(external_name);
CREATE TABLE IF NOT EXISTS topo(
  program TEXT, unit_id TEXT, ord INTEGER, PRIMARY KEY(program, unit_id)
);
CREATE TABLE IF NOT EXISTS llm_usage_log(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  ts TEXT, program TEXT, model TEXT, trigger_by TEXT,
  calls INTEGER, prompt_tokens INTEGER, completion_tokens INTEGER, cost_usd REAL,
  done INTEGER, failed INTEGER, skipped INTEGER, elapsed_s REAL,
  status TEXT DEFAULT 'done', total INTEGER DEFAULT 0, chunks INTEGER DEFAULT 0
);
CREATE TABLE IF NOT EXISTS logic_chunk(
  chunk_id TEXT PRIMARY KEY,
  program TEXT, include TEXT, unit_id TEXT, seq INTEGER,
  line_start INTEGER, line_end INTEGER, code_hash TEXT,
  kind TEXT, purpose_ko TEXT, purpose_en TEXT,
  keywords_ko TEXT, keywords_en TEXT, sap_objects TEXT,
  tables_read TEXT, tables_write TEXT, calls TEXT,
  confidence REAL, prompt_version INTEGER, extracted_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_chunk_program ON logic_chunk(program);
CREATE INDEX IF NOT EXISTS idx_chunk_unit ON logic_chunk(unit_id);
CREATE TABLE IF NOT EXISTS ddic_table(name TEXT PRIMARY KEY, text_ko TEXT);
CREATE TABLE IF NOT EXISTS ddic_field(tabname TEXT, field TEXT, text_ko TEXT, data_element TEXT);
CREATE TABLE IF NOT EXISTS tcode(tcode TEXT PRIMARY KEY, program TEXT, text_ko TEXT);
-- 프로그램 색인은 **두 테이블로 나눈다**. bm25 는 행 전체 길이로 정규화하므로, 짧은 타이틀과
-- 긴 요약을 같은 행에 넣으면 요약이 있는 프로그램이 오히려 밀린다 (실측: 타이틀이 "총계정원장 조회"인
-- ZFIR10070 이 요약을 붙인 뒤 1위→33위. 행 길이 1,872자 vs 타이틀만 있는 행 40자).
-- 나누면 각 테이블 안에서 행 길이가 비슷해져 요약은 **점수를 더하기만** 하고 이름/타이틀 일치를
-- 밀어내지 않는다.
CREATE VIRTUAL TABLE IF NOT EXISTS program_fts USING fts5(
  name, title, keywords, bigrams, tokenize='unicode61'
);
-- 서술 색인 — LLM 요약·조각 키워드·테이블·텍스트 심볼 (수정사항 1·3번)
CREATE VIRTUAL TABLE IF NOT EXISTS program_desc_fts USING fts5(
  name UNINDEXED, purpose, keywords, objects, bigrams, tokenize='unicode61'
);
CREATE VIRTUAL TABLE IF NOT EXISTS unit_fts USING fts5(
  unit_id UNINDEXED, program UNINDEXED, name, purpose, keywords, bigrams, tokenize='unicode61'
);
CREATE VIRTUAL TABLE IF NOT EXISTS chunk_fts USING fts5(
  chunk_id UNINDEXED, program UNINDEXED, purpose, keywords, objects, bigrams, tokenize='unicode61'
);
"""


def db_path() -> Path:
    url = settings.database_url
    if not url.startswith("sqlite:///"):
        raise RuntimeError(
            f"현재 빌드는 SQLite 백엔드만 구현되어 있다 (DATABASE_URL={url}). ASSUMPTIONS.md 참고."
        )
    return Path(url.replace("sqlite:///", ""))


# 기존 DB에 나중에 추가된 컬럼 — CREATE TABLE IF NOT EXISTS 는 컬럼을 추가하지 못한다
_COLUMN_MIGRATIONS = [
    ("unit", "summary_error", "TEXT"),
    ("unit", "chunk_count", "INTEGER DEFAULT 0"),           # 로직 조각 수 (docs/logic-chunk-design.md)
    ("program", "text_symbols_json", "TEXT"),               # TEXT-nnn → 한국어 텍스트 (요약 문맥용)
    ("llm_usage_log", "status", "TEXT DEFAULT 'done'"),
    ("llm_usage_log", "total", "INTEGER DEFAULT 0"),
    ("llm_usage_log", "chunks", "INTEGER DEFAULT 0"),
]


def _migrate(con: sqlite3.Connection) -> None:
    for table, column, decl in _COLUMN_MIGRATIONS:
        cols = {r[1] for r in con.execute(f"PRAGMA table_info({table})")}
        if column not in cols:
            con.execute(f"ALTER TABLE {table} ADD COLUMN {column} {decl}")
    _migrate_fts(con)


# FTS5 가상테이블은 ALTER 로 컬럼을 추가할 수 없다. 모두 파생 데이터(program/unit/logic_chunk 에서
# 재생성 가능)이므로 컬럼 구성이 달라지면 DROP 후 다시 만든다. 재적재는 호출 측 책임:
#   program_fts → index.loader.refresh_program_fts(con)
#   chunk_fts / unit_fts → 해당 프로그램 재적재 또는 요약 재실행
_FTS_EXPECTED = {
    "program_fts": ["name", "title", "keywords", "bigrams"],
    "program_desc_fts": ["name", "purpose", "keywords", "objects", "bigrams"],
    "unit_fts": ["unit_id", "program", "name", "purpose", "keywords", "bigrams"],
    "chunk_fts": ["chunk_id", "program", "purpose", "keywords", "objects", "bigrams"],
}


def _migrate_fts(con: sqlite3.Connection) -> list[str]:
    """컬럼 구성이 바뀐 FTS 테이블을 재생성하고, 재생성한 테이블 이름을 돌려준다."""
    rebuilt: list[str] = []
    for table, expected in _FTS_EXPECTED.items():
        cols = [r[1] for r in con.execute(f"PRAGMA table_info({table})")]
        if cols and cols != expected:
            con.execute(f"DROP TABLE {table}")
            con.executescript(SCHEMA)  # IF NOT EXISTS — 빠진 테이블만 다시 만든다
            rebuilt.append(table)
    return rebuilt


def connect() -> sqlite3.Connection:
    path = db_path()
    path.parent.mkdir(parents=True, exist_ok=True)
    con = sqlite3.connect(str(path))
    con.row_factory = sqlite3.Row
    con.execute("PRAGMA journal_mode=WAL")
    con.executescript(SCHEMA)
    _migrate(con)
    return con


_HANGUL_SEQ = re.compile(r"[가-힣]{2,}")
_TEXT_SYM_REF = re.compile(r"TEXT-(\w{3})", re.I)


def clean_comment(text: str | None) -> str:
    """헤더 주석에서 장식 문자(----, ****, &, *&)를 제거해 사람이 읽을 문장만 남긴다.

    적재 시점의 unit_fts.purpose 와 요약 프롬프트가 모두 이걸 거쳐야 한다 — 거치지 않으면
    색인에 '----------------' 만 들어간다 (수정사항 1번).
    """
    if not text:
        return ""
    parts = re.split(r"[*&\-]{3,}|\s\*\s|&", text)
    words = " ".join(p.strip() for p in parts if p.strip())
    words = re.sub(r"^(Form|FORM|Include|INCLUDE)\s+", "", words).strip()
    return words[:150]


def text_symbol_phrases(code: str, text_symbols: dict[str, str]) -> list[str]:
    """코드에 등장하는 TEXT-nnn 을 실제 한국어 텍스트로 치환한 목록 (수정사항 3번).

    LLM 없이도 조각·프로그램에 자연어 앵커가 생기는 가장 싼 경로라 FTS 에도 넣는다.
    """
    if not code or not text_symbols:
        return []
    out: list[str] = []
    for m in _TEXT_SYM_REF.finditer(code):
        val = text_symbols.get(m.group(1).upper())
        if val and val not in out:
            out.append(val)
    return out


def bigrams(text: str) -> str:
    """한국어 형태소 분석기 없이 부분일치를 위해 한글 연속열의 2-gram을 공백 구분으로 나열."""
    grams: list[str] = []
    for seq in _HANGUL_SEQ.findall(text or ""):
        grams.extend(seq[i : i + 2] for i in range(len(seq) - 1))
    return " ".join(dict.fromkeys(grams))


# FTS5 컬럼 가중치. 기본 `rank`(= 모든 컬럼 가중치 1)를 쓰면 안 된다:
# bm25 는 문서 길이로 정규화하므로, 요약·키워드가 붙어 **길어진** 행이 타이틀만 있는 행보다
# 낮게 나온다. 요약이 있는 프로그램이 오히려 밀리는 역전이 생긴다(실측: 타이틀이 "총계정원장 조회"인
# ZFIR10070 이 요약을 붙인 뒤 top5 밖으로 밀려났다). 이름·타이틀을 강하게 가중해 바로잡는다.
_BM25_WEIGHTS = {
    # 컬럼 순서와 같아야 한다 (UNINDEXED 컬럼도 자리를 차지한다)
    "program_fts": (12.0, 8.0, 1.5, 1.0),                    # name title keywords bigrams
    "program_desc_fts": (0.0, 3.0, 2.0, 1.5, 1.0),           # name purpose keywords objects bigrams
    "unit_fts": (0.0, 0.0, 6.0, 2.0, 1.5, 1.0),              # unit_id program name purpose keywords bigrams
    "chunk_fts": (0.0, 0.0, 4.0, 2.0, 1.5, 1.0),             # chunk_id program purpose keywords objects bigrams
}


def bm25_rank(table: str) -> str:
    """`ORDER BY` / `SELECT` 에 넣을 가중 bm25 식. 점수는 음수이고 작을수록 관련도가 높다."""
    weights = ", ".join(f"{w}" for w in _BM25_WEIGHTS[table])
    return f"bm25({table}, {weights})"


def query_tokens(q: str) -> list[str]:
    """질의 문자열 → 검색 토큰 목록 (구두점·기호 제거)."""
    return [t for t in re.split(r"[^\w가-힣]+", q or "") if t]


def fts_or(terms: list[str]) -> str:
    """임의 용어 목록 → FTS5 MATCH 식. 용어마다 큰따옴표로 감싸고 한글 2-gram 을 함께 OR 결합.

    질의 확장(query/expand.py)이 원질의와 동의어를 서로 다른 MATCH 식으로 나눠 쓰기 위해
    fts_escape 에서 분리했다.
    """
    parts: list[str] = []
    for t in terms:
        t = t.strip()
        if not t:
            continue
        parts.append('"' + t.replace('"', "") + '"')
        for g in bigrams(t).split():
            parts.append(f'"{g}"')
    return " OR ".join(dict.fromkeys(parts)) if parts else '""'


def fts_escape(q: str) -> str:
    """FTS5 MATCH 질의에 안전한 형태로: 토큰별 큰따옴표 감싸고 OR 결합."""
    return fts_or(query_tokens(q))


def loads(v: str | None):
    return json.loads(v) if v else None