Self-Correction: 왜 3번까지만 재시도하는가 [4]

class CorrectionState(TypedDict):
    sql: str
    question: str
    tables: list      # 선택된 테이블 정보
    dialect: str       # MySQL / PostgreSQL
    issues: list[str]  # 발견된 문제들
    attempt: int       # 현재 시도 횟수
    max_attempts: int  # 최대 3
    status: str        # VALID / FIXED / UNFIXABLE
    original_sql: str  # 최초 생성 SQL (비교용)

if history[-1].get("message_type") == "clarification":
    is_followup = True
    # 이전 질문과 답변을 합쳐서 enriched question 생성
    enriched_question = f"{현재_답변} {이전_질문}"

if not result.sql and not result.clarification_question:
    # 비정상 상태 로깅
    return {"clarification_question": "질문을 더 구체적으로 해주시겠어요?"}

new_history.append({
    "role": "assistant",
    "content": assistant_content,
    "message_type": "answer" | "clarification",
    "selected_tables": state.get("selected_table_names", []),
})

# before
graph.add_edge("guardrail", "validation")
 
# after
graph.add_conditional_edges(
    "guardrail",
    self._route_after_guardrail,
    {"validation": "validation", "end_error": "save"},
)
 
def _route_after_guardrail(self, state: GraphState) -> str:
    if state.get("validation_status") == "UNFIXABLE" or state.get("error"):
        return "end_error"
    return "validation"

class TestRouteAfterGuardrail:
    def test_unfixable_routes_to_end_error(self):
        """UNFIXABLE → end_error"""
        state = {"validation_status": "UNFIXABLE"}
        assert _route_after_guardrail(state) == "end_error"
 
    def test_error_routes_to_end_error(self):
        """error 존재 → end_error"""
        state = {"error": "some error"}
        assert _route_after_guardrail(state) == "end_error"
 
    def test_safe_routes_to_validation(self):
        """정상 → validation"""
        state = {"validation_status": "VALID"}
        assert _route_after_guardrail(state) == "validation"
 
class TestDecomposerMultiTable:
    def test_multi_table_data_types_decomposed(self):
        """서로 다른 테이블 데이터 → 분해"""
        result = decompose("A 정보와 B 정보를 보여줘")
        assert len(result.sub_questions) == 2

class TestDMLBlocked:
    def test_insert_blocked(self):
        result = check_sql_safety("INSERT INTO users (name) VALUES ('test')")
        assert result.is_safe is False
 
    def test_drop_blocked(self):
        result = check_sql_safety("DROP TABLE users")
        assert result.is_safe is False
 
    def test_grant_blocked(self):
        result = check_sql_safety("GRANT ALL ON users TO admin")
        assert result.is_safe is False

def test_semicolon_multi_query(self):
    sql = "SELECT * FROM users; DROP TABLE users;"
    result = check_sql_safety(sql)
    assert result.is_safe is False
    assert any("다중 쿼리" in v for v in result.violations)
 
def test_trailing_semicolon_ok(self):
    result = check_sql_safety("SELECT * FROM users;")
    assert result.is_safe is True  # 끝의 세미콜론은 허용

class TestLLMGeneratedLimitPreservation:
    def test_preserves_limit_1(self):
        """'가장 최근' → LIMIT 1 보존"""
        result = check_sql_safety(
            "SELECT * FROM tb_evaluation ORDER BY created_at DESC LIMIT 1"
        )
        assert "LIMIT 1" in result.filtered_sql
        assert "LIMIT 1000" not in result.filtered_sql
 
    def test_preserves_limit_5(self):
        """'상위 5개' → LIMIT 5 보존"""
        result = check_sql_safety(
            "SELECT dept_name, COUNT(*) cnt FROM departments "
            "GROUP BY dept_name ORDER BY cnt DESC LIMIT 5"
        )
        assert "LIMIT 5" in result.filtered_sql
 
    def test_caps_limit_exceeding_max(self):
        """LIMIT 5000 → 1000으로 캡"""
        result = check_sql_safety("SELECT * FROM employees LIMIT 5000")
        assert "LIMIT 1000" in result.filtered_sql
        assert "5000" not in result.filtered_sql

def test_cte_name_not_blocked(self):
    """CTE 이름(ranked)이 화이트리스트에 없어도 차단되지 않음"""
    sql = "WITH ranked AS (SELECT * FROM users) SELECT * FROM ranked"
    result = check_sql_safety(sql, allowed_tables=["users"])
    assert result.is_safe is True
 
def test_multiple_cte_not_blocked(self):
    """다중 CTE + JOIN에서도 실제 테이블만 검사"""
    sql = ("WITH a AS (SELECT * FROM users), "
           "b AS (SELECT * FROM orders) "
           "SELECT * FROM a JOIN b ON a.id = b.user_id")
    result = check_sql_safety(sql, allowed_tables=["users", "orders"])
    assert result.is_safe is True
 
def test_recursive_cte_not_blocked(self):
    """WITH RECURSIVE도 처리"""
    sql = ("WITH RECURSIVE cte AS (SELECT 1 UNION ALL SELECT * FROM users) "
           "SELECT * FROM cte")
    result = check_sql_safety(sql, allowed_tables=["users"])
    assert result.is_safe is True

def test_cross_schema_table_allowed_by_bare_name(self):
    """erp.memberorgrelation → bare name으로 매칭"""
    result = check_sql_safety(
        "SELECT * FROM erp.memberorgrelation",
        allowed_tables=["memberorgrelation", "mit_common.member"],
    )
    assert result.is_safe is True
 
def test_cross_schema_table_allowed_by_qualified_name(self):
    """schema.table 전체 이름으로도 매칭"""
    result = check_sql_safety(
        "SELECT * FROM mit_common.member",
        allowed_tables=["mit_common.member"],
    )
    assert result.is_safe is True

def test_qualify_cross_schema(self):
    """bare name → schema-qualified name 교정"""
    tables = self._make_tables(["mit_common.member", "employee"])
    sql = "SELECT * FROM member"
    result = qualify_table_names(sql, tables)
    assert "FROM mit_common.member" in result
 
def test_qualify_no_conflict(self):
    """primary에 같은 bare name 있으면 교정 스킵"""
    tables = self._make_tables(["mit_common.member", "member"])
    sql = "SELECT * FROM member"
    result = qualify_table_names(sql, tables)
    assert "mit_common.member" not in result  # 충돌이므로 교정 안 함
 
def test_qualify_join(self):
    """JOIN 절의 bare name도 교정"""
    tables = self._make_tables(["mit_common.member", "employee"])
    sql = "SELECT * FROM employee JOIN member ON employee.id = member.emp_id"
    result = qualify_table_names(sql, tables)
    assert "JOIN mit_common.member" in result

def test_valid_returns_end_valid(self):
    state = {"status": "VALID", "attempt": 1, "max_attempts": 3,
             "issues": [], "previous_issues": []}
    assert _route_after_validate(state) == "end_valid"
 
def test_max_attempts_reached_returns_end_unfixable(self):
    state = {"status": "PENDING", "attempt": 3, "max_attempts": 3,
             "issues": ["error"], "previous_issues": []}
    assert _route_after_validate(state) == "end_unfixable"
 
def test_pending_under_max_returns_correct(self):
    state = {"status": "PENDING", "attempt": 1, "max_attempts": 3,
             "issues": ["error A"], "previous_issues": []}
    assert _route_after_validate(state) == "correct"

def test_same_error_twice_returns_end_unfixable(self):
    """동일 에러 2회 연속 → 조기 탈출"""
    state = {"status": "PENDING", "attempt": 2, "max_attempts": 3,
             "issues": ["column X not found"],
             "previous_issues": ["column X not found"]}
    assert _route_after_validate(state) == "end_unfixable"
 
def test_different_error_continues(self):
    """다른 에러면 계속 시도"""
    state = {"status": "PENDING", "attempt": 2, "max_attempts": 3,
             "issues": ["column Y not found"],
             "previous_issues": ["column X not found"]}
    assert _route_after_validate(state) == "correct"

def test_non_selected_table_detected(self):
    sql = ("SELECT m.name FROM mit_common.member m "
           "JOIN mit_common.memberdetail md ON m.id = md.member_id")
    selected = ["mit_common.member"]
    issues = _check_selected_tables_only(sql, selected)
    assert len(issues) == 1
    assert "memberdetail" in issues[0]
 
def test_cte_not_flagged(self):
    """CTE 이름은 false positive 아님"""
    sql = ("WITH dept_stats AS (SELECT dept_id, COUNT(*) cnt "
           "FROM mit_hr.department GROUP BY dept_id) "
           "SELECT * FROM dept_stats")
    selected = ["mit_hr.department"]
    issues = _check_selected_tables_only(sql, selected)
    assert issues == []

def test_phone_with_dashes(self, masker):
    rows = [{"phone": "010-1234-5678"}]
    result = masker.mask_results(rows)
    assert result.pii_found is True
    assert "****" in result.masked_rows[0]["phone"]
    assert "5678" in result.masked_rows[0]["phone"]  # 뒤 4자리 보존
    assert "010" in result.masked_rows[0]["phone"]    # 앞 3자리 보존
 
def test_ssn_with_dash(self, masker):
    rows = [{"data": "900101-1234567"}]
    result = masker.mask_results(rows)
    assert "900101" in masked    # 생년월일 보존
    assert "*******" in masked   # 뒷자리 마스킹
 
def test_card_with_dashes(self, masker):
    rows = [{"data": "1234-5678-9012-3456"}]
    result = masker.mask_results(rows)
    assert "3456" in masked         # 마지막 4자리 보존
    assert "****-****-****-" in masked  # 나머지 마스킹
 
def test_email_pattern(self, masker):
    rows = [{"data": "john@example.com"}]
    result = masker.mask_results(rows)
    assert "j***@example.com" == masked  # 첫 글자만 보존

def test_name_column(self, masker):
    rows = [{"user_name": "홍길동"}]
    result = masker.mask_results(rows)
    assert result.pii_found is True
    assert result.detections[0].pii_type == PIIType.NAME
    assert masked[0] == "홍"     # 성은 보존
    assert "**" in masked        # 이름은 마스킹
 
def test_address_column(self, masker):
    rows = [{"address": "서울시 강남구 역삼동 123"}]
    result = masker.mask_results(rows)
    assert masked.startswith("서울시")  # 시 단위 보존
    assert "****" in masked             # 상세 주소 마스킹
 
def test_phone_column_name(self, masker):
    rows = [{"전화번호": "anything"}]
    result = masker.mask_results(rows)
    assert result.detections[0].pii_type == PIIType.PHONE

def test_disabled_returns_original(self):
    masker = PIIMasker(enabled=False)
    rows = [{"phone": "010-1234-5678", "email": "test@test.com"}]
    result = masker.mask_results(rows)
    assert result.pii_found is False
    assert result.masked_rows == rows  # 원본 그대로