diff --git a/.env.example b/.env.example index a675040..ef000ce 100644 --- a/.env.example +++ b/.env.example @@ -18,6 +18,7 @@ RAG_RESULT_TOP_K=5 RAG_RRF_K=60 REDIS_URL=redis://localhost:6379/0 ANALYSIS_QUEUE_NAME=analysis +ANALYSIS_JOB_TIMEOUT_SECONDS=1800 GITHUB_TOKEN_TTL_SECONDS=3600 SPRING_CALLBACK_BASE_URL=http://localhost:8080 SPRING_SUCCESS_CALLBACK_PATH=/api/internal/analyses/{analysisId}/result diff --git a/app/api/routes/analyze.py b/app/api/routes/analyze.py index ee16fd1..b3dea8f 100644 --- a/app/api/routes/analyze.py +++ b/app/api/routes/analyze.py @@ -4,6 +4,7 @@ from redis.exceptions import RedisError from app.api.dependencies import verify_internal_token +from app.core.config import settings from app.jobs.analysis_job import run_analysis_job from app.jobs.queue import get_analysis_queue from app.jobs.secret_store import delete_github_token_reference, store_github_token @@ -45,6 +46,7 @@ async def analyze(request: AnalyzeRequest): job = queue.enqueue( run_analysis_job, request.to_job_payload(github_token_reference), + job_timeout=settings.ANALYSIS_JOB_TIMEOUT_SECONDS, ) except RedisError as exc: cleanup_token_reference(github_token_reference, request.analysis_id) diff --git a/app/core/config.py b/app/core/config.py index a15ea14..549a229 100644 --- a/app/core/config.py +++ b/app/core/config.py @@ -35,6 +35,7 @@ class Settings(BaseSettings): # Queue REDIS_URL: str = "redis://localhost:6379/0" ANALYSIS_QUEUE_NAME: str = "analysis" + ANALYSIS_JOB_TIMEOUT_SECONDS: int = 1800 GITHUB_TOKEN_TTL_SECONDS: int = 3600 # Spring Callback diff --git a/app/jobs/analysis_job.py b/app/jobs/analysis_job.py index 87434a2..3fdc5f8 100644 --- a/app/jobs/analysis_job.py +++ b/app/jobs/analysis_job.py @@ -17,7 +17,7 @@ from app.services.llm.explanation_generator import enrich_finding_with_explanation from app.services.normalizer.deduplicator import deduplicate_findings from app.services.normalizer.finding_normalizer import normalize_findings -from app.services.rag.hybrid_search import search_evidence_for_finding +from app.services.rag.hybrid_search import search_evidence_for_findings from app.services.scanner.base import AnalyzerContext, AnalyzerRunner, RawFinding from app.services.scanner.codeql_runner import CodeQLRunner from app.services.scanner.semgrep_runner import SemgrepRunner @@ -63,14 +63,37 @@ def run_analysis_job(payload: dict[str, Any]) -> dict[str, Any]: context.normalized_findings = normalize_findings(raw_findings) #정규화 후 context.normalized_findings = deduplicate_findings(context.normalized_findings) #중복 + evidence_by_finding = search_evidence_for_findings( + context.normalized_findings, + analysis_id=context.analysis_id, + ) + #LLM 설명 추가 - context.enriched_findings = [ - enrich_finding_with_explanation( + context.enriched_findings = [] + finding_total = len(context.normalized_findings) + for finding_index, (finding, evidence_documents) in enumerate( + zip(context.normalized_findings, evidence_by_finding), + start=1, + ): + logger.info( + "Finding enrichment started. analysis_id=%s finding_index=%s finding_total=%s tool=%s type=%s", + context.analysis_id, + finding_index, + finding_total, + finding.tool, + finding.type, + ) + enriched_finding = enrich_finding_with_explanation( finding, - search_evidence_for_finding(finding), + evidence_documents, + ) + context.enriched_findings.append(enriched_finding) + logger.info( + "Finding enrichment completed. analysis_id=%s finding_index=%s finding_total=%s", + context.analysis_id, + finding_index, + finding_total, ) - for finding in context.normalized_findings - ] success_payload = build_success_callback_payload(context) callback_client.send_success(success_payload) diff --git a/app/services/rag/hybrid_search.py b/app/services/rag/hybrid_search.py index 033f8db..c3d1880 100644 --- a/app/services/rag/hybrid_search.py +++ b/app/services/rag/hybrid_search.py @@ -34,10 +34,54 @@ def search_evidence_for_finding(finding: Finding) -> list[EvidenceDocument]: return search_evidence(build_rag_query(finding)) +# 동기 분석 job에서 모든 Finding을 하나의 event loop로 검색 +def search_evidence_for_findings( + findings: list[Finding], + analysis_id: int | None = None, +) -> list[list[EvidenceDocument]]: + queries = [build_rag_query(finding) for finding in findings] + return asyncio.run(search_evidence_batch_async(queries, analysis_id)) + + # 동기 분석 job에서 호출하기 위한 검색 진입점 def search_evidence(query: RagQuery) -> list[EvidenceDocument]: + return asyncio.run(search_evidence_safely_async(query)) + + +# 여러 query를 같은 event loop에서 순차적으로 검색 +async def search_evidence_batch_async( + queries: list[RagQuery], + analysis_id: int | None = None, +) -> list[list[EvidenceDocument]]: + results: list[list[EvidenceDocument]] = [] + finding_total = len(queries) + + for finding_index, query in enumerate(queries, start=1): + logger.info( + "Finding RAG search started. analysis_id=%s finding_index=%s finding_total=%s type=%s cwe_id=%s", + analysis_id, + finding_index, + finding_total, + query.vulnerability_type, + query.cwe_id, + ) + evidence_documents = await search_evidence_safely_async(query) + results.append(evidence_documents) + logger.info( + "Finding RAG search completed. analysis_id=%s finding_index=%s finding_total=%s evidence_count=%s", + analysis_id, + finding_index, + finding_total, + len(evidence_documents), + ) + + return results + + +# query 단위 검색 실패가 다른 Finding 처리를 막지 않도록 빈 결과로 fallback +async def search_evidence_safely_async(query: RagQuery) -> list[EvidenceDocument]: try: - return asyncio.run(search_evidence_async(query)) + return await search_evidence_async(query) except Exception: logger.exception( "RAG search failed. cwe_id=%s type=%s", diff --git a/tests/api/test_analyze.py b/tests/api/test_analyze.py new file mode 100644 index 0000000..390e8e2 --- /dev/null +++ b/tests/api/test_analyze.py @@ -0,0 +1,42 @@ +import unittest +from unittest.mock import Mock, patch + +from app.api.routes.analyze import analyze +from app.jobs.analysis_job import run_analysis_job +from app.schemas.analyze import AnalyzeRequest + + +class AnalyzeRouteTest(unittest.IsolatedAsyncioTestCase): + async def test_enqueue_uses_configured_job_timeout(self): + queue = Mock() + queue.enqueue.return_value = Mock(id="job-id") + request = AnalyzeRequest( + analysis_id=11, + repository_id=22, + repository_url="https://github.com/example/repository.git", + branch="main", + github_token="github-token", + ) + + with patch( + "app.api.routes.analyze.get_analysis_queue", + return_value=queue, + ), patch( + "app.api.routes.analyze.store_github_token", + return_value="token-reference", + ), patch( + "app.api.routes.analyze.settings.ANALYSIS_JOB_TIMEOUT_SECONDS", + 1800, + ): + response = await analyze(request) + + self.assertTrue(response.accepted) + queue.enqueue.assert_called_once_with( + run_analysis_job, + request.to_job_payload("token-reference"), + job_timeout=1800, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/jobs/test_analysis_job.py b/tests/jobs/test_analysis_job.py index 3fa2a80..5870e00 100644 --- a/tests/jobs/test_analysis_job.py +++ b/tests/jobs/test_analysis_job.py @@ -7,8 +7,8 @@ get_analyzer_runners, run_analysis_job, ) -from app.schemas.finding import FindingTool -from app.services.scanner.base import AnalyzerError +from app.schemas.finding import FindingSeverity, FindingTool +from app.services.scanner.base import AnalyzerError, RawFinding class AnalysisJobTest(unittest.TestCase): @@ -65,6 +65,31 @@ def test_stub_infra_runner_is_not_registered(self): [FindingTool.SEMGREP, FindingTool.CODEQL], ) + def test_multiple_findings_log_enrichment_progress(self): + runner = Mock(tool=FindingTool.SEMGREP) + runner.run.return_value = [ + build_raw_finding("FIRST", 10), + build_raw_finding("SECOND", 20), + ] + + with self.pipeline_patches([runner]), patch( + "app.jobs.analysis_job.search_evidence_for_findings", + return_value=[[], []], + ), patch( + "app.jobs.analysis_job.enrich_finding_with_explanation", + side_effect=lambda finding, _evidence: finding, + ), self.assertLogs( + "app.jobs.analysis_job", + level="INFO", + ) as logs: + result = run_analysis_job(self.payload) + + self.assertEqual(result["finding_count"], 2) + output = "\n".join(logs.output) + self.assertIn("finding_index=1 finding_total=2", output) + self.assertIn("finding_index=2 finding_total=2", output) + self.assertIn("Finding enrichment completed", output) + def test_missing_analysis_id_skips_failure_callback(self): callback = Mock() @@ -116,5 +141,16 @@ def __exit__(self, exc_type, exc_value, traceback): return self.stack.__exit__(exc_type, exc_value, traceback) +def build_raw_finding(finding_type: str, line_start: int) -> RawFinding: + return RawFinding( + tool=FindingTool.SEMGREP, + type=finding_type, + severity=FindingSeverity.HIGH, + file_path="src/example.py", + message="example finding", + line_start=line_start, + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/services/rag/__init__.py b/tests/services/rag/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/services/rag/test_hybrid_search.py b/tests/services/rag/test_hybrid_search.py new file mode 100644 index 0000000..ec1f391 --- /dev/null +++ b/tests/services/rag/test_hybrid_search.py @@ -0,0 +1,118 @@ +import asyncio +import unittest +from unittest.mock import AsyncMock, patch + +from app.schemas.finding import Finding, FindingSeverity, FindingTool +from app.services.rag.hybrid_search import search_evidence_for_findings + + +class HybridSearchTest(unittest.TestCase): + def test_multiple_findings_use_one_asyncio_run(self): + findings = [build_finding("FIRST"), build_finding("SECOND")] + real_asyncio_run = asyncio.run + search_loop_ids: list[int] = [] + + async def record_search_loop(_query): + search_loop_ids.append(id(asyncio.get_running_loop())) + return [] + + with patch( + "app.services.rag.hybrid_search.search_evidence_async", + new_callable=AsyncMock, + side_effect=record_search_loop, + ) as search_evidence_async, patch( + "app.services.rag.hybrid_search.asyncio.run", + wraps=real_asyncio_run, + ) as asyncio_run: + results = search_evidence_for_findings(findings) + + self.assertEqual(results, [[], []]) + asyncio_run.assert_called_once() + self.assertEqual(search_evidence_async.await_count, 2) + self.assertEqual(len(set(search_loop_ids)), 1) + + def test_multiple_findings_query_database_on_same_event_loop(self): + findings = [build_finding("FIRST"), build_finding("SECOND")] + database_loop_ids: list[int] = [] + session_factory = FakeAsyncSessionFactory(database_loop_ids) + + with patch( + "app.core.database.AsyncSessionLocal", + session_factory, + ), patch( + "app.services.rag.hybrid_search.OpenAIEmbedder", + ) as embedder: + embedder.return_value.embed_query.return_value = [0.0] + results = search_evidence_for_findings(findings, analysis_id=11) + + self.assertEqual(results, [[], []]) + self.assertEqual(session_factory.call_count, 2) + self.assertEqual(len(database_loop_ids), 4) + self.assertEqual(len(set(database_loop_ids)), 1) + + def test_finding_progress_is_logged(self): + findings = [build_finding("FIRST"), build_finding("SECOND")] + + with patch( + "app.services.rag.hybrid_search.search_evidence_async", + new_callable=AsyncMock, + return_value=[], + ), self.assertLogs( + "app.services.rag.hybrid_search", + level="INFO", + ) as logs: + search_evidence_for_findings(findings, analysis_id=11) + + output = "\n".join(logs.output) + self.assertIn("analysis_id=11 finding_index=1 finding_total=2", output) + self.assertIn("analysis_id=11 finding_index=2 finding_total=2", output) + self.assertIn("Finding RAG search completed", output) + + +class FakeAsyncSessionFactory: + def __init__(self, loop_ids: list[int]): + self.loop_ids = loop_ids + self.call_count = 0 + + def __call__(self): + self.call_count += 1 + return FakeAsyncSessionContext(self.loop_ids) + + +class FakeAsyncSessionContext: + def __init__(self, loop_ids: list[int]): + self.session = FakeAsyncSession(loop_ids) + + async def __aenter__(self): + return self.session + + async def __aexit__(self, exc_type, exc_value, traceback): + return False + + +class FakeAsyncSession: + def __init__(self, loop_ids: list[int]): + self.loop_ids = loop_ids + + async def execute(self, _query, _parameters): + self.loop_ids.append(id(asyncio.get_running_loop())) + return FakeQueryResult() + + +class FakeQueryResult: + def mappings(self): + return [] + + +def build_finding(finding_type: str) -> Finding: + return Finding( + tool=FindingTool.SEMGREP, + type=finding_type, + severity=FindingSeverity.HIGH, + file_path="src/example.py", + message="example finding", + ) + + +if __name__ == "__main__": + unittest.main()