Added backend supported csv/json export functionality

This commit is contained in:
Donkie
2024-10-27 12:36:39 +01:00
parent a95ddcedd1
commit 86f69d470f
6 changed files with 165 additions and 4 deletions

85
spoolman/api/v1/export.py Normal file
View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

View File

@@ -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)

66
spoolman/export.py Normal file
View File

@@ -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)