99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from .locks import WorkspaceLockManager
|
|
from .models import QueryRequest, QueryResponse, SyncRequest, SyncResponse
|
|
from .qmd_client import QMDClient
|
|
from .sync_service import SyncService
|
|
|
|
|
|
class QueryService:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
sync_service: SyncService,
|
|
qmd_client: QMDClient,
|
|
lock_manager: WorkspaceLockManager,
|
|
) -> None:
|
|
self.sync_service = sync_service
|
|
self.qmd_client = qmd_client
|
|
self.lock_manager = lock_manager
|
|
|
|
def handle_query(self, request: QueryRequest) -> QueryResponse:
|
|
resolved = self.sync_service.workspace_manager.resolve_workspace(
|
|
branch=request.branch,
|
|
memory_profile=request.memory_profile,
|
|
)
|
|
|
|
with self.lock_manager.acquire(resolved.workspace_name) as lock_stats:
|
|
sync_meta = self.sync_service.sync_for_query(
|
|
branch=resolved.branch,
|
|
memory_profile=None,
|
|
require_latest=request.require_latest,
|
|
)
|
|
|
|
results, qmd_command = self.qmd_client.run_query(
|
|
index_name=sync_meta.qmd_index,
|
|
collection_name=sync_meta.qmd_collection,
|
|
query_type=request.query_type,
|
|
query=request.query,
|
|
n=request.n,
|
|
)
|
|
|
|
debug: dict[str, Any] | None = None
|
|
if request.debug:
|
|
debug = {
|
|
"lock_wait_ms": lock_stats.wait_ms,
|
|
"workspace_changed": sync_meta.changed,
|
|
"previous_commit": sync_meta.previous_commit,
|
|
"qmd_index": sync_meta.qmd_index,
|
|
"qmd_command": qmd_command,
|
|
"update_ran": sync_meta.update_ran,
|
|
"embed_ran": sync_meta.embed_ran,
|
|
"embed_error": sync_meta.embed_error,
|
|
}
|
|
|
|
return QueryResponse(
|
|
ok=True,
|
|
branch=sync_meta.branch,
|
|
resolved_workspace=sync_meta.workspace_relative_path,
|
|
commit_hash=sync_meta.commit_hash,
|
|
synced_at=sync_meta.synced_at,
|
|
query_type=request.query_type,
|
|
results=results,
|
|
qmd_collection=sync_meta.qmd_collection,
|
|
debug=debug,
|
|
)
|
|
|
|
def handle_sync(self, request: SyncRequest) -> SyncResponse:
|
|
resolved = self.sync_service.workspace_manager.resolve_workspace(
|
|
branch=request.branch,
|
|
memory_profile=request.memory_profile,
|
|
)
|
|
|
|
with self.lock_manager.acquire(resolved.workspace_name) as lock_stats:
|
|
sync_meta = self.sync_service.sync_explicit(
|
|
branch=resolved.branch,
|
|
memory_profile=None,
|
|
require_latest=request.require_latest,
|
|
)
|
|
debug = {
|
|
"lock_wait_ms": lock_stats.wait_ms,
|
|
"workspace_changed": sync_meta.changed,
|
|
"previous_commit": sync_meta.previous_commit,
|
|
"qmd_index": sync_meta.qmd_index,
|
|
"update_ran": sync_meta.update_ran,
|
|
"embed_ran": sync_meta.embed_ran,
|
|
"embed_error": sync_meta.embed_error,
|
|
}
|
|
return SyncResponse(
|
|
ok=True,
|
|
branch=sync_meta.branch,
|
|
resolved_workspace=sync_meta.workspace_relative_path,
|
|
commit_hash=sync_meta.commit_hash,
|
|
synced_at=sync_meta.synced_at,
|
|
qmd_collection=sync_meta.qmd_collection,
|
|
debug=debug,
|
|
)
|