Browse
+
+
+
+ Select a connector to browse.
+
+ Selected items (0)
+
+
+
+
+
+
+
+
+ OAuth completed for project {project_id}. You can close this window.
+Next: POST /collect/projects/{project_id}/connect/{connector_id} or browse the remote tree.
{message}
+""" + + +def create_oauth_router(oauth_manager: OAuthManager) -> APIRouter: + router = APIRouter(tags=["collect-oauth"]) + + @router.get("/collect/oauth/{connector_id}/authorize") + async def oauth_authorize( + connector_id: str, + project_id: str, + redirect_after: Optional[str] = None, + ): + if connector_id not in list_connectors(): + raise HTTPException(status_code=404, detail=f"Unknown connector: {connector_id}") + if not oauth_manager.supports_connector(connector_id): + raise HTTPException( + status_code=400, + detail=f"OAuth not configured for connector: {connector_id}", + ) + try: + url = oauth_manager.build_authorize_url( + connector_id=connector_id, + project_id=project_id, + redirect_after=redirect_after, + ) + except OAuthError as exc: + raise HTTPException(status_code=503, detail=str(exc)) from exc + return RedirectResponse(url, status_code=302) + + @router.get("/collect/oauth/{connector_id}/callback") + async def oauth_callback( + connector_id: str, + request: Request, + code: Optional[str] = None, + state: Optional[str] = None, + error: Optional[str] = None, + error_description: Optional[str] = None, + format: Optional[str] = None, + ): + if connector_id not in list_connectors(): + raise HTTPException(status_code=404, detail=f"Unknown connector: {connector_id}") + + if error: + message = error_description or error + if _wants_json(request, format): + raise HTTPException(status_code=400, detail=message) + return HTMLResponse(_error_html(message), status_code=400) + + try: + auth, project_id, redirect_after = oauth_manager.handle_callback( + connector_id=connector_id, + code=code or "", + state=state or "", + ) + except OAuthError as exc: + if _wants_json(request, format): + raise HTTPException(status_code=400, detail=str(exc)) from exc + return HTMLResponse(_error_html(str(exc)), status_code=400) + + if redirect_after: + return RedirectResponse(redirect_after, status_code=302) + + payload = { + "status": "connected", + "project_id": project_id, + "connector_id": connector_id, + "has_refresh_token": bool(auth.refresh_token), + "scopes": auth.scopes, + } + if _wants_json(request, format): + return JSONResponse(payload) + return HTMLResponse(_success_html(connector_id, project_id)) + + return router diff --git a/api/collectors/core/orchestrator.py b/api/collectors/core/orchestrator.py new file mode 100644 index 000000000..81c2991ca --- /dev/null +++ b/api/collectors/core/orchestrator.py @@ -0,0 +1,163 @@ +"""Collection run orchestration.""" + +from __future__ import annotations + +import logging +import os +from datetime import datetime +from typing import Any, Dict, Optional + +from collectors.base.auth_context import AuthContext +from collectors.base.types import CollectionJob, CollectionJobStatus +from collectors.core.logging_manager import LoggingManager +from collectors.core.registry import get_connector_class +from collectors.core.security_manager import SecurityManager +from collectors.core.state_manager import StateManager +from collectors.ingestion.ingestion_pipeline import IngestionPipeline +from collectors.processing.processing_engine_client import ProcessingEngineClient +from collectors.storage.filesystem_store import FilesystemStore +from collectors.storage.sqlite_store import SQLiteMetadataStore + + +class Orchestrator: + """Coordinates connector auth, ingestion, and optional processing.""" + + def __init__( + self, + state_manager: Optional[StateManager] = None, + security_manager: Optional[SecurityManager] = None, + base_output_dir: Optional[str] = None, + ): + self.state = state_manager or StateManager() + self.security = security_manager or SecurityManager() + self.logger = LoggingManager.configure() + self.base_output_dir = base_output_dir or os.getenv( + "COLLECTOR_OUTPUT_BASE", "/data/input/collections" + ) + self.processing_client = ProcessingEngineClient() + + def run_collection(self, job: CollectionJob) -> CollectionJob: + connector_id = job.connector_id + if self.state.is_circuit_open(connector_id): + job.status = CollectionJobStatus.FAILED + job.error_message = "Circuit breaker open for connector" + job.completed_at = datetime.utcnow() + self.state.save_job(job) + return job + + job.status = CollectionJobStatus.RUNNING + job.started_at = datetime.utcnow() + output_dir = os.path.join(self.base_output_dir, job.job_id) + job.output_dir = output_dir + self.state.save_job(job) + + try: + auth = self.security.build_auth_context(connector_id) + auth = self.security.refresh_auth_if_needed(auth) + connector_cls = get_connector_class(connector_id) + collector = connector_cls(auth=auth, config=job.config) + + health = collector.health_check() + if not health.healthy: + raise RuntimeError(f"Connector unhealthy: {health.message}") + + filesystem = FilesystemStore(output_dir) + metadata_store = SQLiteMetadataStore( + db_path=os.path.join(output_dir, "metadata.db") + ) + pipeline = IngestionPipeline( + collector=collector, + filesystem=filesystem, + metadata_store=metadata_store, + state_manager=self.state, + logger=self.logger, + ) + + since = self._parse_since(job.config) + result = pipeline.run(since=since) + + job.items_collected = result["items_collected"] + job.items_failed = result["items_failed"] + job.status = CollectionJobStatus.COMPLETED + job.completed_at = datetime.utcnow() + self.state.record_success(connector_id, sync_token=result.get("sync_token")) + self.state.save_job(job) + + if job.auto_chain: + self._chain_processing(job, output_dir) + + LoggingManager.log_event( + self.logger, + "collection_completed", + job_id=job.job_id, + connector_id=connector_id, + items_collected=job.items_collected, + ) + except Exception as exc: + job.status = CollectionJobStatus.FAILED + job.error_message = str(exc) + job.completed_at = datetime.utcnow() + self.state.record_failure(connector_id) + self.state.save_job(job) + LoggingManager.log_event( + self.logger, + "collection_failed", + level=logging.ERROR, + job_id=job.job_id, + connector_id=connector_id, + error=str(exc), + ) + + return job + + def _chain_processing(self, job: CollectionJob, input_dir: str) -> None: + project_config = job.config.get("processing", {}) + self.processing_client.trigger_processing( + job_id=job.job_id, + input_dir=input_dir, + project_config=project_config, + ) + + def _parse_since(self, config: Dict[str, Any]) -> Optional[datetime]: + since_str = config.get("since") + if not since_str: + return None + return datetime.fromisoformat(since_str) + + def get_connector_health(self, connector_id: str) -> Dict[str, Any]: + try: + auth = self.security.build_auth_context(connector_id) + connector_cls = get_connector_class(connector_id) + collector = connector_cls(auth=auth) + health = collector.health_check() + state = self.state.get_connector_state(connector_id) + return { + "connector_id": connector_id, + "healthy": health.healthy, + "message": health.message, + "circuit_breaker": state.get("circuit_breaker_state"), + "last_successful_run_at": state.get("last_successful_run_at"), + } + except Exception as exc: + return { + "connector_id": connector_id, + "healthy": False, + "message": str(exc), + } + + def subsystem_status(self) -> Dict[str, Any]: + from collectors.core.registry import list_connectors + + connectors = list_connectors() + open_circuits = [ + cid + for cid in connectors + if self.state.is_circuit_open(cid, cooldown_seconds=0) + and self.state.get_connector_state(cid)["circuit_breaker_state"] == "open" + ] + return { + "status": "healthy" if not open_circuits else "degraded", + "registered_connectors": connectors, + "open_circuits": open_circuits, + "processing_busy": self.processing_client.is_busy(), + } diff --git a/api/collectors/core/registry.py b/api/collectors/core/registry.py new file mode 100644 index 000000000..69dc7eead --- /dev/null +++ b/api/collectors/core/registry.py @@ -0,0 +1,37 @@ +"""Connector registry with decorator-based registration.""" + +from __future__ import annotations + +from typing import Callable, Dict, List, Type + +from collectors.base.base_collector import BaseCollector + +_REGISTRY: Dict[str, Type[BaseCollector]] = {} + + +def register_connector(connector_id: str) -> Callable[[Type[BaseCollector]], Type[BaseCollector]]: + """Decorator to register a connector class by ID.""" + + def decorator(cls: Type[BaseCollector]) -> Type[BaseCollector]: + if connector_id in _REGISTRY: + raise ValueError(f"Connector already registered: {connector_id}") + _REGISTRY[connector_id] = cls + cls._registered_connector_id = connector_id # type: ignore[attr-defined] + return cls + + return decorator + + +def get_connector_class(connector_id: str) -> Type[BaseCollector]: + if connector_id not in _REGISTRY: + raise KeyError(f"Unknown connector: {connector_id}") + return _REGISTRY[connector_id] + + +def list_connectors() -> List[str]: + return sorted(_REGISTRY.keys()) + + +def clear_registry() -> None: + """Test helper to reset registry.""" + _REGISTRY.clear() diff --git a/api/collectors/core/security_manager.py b/api/collectors/core/security_manager.py new file mode 100644 index 000000000..cc929999e --- /dev/null +++ b/api/collectors/core/security_manager.py @@ -0,0 +1,163 @@ +"""Secrets management and log sanitization.""" + +from __future__ import annotations + +import json +import os +import re +from pathlib import Path +from typing import TYPE_CHECKING, Any, Callable, Dict, Optional + +from cryptography.fernet import Fernet, InvalidToken + +from collectors.base.auth_context import AuthContext + +if TYPE_CHECKING: + from collectors.core.token_store import TokenStore + + +class SecurityManager: + """Load secrets from env/files, encrypt tokens, sanitize logs.""" + + SENSITIVE_KEYS = frozenset( + { + "access_token", + "refresh_token", + "client_secret", + "password", + "api_key", + "authorization", + } + ) + + def __init__( + self, + encryption_key: Optional[str] = None, + secrets_dir: Optional[str] = None, + ): + key = encryption_key or os.getenv("COLLECTOR_ENCRYPTION_KEY") + self._fernet: Optional[Fernet] = None + if key: + self._fernet = Fernet(key.encode() if isinstance(key, str) else key) + self.secrets_dir = Path(secrets_dir or os.getenv("COLLECTOR_SECRETS_DIR", "/app/secrets")) + self._refresh_hooks: Dict[str, Callable[[AuthContext], AuthContext]] = {} + + def encrypt_token(self, plaintext: str) -> str: + if not self._fernet: + return plaintext + return self._fernet.encrypt(plaintext.encode()).decode() + + def decrypt_token(self, ciphertext: str) -> str: + if not self._fernet: + return ciphertext + try: + return self._fernet.decrypt(ciphertext.encode()).decode() + except InvalidToken as exc: + raise ValueError("Failed to decrypt token") from exc + + def get_env_secret(self, key: str, default: Optional[str] = None) -> Optional[str]: + value = os.getenv(key, default) + if value and value.startswith("enc:"): + return self.decrypt_token(value[4:]) + return value + + def load_connector_secrets(self, connector_id: str) -> Dict[str, Any]: + """Load secrets from env prefix or JSON file.""" + prefix = f"COLLECTOR_{connector_id.upper().replace('-', '_')}_" + secrets: Dict[str, Any] = {} + for env_key, env_val in os.environ.items(): + if env_key.startswith(prefix): + field = env_key[len(prefix) :].lower() + secrets[field] = self._maybe_decrypt(env_val) + + file_path = self.secrets_dir / f"{connector_id}.json" + if file_path.is_file(): + with open(file_path, encoding="utf-8") as fh: + file_secrets = json.load(fh) + for k, v in file_secrets.items(): + if isinstance(v, str): + secrets[k] = self._maybe_decrypt(v) + else: + secrets[k] = v + return secrets + + def _maybe_decrypt(self, value: str) -> str: + if value.startswith("enc:"): + return self.decrypt_token(value[4:]) + return value + + def get_connector_credentials( + self, + connector_id: str, + project_id: Optional[str] = None, + token_store: Optional["TokenStore"] = None, + ) -> Dict[str, Any]: + """Resolve credentials: project OAuth tokens first, then global secrets.""" + secrets = self.load_connector_secrets(connector_id) + if project_id and token_store: + stored = token_store.get_auth_context(project_id, connector_id, self) + if stored and (stored.access_token or stored.refresh_token): + merged = dict(secrets) + if stored.access_token: + merged["access_token"] = stored.access_token + if stored.refresh_token: + merged["refresh_token"] = stored.refresh_token + merged["_project_id"] = project_id + return merged + return secrets + + def build_auth_context( + self, + connector_id: str, + project_id: Optional[str] = None, + token_store: Optional["TokenStore"] = None, + ) -> AuthContext: + secrets = self.get_connector_credentials(connector_id, project_id, token_store) + project = secrets.pop("_project_id", None) or project_id + extra = {k: v for k, v in secrets.items() if k not in ("access_token", "refresh_token")} + if project: + extra["project_id"] = project + return AuthContext( + connector_id=connector_id, + access_token=secrets.get("access_token"), + refresh_token=secrets.get("refresh_token"), + extra=extra, + ) + + def register_refresh_hook( + self, + connector_id: str, + hook: Callable[[AuthContext], AuthContext], + ) -> None: + self._refresh_hooks[connector_id] = hook + + def refresh_auth_if_needed(self, auth: AuthContext) -> AuthContext: + if not auth.is_expired(): + return auth + hook = self._refresh_hooks.get(auth.connector_id) + if hook: + return hook(auth) + return auth + + def sanitize_for_log(self, data: Any) -> Any: + if isinstance(data, dict): + return { + k: ("***REDACTED***" if k.lower() in self.SENSITIVE_KEYS else self.sanitize_for_log(v)) + for k, v in data.items() + } + if isinstance(data, list): + return [self.sanitize_for_log(item) for item in data] + if isinstance(data, str): + return self._sanitize_string(data) + return data + + def _sanitize_string(self, text: str) -> str: + patterns = [ + (r"(Bearer\s+)[A-Za-z0-9\-._~+/]+=*", r"\1***REDACTED***"), + (r'("access_token"\s*:\s*")[^"]+(")', r"\1***REDACTED***\2"), + (r'("refresh_token"\s*:\s*")[^"]+(")', r"\1***REDACTED***\2"), + ] + result = text + for pattern, repl in patterns: + result = re.sub(pattern, repl, result) + return result diff --git a/api/collectors/core/selection_store.py b/api/collectors/core/selection_store.py new file mode 100644 index 000000000..af004cb41 --- /dev/null +++ b/api/collectors/core/selection_store.py @@ -0,0 +1,133 @@ +"""Persist per-project browse selections (not downloaded files).""" + +from __future__ import annotations + +import json +import os +import sqlite3 +import threading +from contextlib import contextmanager +from datetime import datetime +from typing import Any, Dict, List, Optional + +from collectors.base.types import BrowseItemType, SelectedItem + + +class SelectionStore: + """SQLite store keyed by (project_id, connector_id).""" + + def __init__(self, db_path: Optional[str] = None): + self.db_path = db_path or os.getenv("COLLECTOR_STATE_DB", "/app/db/collector_state.db") + self._lock = threading.Lock() + self._init_db() + + @contextmanager + def _connection(self): + conn = sqlite3.connect(self.db_path, check_same_thread=False) + conn.row_factory = sqlite3.Row + try: + yield conn + conn.commit() + finally: + conn.close() + + def _init_db(self) -> None: + os.makedirs(os.path.dirname(self.db_path), exist_ok=True) + with self._lock, self._connection() as conn: + conn.execute( + """ + CREATE TABLE IF NOT EXISTS project_selections ( + project_id TEXT NOT NULL, + connector_id TEXT NOT NULL, + items TEXT NOT NULL DEFAULT '[]', + updated_at TEXT, + PRIMARY KEY (project_id, connector_id) + ) + """ + ) + + def save_selections( + self, + project_id: str, + connector_id: str, + items: List[SelectedItem], + ) -> Dict[str, Any]: + payload = [self._item_to_dict(item) for item in items] + updated_at = datetime.utcnow().isoformat() + with self._lock, self._connection() as conn: + conn.execute( + """ + INSERT INTO project_selections (project_id, connector_id, items, updated_at) + VALUES (?, ?, ?, ?) + ON CONFLICT(project_id, connector_id) DO UPDATE SET + items = excluded.items, + updated_at = excluded.updated_at + """, + (project_id, connector_id, json.dumps(payload), updated_at), + ) + return { + "project_id": project_id, + "connector_id": connector_id, + "items": payload, + "updated_at": updated_at, + } + + def get_selections( + self, + project_id: str, + connector_id: Optional[str] = None, + ) -> List[Dict[str, Any]]: + with self._lock, self._connection() as conn: + if connector_id: + rows = conn.execute( + "SELECT * FROM project_selections WHERE project_id = ? AND connector_id = ?", + (project_id, connector_id), + ).fetchall() + else: + rows = conn.execute( + "SELECT * FROM project_selections WHERE project_id = ?", + (project_id,), + ).fetchall() + return [self._row_to_dict(row) for row in rows] + + def get_selected_items( + self, + project_id: str, + connector_id: str, + ) -> List[SelectedItem]: + rows = self.get_selections(project_id, connector_id) + if not rows: + return [] + return [self._dict_to_item(entry) for entry in rows[0]["items"]] + + @staticmethod + def _item_to_dict(item: SelectedItem) -> Dict[str, Any]: + return { + "id": item.id, + "type": item.type.value if isinstance(item.type, BrowseItemType) else item.type, + "name": item.name, + "path": item.path, + } + + @staticmethod + def _dict_to_item(data: Dict[str, Any]) -> SelectedItem: + item_type = data.get("type", "file") + if isinstance(item_type, BrowseItemType): + parsed_type = item_type + else: + parsed_type = BrowseItemType(item_type) + return SelectedItem( + id=str(data["id"]), + type=parsed_type, + name=str(data.get("name", "")), + path=str(data.get("path", "")), + ) + + @staticmethod + def _row_to_dict(row: sqlite3.Row) -> Dict[str, Any]: + return { + "project_id": row["project_id"], + "connector_id": row["connector_id"], + "items": json.loads(row["items"] or "[]"), + "updated_at": row["updated_at"], + } diff --git a/api/collectors/core/state_manager.py b/api/collectors/core/state_manager.py new file mode 100644 index 000000000..1daae2310 --- /dev/null +++ b/api/collectors/core/state_manager.py @@ -0,0 +1,275 @@ +"""SQLite-backed connector state persistence.""" + +from __future__ import annotations + +import json +import os +import sqlite3 +import threading +from contextlib import contextmanager +from datetime import datetime +from typing import Any, Dict, List, Optional + +from collectors.base.types import CircuitBreakerState, CollectionJob, CollectionJobStatus + + +class StateManager: + """Persists sync tokens, job state, and circuit breaker status.""" + + def __init__(self, db_path: Optional[str] = None): + self.db_path = db_path or os.getenv("COLLECTOR_STATE_DB", "/app/db/collector_state.db") + self._lock = threading.Lock() + self._init_db() + + @contextmanager + def _connection(self): + conn = sqlite3.connect(self.db_path, check_same_thread=False) + conn.row_factory = sqlite3.Row + try: + yield conn + conn.commit() + finally: + conn.close() + + def _init_db(self) -> None: + os.makedirs(os.path.dirname(self.db_path), exist_ok=True) + with self._lock, self._connection() as conn: + conn.executescript( + """ + CREATE TABLE IF NOT EXISTS connector_state ( + connector_id TEXT PRIMARY KEY, + last_sync_token TEXT, + last_successful_run_at TEXT, + failed_item_ids TEXT DEFAULT '[]', + circuit_breaker_state TEXT DEFAULT 'closed', + circuit_failure_count INTEGER DEFAULT 0, + circuit_opened_at TEXT + ); + CREATE TABLE IF NOT EXISTS collection_jobs ( + job_id TEXT PRIMARY KEY, + connector_id TEXT NOT NULL, + status TEXT NOT NULL, + created_at TEXT, + started_at TEXT, + completed_at TEXT, + output_dir TEXT, + items_collected INTEGER DEFAULT 0, + items_failed INTEGER DEFAULT 0, + error_message TEXT, + auto_chain INTEGER DEFAULT 0, + config TEXT DEFAULT '{}' + ); + CREATE TABLE IF NOT EXISTS schedules ( + schedule_id TEXT PRIMARY KEY, + connector_id TEXT NOT NULL, + cron_expression TEXT NOT NULL, + auto_chain INTEGER DEFAULT 0, + config TEXT DEFAULT '{}', + enabled INTEGER DEFAULT 1 + ); + """ + ) + + def get_connector_state(self, connector_id: str) -> Dict[str, Any]: + with self._lock, self._connection() as conn: + row = conn.execute( + "SELECT * FROM connector_state WHERE connector_id = ?", + (connector_id,), + ).fetchone() + if not row: + return { + "connector_id": connector_id, + "last_sync_token": None, + "last_successful_run_at": None, + "failed_item_ids": [], + "circuit_breaker_state": CircuitBreakerState.CLOSED.value, + "circuit_failure_count": 0, + } + return { + "connector_id": row["connector_id"], + "last_sync_token": row["last_sync_token"], + "last_successful_run_at": row["last_successful_run_at"], + "failed_item_ids": json.loads(row["failed_item_ids"] or "[]"), + "circuit_breaker_state": row["circuit_breaker_state"], + "circuit_failure_count": row["circuit_failure_count"], + "circuit_opened_at": row["circuit_opened_at"], + } + + def update_connector_state(self, connector_id: str, **fields: Any) -> None: + state = self.get_connector_state(connector_id) + state.update(fields) + if "failed_item_ids" in fields and isinstance(fields["failed_item_ids"], list): + state["failed_item_ids"] = fields["failed_item_ids"] + with self._lock, self._connection() as conn: + conn.execute( + """ + INSERT INTO connector_state ( + connector_id, last_sync_token, last_successful_run_at, + failed_item_ids, circuit_breaker_state, circuit_failure_count, + circuit_opened_at + ) VALUES (?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(connector_id) DO UPDATE SET + last_sync_token = excluded.last_sync_token, + last_successful_run_at = excluded.last_successful_run_at, + failed_item_ids = excluded.failed_item_ids, + circuit_breaker_state = excluded.circuit_breaker_state, + circuit_failure_count = excluded.circuit_failure_count, + circuit_opened_at = excluded.circuit_opened_at + """, + ( + connector_id, + state.get("last_sync_token"), + state.get("last_successful_run_at"), + json.dumps(state.get("failed_item_ids", [])), + state.get("circuit_breaker_state", CircuitBreakerState.CLOSED.value), + state.get("circuit_failure_count", 0), + state.get("circuit_opened_at"), + ), + ) + + def record_success(self, connector_id: str, sync_token: Optional[str] = None) -> None: + self.update_connector_state( + connector_id, + last_sync_token=sync_token, + last_successful_run_at=datetime.utcnow().isoformat(), + circuit_breaker_state=CircuitBreakerState.CLOSED.value, + circuit_failure_count=0, + circuit_opened_at=None, + ) + + def record_failure(self, connector_id: str, threshold: int = 5) -> CircuitBreakerState: + state = self.get_connector_state(connector_id) + count = int(state.get("circuit_failure_count", 0)) + 1 + cb_state = CircuitBreakerState.CLOSED + opened_at = state.get("circuit_opened_at") + if count >= threshold: + cb_state = CircuitBreakerState.OPEN + opened_at = datetime.utcnow().isoformat() + self.update_connector_state( + connector_id, + circuit_failure_count=count, + circuit_breaker_state=cb_state.value, + circuit_opened_at=opened_at, + ) + return cb_state + + def is_circuit_open(self, connector_id: str, cooldown_seconds: int = 300) -> bool: + state = self.get_connector_state(connector_id) + if state["circuit_breaker_state"] != CircuitBreakerState.OPEN.value: + return False + opened_at = state.get("circuit_opened_at") + if not opened_at: + return True + elapsed = (datetime.utcnow() - datetime.fromisoformat(opened_at)).total_seconds() + if elapsed >= cooldown_seconds: + self.update_connector_state( + connector_id, + circuit_breaker_state=CircuitBreakerState.HALF_OPEN.value, + ) + return False + return True + + def add_failed_item(self, connector_id: str, item_id: str) -> None: + state = self.get_connector_state(connector_id) + failed: List[str] = list(state.get("failed_item_ids", [])) + if item_id not in failed: + failed.append(item_id) + self.update_connector_state(connector_id, failed_item_ids=failed) + + def save_job(self, job: CollectionJob) -> None: + with self._lock, self._connection() as conn: + conn.execute( + """ + INSERT INTO collection_jobs ( + job_id, connector_id, status, created_at, started_at, + completed_at, output_dir, items_collected, items_failed, + error_message, auto_chain, config + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(job_id) DO UPDATE SET + status = excluded.status, + started_at = excluded.started_at, + completed_at = excluded.completed_at, + output_dir = excluded.output_dir, + items_collected = excluded.items_collected, + items_failed = excluded.items_failed, + error_message = excluded.error_message, + auto_chain = excluded.auto_chain, + config = excluded.config + """, + ( + job.job_id, + job.connector_id, + job.status.value, + job.created_at.isoformat() if job.created_at else None, + job.started_at.isoformat() if job.started_at else None, + job.completed_at.isoformat() if job.completed_at else None, + job.output_dir, + job.items_collected, + job.items_failed, + job.error_message, + 1 if job.auto_chain else 0, + json.dumps(job.config), + ), + ) + + def get_job(self, job_id: str) -> Optional[CollectionJob]: + with self._lock, self._connection() as conn: + row = conn.execute( + "SELECT * FROM collection_jobs WHERE job_id = ?", + (job_id,), + ).fetchone() + if not row: + return None + return CollectionJob( + job_id=row["job_id"], + connector_id=row["connector_id"], + status=CollectionJobStatus(row["status"]), + created_at=_parse_dt(row["created_at"]), + started_at=_parse_dt(row["started_at"]), + completed_at=_parse_dt(row["completed_at"]), + output_dir=row["output_dir"], + items_collected=row["items_collected"], + items_failed=row["items_failed"], + error_message=row["error_message"], + auto_chain=bool(row["auto_chain"]), + config=json.loads(row["config"] or "{}"), + ) + + def save_schedule( + self, + schedule_id: str, + connector_id: str, + cron_expression: str, + auto_chain: bool = False, + config: Optional[Dict[str, Any]] = None, + ) -> None: + with self._lock, self._connection() as conn: + conn.execute( + """ + INSERT INTO schedules (schedule_id, connector_id, cron_expression, auto_chain, config) + VALUES (?, ?, ?, ?, ?) + ON CONFLICT(schedule_id) DO UPDATE SET + connector_id = excluded.connector_id, + cron_expression = excluded.cron_expression, + auto_chain = excluded.auto_chain, + config = excluded.config + """, + ( + schedule_id, + connector_id, + cron_expression, + 1 if auto_chain else 0, + json.dumps(config or {}), + ), + ) + + def list_schedules(self) -> List[Dict[str, Any]]: + with self._lock, self._connection() as conn: + rows = conn.execute("SELECT * FROM schedules WHERE enabled = 1").fetchall() + return [dict(row) for row in rows] + + +def _parse_dt(value: Optional[str]) -> Optional[datetime]: + if not value: + return None + return datetime.fromisoformat(value) diff --git a/api/collectors/core/token_store.py b/api/collectors/core/token_store.py new file mode 100644 index 000000000..bbc3300ba --- /dev/null +++ b/api/collectors/core/token_store.py @@ -0,0 +1,212 @@ +"""Project-scoped OAuth token persistence.""" + +from __future__ import annotations + +import json +import os +import sqlite3 +import threading +from contextlib import contextmanager +from datetime import datetime +from typing import Any, Dict, Optional, TYPE_CHECKING + +from collectors.base.auth_context import AuthContext + +if TYPE_CHECKING: + from collectors.core.security_manager import SecurityManager + + +class TokenStore: + """SQLite store for OAuth tokens keyed by (project_id, connector_id).""" + + def __init__(self, db_path: Optional[str] = None): + self.db_path = db_path or os.getenv("COLLECTOR_STATE_DB", "/app/db/collector_state.db") + self._lock = threading.Lock() + self._init_db() + + @contextmanager + def _connection(self): + conn = sqlite3.connect(self.db_path, check_same_thread=False) + conn.row_factory = sqlite3.Row + try: + yield conn + conn.commit() + finally: + conn.close() + + def _init_db(self) -> None: + os.makedirs(os.path.dirname(self.db_path), exist_ok=True) + with self._lock, self._connection() as conn: + conn.executescript( + """ + CREATE TABLE IF NOT EXISTS oauth_tokens ( + project_id TEXT NOT NULL, + connector_id TEXT NOT NULL, + access_token TEXT, + refresh_token TEXT, + token_type TEXT DEFAULT 'Bearer', + expires_at TEXT, + scopes TEXT DEFAULT '[]', + updated_at TEXT, + PRIMARY KEY (project_id, connector_id) + ); + CREATE TABLE IF NOT EXISTS oauth_states ( + state TEXT PRIMARY KEY, + project_id TEXT NOT NULL, + connector_id TEXT NOT NULL, + code_verifier TEXT NOT NULL, + redirect_after TEXT, + expires_at TEXT NOT NULL, + created_at TEXT NOT NULL + ); + """ + ) + + def save_auth_context( + self, + project_id: str, + connector_id: str, + auth: AuthContext, + security: SecurityManager, + ) -> None: + access = self._encrypt_field(auth.access_token, security) + refresh = self._encrypt_field(auth.refresh_token, security) + updated_at = datetime.utcnow().isoformat() + with self._lock, self._connection() as conn: + conn.execute( + """ + INSERT INTO oauth_tokens ( + project_id, connector_id, access_token, refresh_token, + token_type, expires_at, scopes, updated_at + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?) + ON CONFLICT(project_id, connector_id) DO UPDATE SET + access_token = excluded.access_token, + refresh_token = excluded.refresh_token, + token_type = excluded.token_type, + expires_at = excluded.expires_at, + scopes = excluded.scopes, + updated_at = excluded.updated_at + """, + ( + project_id, + connector_id, + access, + refresh, + auth.token_type, + auth.expires_at.isoformat() if auth.expires_at else None, + json.dumps(auth.scopes), + updated_at, + ), + ) + + def get_auth_context( + self, + project_id: str, + connector_id: str, + security: SecurityManager, + ) -> Optional[AuthContext]: + with self._lock, self._connection() as conn: + row = conn.execute( + """ + SELECT * FROM oauth_tokens + WHERE project_id = ? AND connector_id = ? + """, + (project_id, connector_id), + ).fetchone() + if not row: + return None + access = self._decrypt_field(row["access_token"], security) + refresh = self._decrypt_field(row["refresh_token"], security) + if not access and not refresh: + return None + expires_at = None + if row["expires_at"]: + expires_at = datetime.fromisoformat(row["expires_at"]) + scopes = json.loads(row["scopes"] or "[]") + return AuthContext( + connector_id=connector_id, + access_token=access, + refresh_token=refresh, + token_type=row["token_type"] or "Bearer", + expires_at=expires_at, + scopes=scopes, + extra={"project_id": project_id}, + ) + + def has_tokens(self, project_id: str, connector_id: str) -> bool: + with self._lock, self._connection() as conn: + row = conn.execute( + """ + SELECT access_token, refresh_token FROM oauth_tokens + WHERE project_id = ? AND connector_id = ? + """, + (project_id, connector_id), + ).fetchone() + if not row: + return False + return bool(row["access_token"] or row["refresh_token"]) + + def save_oauth_state( + self, + state: str, + project_id: str, + connector_id: str, + code_verifier: str, + expires_at: datetime, + redirect_after: Optional[str] = None, + ) -> None: + now = datetime.utcnow().isoformat() + with self._lock, self._connection() as conn: + conn.execute( + """ + INSERT INTO oauth_states ( + state, project_id, connector_id, code_verifier, + redirect_after, expires_at, created_at + ) VALUES (?, ?, ?, ?, ?, ?, ?) + """, + ( + state, + project_id, + connector_id, + code_verifier, + redirect_after, + expires_at.isoformat(), + now, + ), + ) + + def consume_oauth_state(self, state: str, connector_id: str) -> Optional[Dict[str, Any]]: + with self._lock, self._connection() as conn: + row = conn.execute( + "SELECT * FROM oauth_states WHERE state = ? AND connector_id = ?", + (state, connector_id), + ).fetchone() + if not row: + return None + conn.execute("DELETE FROM oauth_states WHERE state = ?", (state,)) + expires_at = datetime.fromisoformat(row["expires_at"]) + if datetime.utcnow() > expires_at: + return None + return { + "project_id": row["project_id"], + "connector_id": row["connector_id"], + "code_verifier": row["code_verifier"], + "redirect_after": row["redirect_after"], + } + + @staticmethod + def _encrypt_field(value: Optional[str], security: SecurityManager) -> Optional[str]: + if not value: + return None + encrypted = security.encrypt_token(value) + if encrypted != value: + return f"enc:{encrypted}" + return value + + @staticmethod + def _decrypt_field(value: Optional[str], security: SecurityManager) -> Optional[str]: + if not value: + return None + if value.startswith("enc:"): + return security.decrypt_token(value[4:]) + return value diff --git a/api/collectors/core/trigger_manager.py b/api/collectors/core/trigger_manager.py new file mode 100644 index 000000000..e437d1d86 --- /dev/null +++ b/api/collectors/core/trigger_manager.py @@ -0,0 +1,119 @@ +"""Manual and scheduled collection triggers.""" + +from __future__ import annotations + +import logging +import threading +import uuid +from datetime import datetime +from typing import Any, Callable, Dict, Optional + +from apscheduler.schedulers.background import BackgroundScheduler +from apscheduler.triggers.cron import CronTrigger + +from collectors.base.types import CollectionJob, CollectionJobStatus +from collectors.core.logging_manager import LoggingManager +from collectors.core.state_manager import StateManager + + +class TriggerManager: + """Handles manual triggers and APScheduler cron jobs.""" + + def __init__( + self, + state_manager: StateManager, + run_callback: Callable[[CollectionJob], None], + ): + self.state = state_manager + self.run_callback = run_callback + self.logger = LoggingManager.configure() + self._scheduler = BackgroundScheduler() + self._scheduler.start() + self._shutdown = threading.Event() + + def trigger_manual( + self, + connector_id: str, + auto_chain: bool = False, + config: Optional[Dict[str, Any]] = None, + ) -> CollectionJob: + job = CollectionJob( + job_id=str(uuid.uuid4()), + connector_id=connector_id, + status=CollectionJobStatus.PENDING, + created_at=datetime.utcnow(), + auto_chain=auto_chain, + config=config or {}, + ) + self.state.save_job(job) + if not self._shutdown.is_set(): + thread = threading.Thread(target=self._execute_job, args=(job,), daemon=True) + thread.start() + return job + + def _execute_job(self, job: CollectionJob) -> None: + try: + self.run_callback(job) + except Exception as exc: + LoggingManager.log_event( + self.logger, + "trigger_job_failed", + level=logging.ERROR, + job_id=job.job_id, + error=str(exc), + ) + + def add_schedule( + self, + schedule_id: str, + connector_id: str, + cron_expression: str, + auto_chain: bool = False, + config: Optional[Dict[str, Any]] = None, + ) -> None: + parts = cron_expression.split() + if len(parts) != 5: + raise ValueError("Cron expression must have 5 fields: min hour dom month dow") + minute, hour, day, month, day_of_week = parts + trigger = CronTrigger( + minute=minute, + hour=hour, + day=day, + month=month, + day_of_week=day_of_week, + ) + self.state.save_schedule(schedule_id, connector_id, cron_expression, auto_chain, config) + + def _scheduled_run() -> None: + if self._shutdown.is_set(): + return + self.trigger_manual(connector_id, auto_chain=auto_chain, config=config) + + self._scheduler.add_job( + _scheduled_run, + trigger=trigger, + id=schedule_id, + replace_existing=True, + ) + LoggingManager.log_event( + self.logger, + "schedule_registered", + schedule_id=schedule_id, + connector_id=connector_id, + cron=cron_expression, + ) + + def load_schedules_from_db(self) -> None: + for row in self.state.list_schedules(): + self.add_schedule( + schedule_id=row["schedule_id"], + connector_id=row["connector_id"], + cron_expression=row["cron_expression"], + auto_chain=bool(row["auto_chain"]), + config=__import__("json").loads(row["config"] or "{}"), + ) + + def shutdown(self) -> None: + self._shutdown.set() + if self._scheduler.running: + self._scheduler.shutdown(wait=False) diff --git a/api/collectors/docs/architecture.md b/api/collectors/docs/architecture.md new file mode 100644 index 000000000..09b27a7fe --- /dev/null +++ b/api/collectors/docs/architecture.md @@ -0,0 +1,40 @@ +# External Data Collection Framework + +FreeEed collectors ingest documents from cloud storage providers into `/data/input/collections/