diff --git a/src/flareio_cli/cli.py b/src/flareio_cli/cli.py index a3b71a1..8734f0f 100644 --- a/src/flareio_cli/cli.py +++ b/src/flareio_cli/cli.py @@ -15,6 +15,9 @@ def create_app() -> typer.Typer: no_args_is_help=True, ) + from flareio_cli.commands.export_identifier_credentials import ( + run_export_identifier_credentials, + ) from flareio_cli.commands.export_tenant_credentials import ( run_export_tenant_credentials, ) @@ -39,6 +42,10 @@ def create_app() -> typer.Typer: name="export-tenant-credentials", callable=run_export_tenant_credentials, ), + Command( + name="export-identifier-credentials", + callable=run_export_identifier_credentials, + ), ] for command in commands: app.command(name=command.name)(command.callable) diff --git a/src/flareio_cli/commands/export_identifier_credentials.py b/src/flareio_cli/commands/export_identifier_credentials.py new file mode 100644 index 0000000..5427a5e --- /dev/null +++ b/src/flareio_cli/commands/export_identifier_credentials.py @@ -0,0 +1,33 @@ +import pathlib +import typer + +from flareio.api_client import FlareApiClient +from flareio_cli.api.client import get_api_client +from flareio_cli.cursor import CursorFile +from flareio_cli.exporters.credentials import export_credentials + +import typing as t + + +def run_export_identifier_credentials( + *, + cursor_file: t.Annotated[pathlib.Path, typer.Option()], + output_file: t.Annotated[pathlib.Path, typer.Option()], + identifier_id: t.Annotated[int, 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) + if cursor.value(): + typer.echo(f"Found existing cursor. Will resume from cursor={cursor.value()}") + + # Run the export + export_credentials( + endpoint=f"/firework/v3/identifiers/{identifier_id}/feed/credentials", + api_client=api_client, + cursor=cursor, + output_file=output_file, + ) diff --git a/src/flareio_cli/commands/export_tenant_credentials.py b/src/flareio_cli/commands/export_tenant_credentials.py index ac42914..4bde068 100644 --- a/src/flareio_cli/commands/export_tenant_credentials.py +++ b/src/flareio_cli/commands/export_tenant_credentials.py @@ -1,82 +1,18 @@ import pathlib -import pydantic import typer -from datetime import timedelta -from flareio._ratelimit import _Limiter from flareio.api_client import FlareApiClient from flareio_cli.api.client import get_api_client -from flareio_cli.api.models.credentials import CredentialItem -from flareio_cli.csv import PydanticCsvWriter from flareio_cli.cursor import CursorFile -from flareio_cli.progress import export_progress +from flareio_cli.exporters.credentials import export_credentials import typing as t -class CsvItem(pydantic.BaseModel): - id: int = pydantic.Field() - identity_name: str = pydantic.Field() - hash: str = pydantic.Field() - source_id: str = pydantic.Field() - - -def _export( - *, - api_client: FlareApiClient, - csv_writer: PydanticCsvWriter[CsvItem], - cursor: CursorFile, -) -> None: - pages_limiter: _Limiter = _Limiter( - tick_interval=timedelta(seconds=1), - ) - - with export_progress( - object_name="credentials", - ) as increment_progress: - for response in api_client.scroll( - method="POST", - url="/firework/v2/me/feed/credentials", - json={ - "from": cursor.value(), - "size": 10, - "order_type": "asc", - }, - ): - pages_limiter.tick() - resp_json = response.json() - - cursor.save(resp_json["next"]) - - for item in resp_json["items"]: - credential_item = CredentialItem.model_validate(item) - - csv_writer.writerow( - row=CsvItem( - identity_name=credential_item.identity_name, - hash=credential_item.hash, - id=credential_item.id, - source_id=credential_item.source_id, - ), - ) - csv_writer.flush() - - increment_progress( - incr_completed=1, - new_cursor=cursor.value(), - ) - - def run_export_tenant_credentials( *, - cursor_file: t.Annotated[ - pathlib.Path, - typer.Option(), - ], - output_file: t.Annotated[ - pathlib.Path, - typer.Option(), - ], + 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 @@ -87,20 +23,10 @@ def run_export_tenant_credentials( if cursor.value(): typer.echo(f"Found existing cursor. Will resume from cursor={cursor.value()}") - is_output_empty: bool = ( - not output_file.exists() or not output_file.read_text().strip() - ) - # Run the export - with open(output_file, "a+", encoding="utf-8") as f_output: - dict_writer = PydanticCsvWriter( - file=f_output, - model=CsvItem, - ) - if is_output_empty: - dict_writer.writeheader() - _export( - api_client=api_client, - csv_writer=dict_writer, - cursor=cursor, - ) + export_credentials( + endpoint="/firework/v2/me/feed/credentials", + api_client=api_client, + cursor=cursor, + output_file=output_file, + ) diff --git a/src/flareio_cli/commands/export_tenant_events.py b/src/flareio_cli/commands/export_tenant_events.py index 19a348c..9607fb6 100644 --- a/src/flareio_cli/commands/export_tenant_events.py +++ b/src/flareio_cli/commands/export_tenant_events.py @@ -72,18 +72,9 @@ def _export( def run_export_tenant_events( *, - cursor_file: t.Annotated[ - pathlib.Path, - typer.Option(), - ], - from_date: t.Annotated[ - datetime.datetime | None, - typer.Option(), - ] = None, - output_file: t.Annotated[ - pathlib.Path, - typer.Option(), - ], + cursor_file: t.Annotated[pathlib.Path, typer.Option()], + from_date: t.Annotated[datetime.datetime | None, typer.Option()] = None, + output_file: t.Annotated[pathlib.Path, typer.Option()], format: t.Literal["csv"] = "csv", ) -> None: # Setup API client diff --git a/src/flareio_cli/exporters/__init__.py b/src/flareio_cli/exporters/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/flareio_cli/exporters/credentials.py b/src/flareio_cli/exporters/credentials.py new file mode 100644 index 0000000..bf07e8d --- /dev/null +++ b/src/flareio_cli/exporters/credentials.py @@ -0,0 +1,75 @@ +import pathlib +import pydantic + +from datetime import timedelta +from flareio._ratelimit import _Limiter +from flareio.api_client import FlareApiClient +from flareio_cli.api.models.credentials import CredentialItem +from flareio_cli.csv import PydanticCsvWriter +from flareio_cli.cursor import CursorFile +from flareio_cli.progress import export_progress + + +class CsvItem(pydantic.BaseModel): + id: int = pydantic.Field() + identity_name: str = pydantic.Field() + hash: str = pydantic.Field() + source_id: str = pydantic.Field() + + +def export_credentials( + *, + output_file: pathlib.Path, + endpoint: str, + api_client: FlareApiClient, + cursor: CursorFile, +) -> None: + pages_limiter: _Limiter = _Limiter( + tick_interval=timedelta(seconds=1), + ) + is_output_empty: bool = ( + not output_file.exists() or not output_file.read_text().strip() + ) + + with ( + open(output_file, "a+", encoding="utf-8") as f_output, + export_progress(object_name="credentials") as increment_progress, + ): + csv_writer = PydanticCsvWriter( + file=f_output, + model=CsvItem, + ) + if is_output_empty: + csv_writer.writeheader() + + for response in api_client.scroll( + method="POST", + url=endpoint, + json={ + "from": cursor.value(), + "size": 10, + "order_type": "asc", + }, + ): + pages_limiter.tick() + resp_json = response.json() + + cursor.save(resp_json["next"]) + + for item in resp_json["items"]: + credential_item = CredentialItem.model_validate(item) + + csv_writer.writerow( + row=CsvItem( + identity_name=credential_item.identity_name, + hash=credential_item.hash, + id=credential_item.id, + source_id=credential_item.source_id, + ), + ) + csv_writer.flush() + + increment_progress( + incr_completed=1, + new_cursor=cursor.value(), + )