From 86f69d470f77a47f80a68efe0c49f83cdbebdf88 Mon Sep 17 00:00:00 2001 From: Donkie Date: Sun, 27 Oct 2024 12:36:39 +0100 Subject: [PATCH] Added backend supported csv/json export functionality --- spoolman/api/v1/export.py | 85 +++++++++++++++++++++++++++++++++++ spoolman/api/v1/router.py | 3 +- spoolman/database/filament.py | 5 ++- spoolman/database/spool.py | 5 ++- spoolman/database/vendor.py | 5 ++- spoolman/export.py | 66 +++++++++++++++++++++++++++ 6 files changed, 165 insertions(+), 4 deletions(-) create mode 100644 spoolman/api/v1/export.py create mode 100644 spoolman/export.py diff --git a/spoolman/api/v1/export.py b/spoolman/api/v1/export.py new file mode 100644 index 0000000..f70fb3a --- /dev/null +++ b/spoolman/api/v1/export.py @@ -0,0 +1,85 @@ +"""Functions for exporting data.""" + +import io +from collections.abc import Iterable +from enum import Enum +from typing import Annotated + +from fastapi import APIRouter, Depends, Response +from sqlalchemy.ext.asyncio import AsyncSession + +from spoolman.database import filament, spool, vendor +from spoolman.database.database import get_db_session +from spoolman.database.models import Base +from spoolman.export import dump_as_csv, dump_as_json + +# ruff: noqa: D103,B008 +router = APIRouter( + prefix="/export", + tags=["export"], +) + + +class ExportFormat(Enum): + CSV = "csv" + JSON = "json" + + +@router.get( + "/spools", + name="Export spools", + description="Export the list of spools in various formats. Filament and vendor data is included.", +) +async def export_spools( + *, + db: Annotated[AsyncSession, Depends(get_db_session)], + fmt: ExportFormat, +) -> Response: + + all_spools, _ = await spool.find(db=db) + return await _export(all_spools, fmt) + + +@router.get( + "/filaments", + name="Export filaments", + description="Export the list of filaments in various formats. Vendor data is included.", +) +async def export_filaments( + *, + db: Annotated[AsyncSession, Depends(get_db_session)], + fmt: ExportFormat, +) -> Response: + all_filaments, _ = await filament.find(db=db) + return await _export(all_filaments, fmt) + + +@router.get( + "/vendors", + name="Export vendors", + description="Export the list of vendors in various formats.", +) +async def export_vendors( + *, + db: Annotated[AsyncSession, Depends(get_db_session)], + fmt: ExportFormat, +) -> Response: + all_vendors, _ = await vendor.find(db=db) + return await _export(all_vendors, fmt) + + +async def _export(objects: Iterable[Base], fmt: ExportFormat) -> Response: + """Export the objects in various formats.""" + buffer = io.StringIO() + media_type = "" + + if fmt == ExportFormat.CSV: + media_type = "text/csv" + await dump_as_csv(objects, buffer) + elif fmt == ExportFormat.JSON: + media_type = "application/json" + await dump_as_json(objects, buffer) + else: + raise ValueError(f"Unknown export format: {fmt}") + + return Response(content=buffer.getvalue(), media_type=media_type) diff --git a/spoolman/api/v1/router.py b/spoolman/api/v1/router.py index ea72938..f8466fa 100644 --- a/spoolman/api/v1/router.py +++ b/spoolman/api/v1/router.py @@ -15,7 +15,7 @@ from spoolman.database.database import backup_global_db from spoolman.exceptions import ItemNotFoundError from spoolman.ws import websocket_manager -from . import externaldb, field, filament, models, other, setting, spool, vendor +from . import export, externaldb, field, filament, models, other, setting, spool, vendor logger = logging.getLogger(__name__) @@ -111,3 +111,4 @@ app.include_router(setting.router) app.include_router(field.router) app.include_router(other.router) app.include_router(externaldb.router) +app.include_router(export.router) diff --git a/spoolman/database/filament.py b/spoolman/database/filament.py index 2613e72..7ff7c65 100644 --- a/spoolman/database/filament.py +++ b/spoolman/database/filament.py @@ -144,7 +144,10 @@ async def find( elif order == SortOrder.DESC: stmt = stmt.order_by(field.desc()) - rows = await db.execute(stmt) + rows = await db.execute( + stmt, + execution_options={"populate_existing": True}, + ) result = list(rows.unique().scalars().all()) if total_count is None: total_count = len(result) diff --git a/spoolman/database/spool.py b/spoolman/database/spool.py index b70aac3..254a85e 100644 --- a/spoolman/database/spool.py +++ b/spoolman/database/spool.py @@ -198,7 +198,10 @@ async def find( # noqa: C901, PLR0912 elif order == SortOrder.DESC: stmt = stmt.order_by(*(f.desc() for f in sorts)) - rows = await db.execute(stmt) + rows = await db.execute( + stmt, + execution_options={"populate_existing": True}, + ) result = list(rows.unique().scalars().all()) if total_count is None: total_count = len(result) diff --git a/spoolman/database/vendor.py b/spoolman/database/vendor.py index a348e48..3b90aea 100644 --- a/spoolman/database/vendor.py +++ b/spoolman/database/vendor.py @@ -83,7 +83,10 @@ async def find( elif order == SortOrder.DESC: stmt = stmt.order_by(field.desc()) - rows = await db.execute(stmt) + rows = await db.execute( + stmt, + execution_options={"populate_existing": True}, + ) result = list(rows.unique().scalars().all()) if total_count is None: total_count = len(result) diff --git a/spoolman/export.py b/spoolman/export.py new file mode 100644 index 0000000..db3bf0a --- /dev/null +++ b/spoolman/export.py @@ -0,0 +1,66 @@ +"""Functionality for exporting data in various format.""" + +import asyncio +import csv +import json +from collections.abc import Iterable +from typing import TYPE_CHECKING, Any + +from spoolman.database import models + +if TYPE_CHECKING: + from _typeshed import SupportsWrite + +banned_attrs = {"awaitable_attrs", "metadata", "registry", "spools", "filaments"} + + +async def flatten_sqlalchemy_object(obj: models.Base, parent_key: str = "", sep: str = ".") -> dict[str, Any]: + """Recursively flattens a SQLAlchemy object into a dictionary with dot-separated keys.""" + fields = {} + for attr in dir(obj): + # Check if the attribute is a column or a relationship + if not attr.startswith("_") and attr not in banned_attrs: + value = await getattr(obj.awaitable_attrs, attr) + + if attr == "extra": + # Handle extra fields + for v in value: + fields[f"{parent_key}extra.{v.key}"] = v.value + continue + + # Handle nested SQLAlchemy objects + if isinstance(value, models.Base): + nested_fields = await flatten_sqlalchemy_object(value, f"{parent_key}{attr}{sep}", sep=sep) + fields.update(nested_fields) + else: + # Use only columns and simple data types + fields[f"{parent_key}{attr}"] = value + return fields + + +async def dump_as_csv(sqlalchemy_objects: Iterable[models.Base], writer: "SupportsWrite[str]") -> None: + """Export a list of objects as CSV to a writer. Nested objects are flattened with dot-separated keys.""" + # Flatten each object and get all column names + all_flattened = await asyncio.gather(*[flatten_sqlalchemy_object(obj) for obj in sqlalchemy_objects]) + + # Collect all unique headers across flattened objects + headers = set() + for flattened_obj in all_flattened: + headers.update(flattened_obj.keys()) + + headers = sorted(headers) # Sort headers for consistent column ordering + + # Write to CSV + csv_writer = csv.DictWriter(writer, fieldnames=headers) + csv_writer.writeheader() + for flattened_obj in all_flattened: + csv_writer.writerow(flattened_obj) + + +async def dump_as_json(sqlalchemy_objects: Iterable[models.Base], writer: "SupportsWrite[str]") -> None: + """Export a list of objects as JSON to a writer. Nested objects are flattened with dot-separated keys.""" + # Flatten each object and get all column names + all_flattened = await asyncio.gather(*[flatten_sqlalchemy_object(obj) for obj in sqlalchemy_objects]) + + # Write to JSON + json.dump(all_flattened, writer, default=str)