@qingflow-tech/qingflow-cli 1.1.59 → 1.1.60

package/src/qingflow_mcp/tools/record_tools.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import csv
 import html
 import mimetypes
 import json
@@ -34,6 +33,7 @@ from ..list_type_labels import (
 )
 from .base import ToolBase
 from .directory_tools import _directory_has_more, _directory_items
+from .csv_export import CsvShardWriter
 
 
 DEFAULT_QUERY_PAGE_SIZE = 50
@@ -2284,57 +2284,22 @@ class RecordTools(ToolBase):
                 }
 
             run_dir = _record_access_run_dir()
-            run_dir.mkdir(parents=True, exist_ok=True)
             header = ["record_id"] + [str(field["column_name"]) for field in fields_payload]
-            files: list[JSONObject] = []
-            shard_part = 0
-            shard_rows = 0
-            row_count = 0
-            file_handle: Any | None = None
-            writer: Any | None = None
-
-            def close_shard() -> None:
-                nonlocal file_handle, writer, shard_rows
-                if file_handle is None:
-                    return
-                path = Path(file_handle.name)
-                file_handle.close()
-                files.append(
-                    {
-                        "local_path": str(path),
-                        "row_count": shard_rows,
-                        "part": len(files) + 1,
-                    }
-                )
-                file_handle = None
-                writer = None
-                shard_rows = 0
-
-            def ensure_writer() -> Any:
-                nonlocal file_handle, writer, shard_part
-                if writer is not None:
-                    return writer
-                shard_part += 1
-                path = run_dir / f"records-{shard_part:04d}.csv"
-                file_handle = path.open("w", newline="", encoding="utf-8-sig")
-                writer = csv.writer(file_handle)
-                writer.writerow(header)
-                return writer
+            shard_writer = CsvShardWriter(
+                output_dir=run_dir,
+                file_prefix="records",
+                header=header,
+                shard_rows=DEFAULT_RECORD_ACCESS_SHARD_ROWS,
+            )
 
             def write_record(apply_id: int | None, answer_list: list[JSONValue]) -> None:
-                nonlocal row_count, shard_rows
                 if apply_id is None:
                     return
-                if shard_rows >= DEFAULT_RECORD_ACCESS_SHARD_ROWS:
-                    close_shard()
-                csv_writer = ensure_writer()
                 values: list[str] = [_public_record_id_text(apply_id) or ""]
                 for field in selected_fields:
                     answer = _find_answer_for_field(answer_list, field)
                     values.append(_record_access_csv_cell(_normalize_answer_field_value_for_output(answer, field)))
-                csv_writer.writerow(values)
-                row_count += 1
-                shard_rows += 1
+                shard_writer.write(values)
 
             current_page = 1
             has_more = False
@@ -2374,7 +2339,7 @@ class RecordTools(ToolBase):
                         break
                     page = fetch_page(current_page, BACKEND_RECORD_ACCESS_PAGE_SIZE)
             finally:
-                close_shard()
+                shard_writer.close_shard()
 
             if stopped_reason == "TIME_BUDGET_EXCEEDED":
                 warnings.append(_record_access_time_budget_warning(reported_total=reported_total, fetched_pages=fetched_pages))
@@ -2399,11 +2364,11 @@ class RecordTools(ToolBase):
                 "view_id": view_route.view_id,
                 "format": "csv",
                 "local_dir": str(run_dir),
-                "row_count": row_count,
+                "row_count": shard_writer.row_count,
                 "complete": complete,
                 "truncated": not complete,
                 "safe_for_final_conclusion": safe_for_final_conclusion,
-                "files": files,
+                "files": shard_writer.files,
                 "fields": fields_payload,
                 "warnings": warnings,
                 "verification": verification_payload,