Added backend supported csv/json export functionality
This commit is contained in:
85
spoolman/api/v1/export.py
Normal file
85
spoolman/api/v1/export.py
Normal 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)
|
||||||
@@ -15,7 +15,7 @@ from spoolman.database.database import backup_global_db
|
|||||||
from spoolman.exceptions import ItemNotFoundError
|
from spoolman.exceptions import ItemNotFoundError
|
||||||
from spoolman.ws import websocket_manager
|
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__)
|
logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
@@ -111,3 +111,4 @@ app.include_router(setting.router)
|
|||||||
app.include_router(field.router)
|
app.include_router(field.router)
|
||||||
app.include_router(other.router)
|
app.include_router(other.router)
|
||||||
app.include_router(externaldb.router)
|
app.include_router(externaldb.router)
|
||||||
|
app.include_router(export.router)
|
||||||
|
|||||||
@@ -144,7 +144,10 @@ async def find(
|
|||||||
elif order == SortOrder.DESC:
|
elif order == SortOrder.DESC:
|
||||||
stmt = stmt.order_by(field.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())
|
result = list(rows.unique().scalars().all())
|
||||||
if total_count is None:
|
if total_count is None:
|
||||||
total_count = len(result)
|
total_count = len(result)
|
||||||
|
|||||||
@@ -198,7 +198,10 @@ async def find( # noqa: C901, PLR0912
|
|||||||
elif order == SortOrder.DESC:
|
elif order == SortOrder.DESC:
|
||||||
stmt = stmt.order_by(*(f.desc() for f in sorts))
|
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())
|
result = list(rows.unique().scalars().all())
|
||||||
if total_count is None:
|
if total_count is None:
|
||||||
total_count = len(result)
|
total_count = len(result)
|
||||||
|
|||||||
@@ -83,7 +83,10 @@ async def find(
|
|||||||
elif order == SortOrder.DESC:
|
elif order == SortOrder.DESC:
|
||||||
stmt = stmt.order_by(field.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())
|
result = list(rows.unique().scalars().all())
|
||||||
if total_count is None:
|
if total_count is None:
|
||||||
total_count = len(result)
|
total_count = len(result)
|
||||||
|
|||||||
66
spoolman/export.py
Normal file
66
spoolman/export.py
Normal 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)
|
||||||
Reference in New Issue
Block a user