remove cursor files

Author: aviau

src/flareio_cli/commands/export_identifier_credentials.py

@@ -7,7 +7,7 @@
 from flareio.api_client import FlareApiClient
 
 from flareio_cli.api.client import get_api_client
-from flareio_cli.cursor import CursorFile
+from flareio_cli.cursor import Cursor
 from flareio_cli.exporters.credentials import export_credentials
 
 
@@ -26,7 +26,7 @@ def export_identifier_credentials(
     api_client: FlareApiClient = get_api_client()
 
     # Load existing cursor if it exists.
-    cursor: CursorFile = CursorFile(path=cursor_file)
+    cursor: Cursor = Cursor.from_csv(path=output_file)
     if cursor.value():
         typer.echo(f"Found existing cursor. Will resume from cursor={cursor.value()}")
 

src/flareio_cli/commands/export_tenant_credentials.py

@@ -7,7 +7,7 @@
 from flareio.api_client import FlareApiClient
 
 from flareio_cli.api.client import get_api_client
-from flareio_cli.cursor import CursorFile
+from flareio_cli.cursor import Cursor
 from flareio_cli.exporters.credentials import export_credentials
 
 
@@ -17,15 +17,14 @@
 @app.command()
 def export_tenant_credentials(
     *,
-    cursor_file: t.Annotated[pathlib.Path, typer.Option()],
     output_file: t.Annotated[pathlib.Path, typer.Option()],
     format: t.Literal["csv"] = "csv",
 ) -> None:
     # Setup API client
     api_client: FlareApiClient = get_api_client()
 
     # Load existing cursor if it exists.
-    cursor: CursorFile = CursorFile(path=cursor_file)
+    cursor: Cursor = Cursor.from_csv(path=output_file)
     if cursor.value():
         typer.echo(f"Found existing cursor. Will resume from cursor={cursor.value()}")
 

src/flareio_cli/commands/export_tenant_events.py

@@ -9,7 +9,7 @@
 from flareio.models import ScrollEventsResult
 
 from flareio_cli.api.client import get_api_client
-from flareio_cli.cursor import CursorFile
+from flareio_cli.cursor import Cursor
 from flareio_cli.exporters.events import export_events
 
 
@@ -28,7 +28,7 @@ def export_tenant_events(
     api_client: FlareApiClient = get_api_client()
 
     # Load existing cursor if it exists.
-    cursor: CursorFile = CursorFile(path=cursor_file)
+    cursor: Cursor = Cursor.from_csv(path=cursor_file)
     if cursor.value():
         typer.echo(f"Found existing cursor. Will resume from cursor={cursor.value()}")
 

src/flareio_cli/csv.py

@@ -4,6 +4,8 @@
 
 import typing as t
 
+from flareio_cli.cursor import Cursor
+
 
 CsvModel = t.TypeVar("CsvModel", bound=BaseModel)
 
@@ -15,21 +17,29 @@ def __init__(
         file: t.TextIO,
         model: t.Type[CsvModel],
     ) -> None:
+        fieldnames = [
+            field_info.serialization_alias or field_name
+            for field_name, field_info in model.model_fields.items()
+        ]
+        fieldnames.append("next")
         self.dict_writer = csv.DictWriter(
             file,
-            fieldnames=[
-                field_info.serialization_alias or field_name
-                for field_name, field_info in model.model_fields.items()
-            ],
+            fieldnames=fieldnames,
             lineterminator="\n",
         )
         self.file = file
 
     def writeheader(self) -> None:
         self.dict_writer.writeheader()
 
-    def writerow(self, row: CsvModel) -> None:
-        dumped = row.model_dump(by_alias=True)
+    def writerow(
+        self,
+        *,
+        row: CsvModel,
+        cursor: Cursor,
+    ) -> None:
+        dumped: dict = row.model_dump(by_alias=True)
+        dumped["next"] = cursor.value()
         self.dict_writer.writerow(dumped)
 
     def flush(self) -> None:

src/flareio_cli/cursor.py

@@ -1,26 +1,38 @@
+import csv
 import pathlib
 
 
-class CursorFile:
+class Cursor:
     def __init__(
         self,
         *,
-        path: pathlib.Path,
+        value: str | None,
     ) -> None:
-        self._path = path
-        self._cached_value: str | None = None
+        self._value: str | None = value
 
     def value(self) -> str | None:
-        if self._cached_value:
-            return self._cached_value
-        if not self._path.exists():
-            return None
-        cursor = self._path.read_text().strip() or None
-        self._cached_value = cursor
-        return cursor
+        return self._value
 
     def save(self, value: str | None) -> None:
         if not value:
             return
-        self._cached_value = value
-        self._path.write_text(value)
+        self._value = value
+
+    @classmethod
+    def from_csv(
+        cls,
+        *,
+        path: pathlib.Path,
+    ) -> "Cursor":
+        cursor = cls(value=None)
+
+        if not path.exists():
+            return cursor
+
+        with open(path, "r", encoding="utf-8") as f:
+            csv_reader = csv.DictReader(f)
+            for row in csv_reader:
+                next_ = row.get("next")
+                cursor.save(next_)
+
+        return cursor

src/flareio_cli/exporters/base.py

@@ -7,11 +7,14 @@
 import typing as t
 
 from flareio_cli.csv import PydanticCsvWriter
-from flareio_cli.cursor import CursorFile
+from flareio_cli.cursor import Cursor
 from flareio_cli.progress import export_progress
 
 
-ExportItem = t.TypeVar("ExportItem", bound=pydantic.BaseModel)
+ExportItem = t.TypeVar(
+    "ExportItem",
+    bound=pydantic.BaseModel,
+)
 
 
 @dataclasses.dataclass(frozen=True)
@@ -24,7 +27,7 @@ def export_to_csv(
     *,
     output_file: pathlib.Path,
     pages: t.Iterator[ExportPage[ExportItem]],
-    cursor: CursorFile,
+    cursor: Cursor,
     object_name: str = "items",
     item_model: type[ExportItem],
 ) -> None:
@@ -45,7 +48,7 @@ def _export_to_csv(
     *,
     output_file: pathlib.Path,
     pages: t.Iterator[ExportPage[ExportItem]],
-    cursor: CursorFile,
+    cursor: Cursor,
     object_name: str = "items",
     item_model: type[ExportItem],
 ) -> None:
@@ -65,12 +68,15 @@ def _export_to_csv(
             writer.writeheader()
 
         for page in pages:
+            cursor.save(page.next)
+
             for row in page.items:
-                writer.writerow(row)
+                writer.writerow(
+                    row=row,
+                    cursor=cursor,
+                )
                 writer.flush()
 
-            cursor.save(page.next)
-
             progress_manager.update_progress(
                 incr_completed=len(page.items),
                 new_cursor=cursor.value(),

src/flareio_cli/exporters/credentials.py

@@ -10,7 +10,7 @@
 from flareio.ratelimit import Limiter
 
 from flareio_cli.api.models.credentials import CredentialItem
-from flareio_cli.cursor import CursorFile
+from flareio_cli.cursor import Cursor
 from flareio_cli.exporters.base import ExportPage
 from flareio_cli.exporters.base import export_to_csv
 
@@ -60,7 +60,7 @@ def export_credentials(
     *,
     output_file: pathlib.Path,
     resp_iterator: t.Iterator[requests.Response],
-    cursor: CursorFile,
+    cursor: Cursor,
 ) -> None:
     export_to_csv(
         output_file=output_file,

src/flareio_cli/exporters/events.py

@@ -8,7 +8,7 @@
 from flareio.models import ScrollEventsResult
 
 from flareio_cli.api.models.events import EventItem
-from flareio_cli.cursor import CursorFile
+from flareio_cli.cursor import Cursor
 from flareio_cli.exporters.base import ExportPage
 from flareio_cli.exporters.base import export_to_csv
 
@@ -45,7 +45,7 @@ def export_events(
     *,
     output_file: pathlib.Path,
     events_iterator: t.Iterator[ScrollEventsResult],
-    cursor: CursorFile,
+    cursor: Cursor,
 ) -> None:
     export_to_csv(
         output_file=output_file,

src/flareio_cli/test_csv.py

@@ -3,6 +3,7 @@
 import pydantic
 
 from flareio_cli.csv import PydanticCsvWriter
+from flareio_cli.cursor import Cursor
 
 
 def test_pydantic_csv_writer() -> None:
@@ -18,13 +19,19 @@ class TestModel(pydantic.BaseModel):
     )
 
     writer.writeheader()
-    writer.writerow(TestModel(first_field=1, second_field="Alice"))
-    writer.writerow(TestModel(first_field=2, second_field="Bob"))
+    writer.writerow(
+        row=TestModel(first_field=1, second_field="Alice"),
+        cursor=Cursor(value=None),
+    )
+    writer.writerow(
+        row=TestModel(first_field=2, second_field="Bob"),
+        cursor=Cursor(value=None),
+    )
 
     assert (
         csv_file_str.getvalue()
-        == """first_field,second.field
-1,Alice
-2,Bob
+        == """first_field,second.field,next
+1,Alice,
+2,Bob,
 """
     )

src/flareio_cli/test_cursor.py

@@ -1,18 +1,49 @@
-import pathlib
 import tempfile
 
-from flareio_cli.cursor import CursorFile
+from pathlib import Path
 
+from flareio_cli.cursor import Cursor
 
-def test_cursor_file() -> None:
-    with tempfile.NamedTemporaryFile() as f:
-        cursor_path = pathlib.Path(f.name)
-        cursor = CursorFile(path=cursor_path)
 
-        assert cursor.value() is None
+def test_cursor() -> None:
+    cursor = Cursor(value=None)
+
+    assert cursor.value() is None
+
+    cursor.save("123")
+    assert cursor.value() == "123"
+
+    cursor.save(None)
+    assert cursor.value() == "123"
 
-        cursor.save("123")
-        assert cursor.value() == "123"
 
-        cursor.save(None)
-        assert cursor.value() == "123"
+def test_cursor_from_csv() -> None:
+    with tempfile.NamedTemporaryFile() as f:
+        f.write(
+            """name,next
+bob,first
+bob,second
+bob,
+bob,third
+bob,
+""".encode()
+        )
+        f.flush()
+        cursor = Cursor.from_csv(
+            path=Path(f.name),
+        )
+        assert cursor.value() == "third"
+
+
+def test_cursor_from_csv_none() -> None:
+    with tempfile.NamedTemporaryFile() as f:
+        f.write(
+            """name,next
+bob,
+""".encode()
+        )
+        f.flush()
+        cursor = Cursor.from_csv(
+            path=Path(f.name),
+        )
+        assert cursor.value() is None