Add force rerun recovery and polish CLI startup
This commit is contained in:
rayd1o
@@ -1,5 +1,6 @@
|
||||
"""Base collector class for all data sources"""
|
||||
|
||||
import asyncio
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, List, Any, Optional
|
||||
from datetime import UTC, datetime
|
||||
@@ -166,6 +167,59 @@ class BaseCollector(ABC):
|
||||
await db.commit()
|
||||
return snapshot.id
|
||||
|
||||
async def _rollback_incomplete_run(
|
||||
self,
|
||||
db: AsyncSession,
|
||||
*,
|
||||
task_id: int,
|
||||
snapshot_id: Optional[int],
|
||||
reason: str,
|
||||
) -> None:
|
||||
from app.models.collected_data import CollectedData
|
||||
from app.models.data_snapshot import DataSnapshot
|
||||
|
||||
await db.execute(CollectedData.__table__.delete().where(CollectedData.task_id == task_id))
|
||||
|
||||
parent_snapshot_id: Optional[int] = None
|
||||
if snapshot_id is not None:
|
||||
snapshot = await db.get(DataSnapshot, snapshot_id)
|
||||
if snapshot:
|
||||
parent_snapshot_id = snapshot.parent_snapshot_id
|
||||
snapshot.status = "cancelled"
|
||||
snapshot.is_current = False
|
||||
snapshot.completed_at = datetime.now(UTC)
|
||||
summary = dict(snapshot.summary or {})
|
||||
summary["rollback"] = True
|
||||
summary["rollback_reason"] = reason
|
||||
snapshot.summary = summary
|
||||
|
||||
await db.execute(
|
||||
text(
|
||||
"""
|
||||
UPDATE collected_data
|
||||
SET is_current = FALSE
|
||||
WHERE source = :source
|
||||
"""
|
||||
),
|
||||
{"source": self.name},
|
||||
)
|
||||
|
||||
if parent_snapshot_id is not None:
|
||||
parent_snapshot = await db.get(DataSnapshot, parent_snapshot_id)
|
||||
if parent_snapshot:
|
||||
parent_snapshot.is_current = True
|
||||
|
||||
await db.execute(
|
||||
text(
|
||||
"""
|
||||
UPDATE collected_data
|
||||
SET is_current = TRUE
|
||||
WHERE snapshot_id = :snapshot_id
|
||||
"""
|
||||
),
|
||||
{"snapshot_id": parent_snapshot_id},
|
||||
)
|
||||
|
||||
async def run(self, db: AsyncSession) -> Dict[str, Any]:
|
||||
"""Full pipeline: fetch -> transform -> save"""
|
||||
from app.services.collectors.registry import collector_registry
|
||||
@@ -227,6 +281,21 @@ class BaseCollector(ABC):
|
||||
"records_processed": records_count,
|
||||
"execution_time_seconds": (datetime.now(UTC) - start_time).total_seconds(),
|
||||
}
|
||||
except asyncio.CancelledError:
|
||||
task.status = "cancelled"
|
||||
task.phase = "cancelled"
|
||||
task.error_message = "Collection cancelled by operator and rolled back"
|
||||
task.completed_at = datetime.now(UTC)
|
||||
if snapshot_id is not None:
|
||||
await self._rollback_incomplete_run(
|
||||
db,
|
||||
task_id=task_id,
|
||||
snapshot_id=snapshot_id,
|
||||
reason="cancelled_by_operator",
|
||||
)
|
||||
await db.commit()
|
||||
await self._publish_task_update(force=True)
|
||||
raise
|
||||
except Exception as e:
|
||||
task.status = "failed"
|
||||
task.phase = "failed"
|
||||
@@ -276,20 +345,34 @@ class BaseCollector(ABC):
|
||||
updated_count = 0
|
||||
unchanged_count = 0
|
||||
seen_entity_keys: set[str] = set()
|
||||
previous_current_keys: set[str] = set()
|
||||
progress_commit_interval = 1000
|
||||
|
||||
previous_current_result = await db.execute(
|
||||
select(CollectedData.entity_key).where(
|
||||
select(CollectedData)
|
||||
.where(
|
||||
CollectedData.source == self.name,
|
||||
CollectedData.is_current == True,
|
||||
)
|
||||
.order_by(CollectedData.entity_key.asc(), CollectedData.collected_at.desc().nullslast(), CollectedData.id.desc())
|
||||
)
|
||||
previous_current_keys = {row[0] for row in previous_current_result.fetchall() if row[0]}
|
||||
previous_current_records = previous_current_result.scalars().all()
|
||||
previous_current_keys = {record.entity_key for record in previous_current_records if record.entity_key}
|
||||
previous_current_map: dict[str, CollectedData] = {}
|
||||
stale_previous_records: list[CollectedData] = []
|
||||
|
||||
for existing_record in previous_current_records:
|
||||
entity_key = existing_record.entity_key
|
||||
if not entity_key:
|
||||
continue
|
||||
if entity_key not in previous_current_map:
|
||||
previous_current_map[entity_key] = existing_record
|
||||
continue
|
||||
stale_previous_records.append(existing_record)
|
||||
|
||||
for stale_record in stale_previous_records:
|
||||
stale_record.is_current = False
|
||||
|
||||
for i, item in enumerate(data):
|
||||
print(
|
||||
f"DEBUG: Saving item {i}: name={item.get('name')}, metadata={item.get('metadata', 'NOT FOUND')}"
|
||||
)
|
||||
raw_metadata = item.get("metadata", {})
|
||||
extra_data = build_dynamic_metadata(
|
||||
raw_metadata,
|
||||
@@ -318,20 +401,9 @@ class BaseCollector(ABC):
|
||||
previous_record = None
|
||||
|
||||
if entity_key and entity_key not in seen_entity_keys:
|
||||
result = await db.execute(
|
||||
select(CollectedData)
|
||||
.where(
|
||||
CollectedData.source == self.name,
|
||||
CollectedData.entity_key == entity_key,
|
||||
CollectedData.is_current == True,
|
||||
)
|
||||
.order_by(CollectedData.collected_at.desc().nullslast(), CollectedData.id.desc())
|
||||
)
|
||||
previous_records = result.scalars().all()
|
||||
if previous_records:
|
||||
previous_record = previous_records[0]
|
||||
for old_record in previous_records:
|
||||
old_record.is_current = False
|
||||
previous_record = previous_current_map.get(entity_key)
|
||||
if previous_record is not None:
|
||||
previous_record.is_current = False
|
||||
|
||||
record = CollectedData(
|
||||
snapshot_id=snapshot_id,
|
||||
@@ -375,7 +447,7 @@ class BaseCollector(ABC):
|
||||
seen_entity_keys.add(entity_key)
|
||||
records_added += 1
|
||||
|
||||
if i % 100 == 0:
|
||||
if (i + 1) % progress_commit_interval == 0:
|
||||
await self.update_progress(i + 1, commit=True)
|
||||
|
||||
if snapshot_id is not None:
|
||||
|
||||
Reference in New Issue
Block a user