rigorous dump cleanup, move logic to model, add tests, fixes #105

Author: monneyboi

poliloom/poliloom/cli.py

@@ -18,6 +18,8 @@
     Country,
     CurrentImportEntity,
     CurrentImportStatement,
+    DownloadAlreadyCompleteError,
+    DownloadInProgressError,
     Language,
     Location,
     Position,
@@ -140,62 +142,68 @@ def dump_download(output, force):
                 last_modified_str, "%a, %d %b %Y %H:%M:%S %Z"
             ).replace(tzinfo=timezone.utc)
 
-        # Check if we already have this dump (completed or in-progress) unless --force is used
+        # Prepare dump record (handles stale detection, force mode, etc.)
         with Session(get_engine()) as session:
-            existing_dump = (
-                session.query(WikidataDump)
-                .filter(WikidataDump.url == url)
-                .filter(WikidataDump.last_modified == last_modified)
-                .first()
-            )
+            try:
+                new_dump = WikidataDump.prepare_for_download(
+                    session, url, last_modified, force=force
+                )
+                session.commit()
 
-            if existing_dump and not force:
-                if existing_dump.downloaded_at:
+                if force:
                     click.echo(
-                        f"❌ Dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC already downloaded"
+                        f"⚠️  Forcing new download for dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC"
                     )
-                    click.echo("No new dump available. Use --force to download anyway.")
-                    raise SystemExit(1)
                 else:
                     click.echo(
-                        f"❌ Download for dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC already in progress"
-                    )
-                    click.echo(
-                        "Another download process is running. Use --force to start new download."
+                        f"📝 New dump found from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC"
                     )
-                    raise SystemExit(1)
-            elif existing_dump and force:
+
+            except DownloadAlreadyCompleteError:
                 click.echo(
-                    f"⚠️  Forcing new download for dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC (bypassing existing check)"
+                    f"❌ Dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC already downloaded"
                 )
+                click.echo("No new dump available. Use --force to download anyway.")
+                raise SystemExit(1)
 
-            # Create new dump record
-            if not existing_dump:
+            except DownloadInProgressError as e:
                 click.echo(
-                    f"📝 New dump found from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC"
+                    f"❌ Download for dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC already in progress"
                 )
-            new_dump = WikidataDump(url=url, last_modified=last_modified)
-            session.add(new_dump)
-            session.commit()
+                click.echo(
+                    f"   Started {e.hours_elapsed:.1f}h ago. Use --force to start new download."
+                )
+                raise SystemExit(1)
 
         # Download the file
         click.echo(f"⏳ Downloading Wikidata dump to {output}...")
         click.echo(
             "This is a large file (~100GB compressed) and may take several hours."
         )
 
-        StorageFactory.download_from_url(url, output)
+        try:
+            StorageFactory.download_from_url(url, output)
 
-        # Mark as downloaded
-        new_dump.downloaded_at = datetime.now(timezone.utc)
-        with Session(get_engine()) as session:
-            session.merge(new_dump)
-            session.commit()
+            # Mark as downloaded
+            with Session(get_engine()) as session:
+                new_dump.mark_downloaded(session)
+                session.commit()
+
+            click.echo(f"✅ Successfully downloaded dump to {output}")
 
-        click.echo(f"✅ Successfully downloaded dump to {output}")
+        except Exception as download_error:
+            # Clean up the dump record on failure to allow retries
+            click.echo(f"❌ Download failed: {download_error}")
+            click.echo("   Cleaning up dump record to allow retry...")
+            with Session(get_engine()) as session:
+                new_dump.cleanup_failed_download(session)
+                session.commit()
+            raise SystemExit(1)
 
+    except SystemExit:
+        raise
     except Exception as e:
-        click.echo(f"❌ Download failed: {e}")
+        click.echo(f"❌ Error: {e}")
         raise SystemExit(1)
 
 

poliloom/poliloom/models/__init__.py

@@ -17,6 +17,8 @@
 from .wikidata import (
     CurrentImportEntity,
     CurrentImportStatement,
+    DownloadAlreadyCompleteError,
+    DownloadInProgressError,
     WikidataEntity,
     WikidataEntityLabel,
     WikidataEntityMixin,
@@ -54,6 +56,8 @@
     # Wikidata
     "CurrentImportEntity",
     "CurrentImportStatement",
+    "DownloadAlreadyCompleteError",
+    "DownloadInProgressError",
     "WikidataEntity",
     "WikidataEntityLabel",
     "WikidataDump",

poliloom/poliloom/models/wikidata.py

@@ -708,11 +708,29 @@ class WikidataRelation(Base, TimestampMixin, SoftDeleteMixin, UpsertMixin):
     )
 
 
+class DownloadAlreadyCompleteError(Exception):
+    """Raised when attempting to download a dump that's already been downloaded."""
+
+    pass
+
+
+class DownloadInProgressError(Exception):
+    """Raised when another download is already in progress for this dump."""
+
+    def __init__(self, message: str, hours_elapsed: float):
+        super().__init__(message)
+        self.hours_elapsed = hours_elapsed
+
+
 class WikidataDump(Base, TimestampMixin):
     """WikidataDump entity for tracking dump download and processing stages."""
 
     __tablename__ = "wikidata_dumps"
 
+    # Default stale threshold: downloads taking longer than 24 hours are considered failed
+    # (typical download time is ~10 hours)
+    STALE_THRESHOLD_HOURS = 24
+
     id = Column(
         UUID(as_uuid=True), primary_key=True, server_default=text("gen_random_uuid()")
     )
@@ -734,6 +752,101 @@ class WikidataDump(Base, TimestampMixin):
         DateTime, nullable=True
     )  # When politicians import completed
 
+    @classmethod
+    def prepare_for_download(
+        cls,
+        session: Session,
+        url: str,
+        last_modified: datetime,
+        force: bool = False,
+    ) -> "WikidataDump":
+        """Prepare a WikidataDump record for downloading.
+
+        Handles checking for existing downloads (completed or in-progress),
+        stale download detection, and record cleanup.
+
+        Args:
+            session: Database session
+            url: URL of the dump file
+            last_modified: Last-Modified timestamp from the server
+            force: If True, bypass existing download checks
+
+        Returns:
+            WikidataDump record ready for download
+
+        Raises:
+            DownloadAlreadyCompleteError: If dump was already downloaded (and not force)
+            DownloadInProgressError: If another download is in progress (and not force/stale)
+        """
+        from datetime import timedelta, timezone
+
+        existing_dump = (
+            session.query(cls)
+            .filter(cls.url == url)
+            .filter(cls.last_modified == last_modified)
+            .first()
+        )
+
+        if existing_dump and not force:
+            if existing_dump.downloaded_at:
+                raise DownloadAlreadyCompleteError(
+                    f"Dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC "
+                    "already downloaded"
+                )
+            else:
+                # Check if the download is stale
+                created_at_utc = existing_dump.created_at.replace(tzinfo=timezone.utc)
+                age = datetime.now(timezone.utc) - created_at_utc
+                hours_elapsed = age.total_seconds() / 3600
+
+                if age > timedelta(hours=cls.STALE_THRESHOLD_HOURS):
+                    # Stale download - clean up and allow retry
+                    session.delete(existing_dump)
+                    session.flush()
+                    existing_dump = None
+                else:
+                    raise DownloadInProgressError(
+                        f"Download for dump from {last_modified.strftime('%Y-%m-%d %H:%M:%S')} UTC "
+                        "already in progress",
+                        hours_elapsed=hours_elapsed,
+                    )
+        elif existing_dump and force:
+            # Force mode - delete existing record
+            session.delete(existing_dump)
+            session.flush()
+            existing_dump = None
+
+        # Create new dump record
+        new_dump = cls(url=url, last_modified=last_modified)
+        session.add(new_dump)
+        session.flush()
+
+        return new_dump
+
+    def mark_downloaded(self, session: Session) -> None:
+        """Mark this dump as successfully downloaded.
+
+        Args:
+            session: Database session
+        """
+        from datetime import timezone
+
+        self.downloaded_at = datetime.now(timezone.utc)
+        session.merge(self)
+        session.flush()
+
+    def cleanup_failed_download(self, session: Session) -> None:
+        """Clean up this dump record after a failed download.
+
+        Removes the record to allow future retry attempts.
+
+        Args:
+            session: Database session
+        """
+        session.merge(self)
+        session.delete(self)
+        session.flush()
+
 
 class CurrentImportEntity(Base):
     """Temporary tracking table for entities seen during current import."""