Skip to content

Commit 3cf0572

Browse files
RijulTPRijulTP
committed
Added metadata to export results, and improved the output
1 parent 81f7381 commit 3cf0572

2 files changed

Lines changed: 180 additions & 102 deletions

File tree

‎export.go‎

Lines changed: 133 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -14,24 +14,34 @@ import (
1414
)
1515

1616
type ExportOptions struct {
17-
Format string
18-
OutDir string
19-
FromDate string
20-
DryRun bool
21-
Overwrite bool
17+
Format string
18+
OutDir string
19+
FromDate string
20+
FullExport bool
21+
DryRun bool
22+
Overwrite bool
2223
}
2324

25+
2426
type ExportResult struct {
25-
JobID int64 `json:"job_id"`
26-
Result string `json:"result"`
27-
CreatedAt time.Time `json:"created_at"`
27+
JobID int64
28+
Response []byte
29+
GroupName *string
30+
CreatedAt time.Time
2831
}
2932

33+
34+
3035
func ExportResults(
3136
ctx context.Context,
3237
dbPool *pgxpool.Pool,
3338
opts ExportOptions,
3439
) (int, error) {
40+
if opts.FullExport && opts.FromDate != "" {
41+
return 0, fmt.Errorf("--full-export cannot be used with --from-date")
42+
}
43+
44+
start := time.Now()
3545

3646
// Ensure directory exists
3747
if err := os.MkdirAll(opts.OutDir, 0755); err != nil {
@@ -51,100 +61,158 @@ func ExportResults(
5161
existing = scanExistingExports(opts.OutDir, opts.Format)
5262
}
5363

54-
// Determine fromTime
55-
var fromTime time.Time
56-
if opts.FromDate != "" {
57-
t, err := time.Parse("2006-01-02", opts.FromDate)
58-
if err != nil {
59-
return 0, fmt.Errorf("invalid --from-date format (expected YYYY-MM-DD)")
60-
}
61-
fromTime = t
64+
var (
65+
query string
66+
args []any
67+
)
68+
69+
if opts.FullExport {
70+
fmt.Println("Mode: full-export (no date filter)")
71+
query = `
72+
SELECT
73+
r.job_id,
74+
r.response,
75+
r.created_at,
76+
g.group_name
77+
FROM dprompts_results r
78+
LEFT JOIN dprompt_groups g
79+
ON r.group_id = g.id
80+
ORDER BY r.created_at ASC
81+
`
6282
} else {
63-
fromTime = time.Now().Add(-24 * time.Hour)
83+
var fromTime time.Time
84+
85+
if opts.FromDate != "" {
86+
fmt.Println("Mode: from-date", opts.FromDate)
87+
t, err := time.Parse("2006-01-02", opts.FromDate)
88+
if err != nil {
89+
return 0, fmt.Errorf("invalid --from-date format (expected YYYY-MM-DD)")
90+
}
91+
fromTime = t
92+
} else {
93+
fmt.Println("Mode: last-24-hours")
94+
fromTime = time.Now().Add(-24 * time.Hour)
95+
}
96+
97+
query = `
98+
SELECT
99+
r.job_id,
100+
r.response,
101+
r.created_at,
102+
g.group_name
103+
FROM dprompts_results r
104+
LEFT JOIN dprompt_groups g
105+
ON r.group_id = g.id
106+
WHERE r.created_at >= $1
107+
ORDER BY r.created_at ASC
108+
`
109+
args = []any{fromTime}
110+
}
111+
112+
// ---- count total matching rows ----
113+
countQuery := fmt.Sprintf("SELECT COUNT(*) FROM (%s) q", query)
114+
var totalMatched int
115+
if err := dbPool.QueryRow(ctx, countQuery, args...).Scan(&totalMatched); err != nil {
116+
return 0, err
64117
}
65118

66-
query := `
67-
SELECT job_id, response, created_at
68-
FROM dprompts_results
69-
WHERE created_at >= $1
70-
ORDER BY created_at ASC
71-
`
72-
fmt.Println(fromTime)
73-
rows, err := dbPool.Query(ctx, query, fromTime)
119+
skipped := len(existing)
120+
toExport := totalMatched - skipped
121+
if toExport < 0 {
122+
toExport = 0
123+
}
124+
125+
fmt.Println("Matched jobs in DB:", totalMatched)
126+
fmt.Println("Already exported (skipped):", skipped)
127+
fmt.Println("Jobs to export:", toExport)
128+
fmt.Println()
129+
130+
rows, err := dbPool.Query(ctx, query, args...)
74131
if err != nil {
75132
return 0, err
76133
}
77134
defer rows.Close()
78135

79-
count := 0
136+
exported := 0
80137

81138
for rows.Next() {
82139
var r ExportResult
83140
if err := rows.Scan(
84141
&r.JobID,
85-
&r.Result,
142+
&r.Response,
86143
&r.CreatedAt,
144+
&r.GroupName,
87145
); err != nil {
88-
return count, err
146+
return exported, err
89147
}
90148

91-
// Skip already exported jobs
92149
if _, ok := existing[r.JobID]; ok {
93150
continue
94151
}
95152

153+
exported++
154+
155+
if exported == 1 || exported%50 == 0 || exported == toExport {
156+
fmt.Printf(
157+
"Exporting %d/%d (job_id=%d)\n",
158+
exported,
159+
toExport,
160+
r.JobID,
161+
)
162+
}
163+
96164
if !opts.DryRun {
97165
if err := writeExportFile(opts, r); err != nil {
98-
return count, err
166+
return exported, err
99167
}
100168
}
101-
102-
count++
103169
}
104170

105-
return count, rows.Err()
171+
duration := time.Since(start)
172+
173+
fmt.Println()
174+
fmt.Println("Export completed")
175+
fmt.Println("Exported:", exported)
176+
fmt.Println("Skipped:", skipped)
177+
fmt.Println("Duration:", duration.Round(time.Millisecond))
178+
179+
return exported, rows.Err()
106180
}
107181

182+
108183
func writeExportFile(opts ExportOptions, r ExportResult) error {
109-
filename := fmt.Sprintf("%d.%s", r.JobID, opts.Format)
184+
filename := fmt.Sprintf("%d.json", r.JobID)
110185
path := filepath.Join(opts.OutDir, filename)
111186

112-
switch opts.Format {
113-
case "json":
114-
var parsed any
115-
if err := json.Unmarshal([]byte(r.Result), &parsed); err != nil {
116-
return fmt.Errorf("invalid JSON in result for job %d: %w", r.JobID, err)
117-
}
118-
119-
normalized := normalizeJSON(parsed)
120-
121-
out := map[string]any{
122-
"job_id": r.JobID,
123-
"created_at": r.CreatedAt,
124-
"result": normalized,
125-
}
126-
127-
data, err := json.MarshalIndent(out, "", " ")
128-
if err != nil {
129-
return err
130-
}
131-
132-
return os.WriteFile(path, data, 0644)
187+
var resultValue any
133188

134-
case "txt":
135-
content := fmt.Sprintf(
136-
"Job ID: %d\nCreated At: %s\n\nResult:\n%s\n",
137-
r.JobID,
138-
r.CreatedAt.Format(time.RFC3339),
139-
r.Result,
140-
)
141-
return os.WriteFile(path, []byte(content), 0644)
189+
// Try JSON first
190+
var parsed any
191+
if err := json.Unmarshal(r.Response, &parsed); err == nil {
192+
resultValue = normalizeJSON(parsed)
193+
} else {
194+
// Not JSON → keep as text
195+
resultValue = string(r.Response)
196+
}
142197

143-
default:
144-
return fmt.Errorf("unsupported export format: %s", opts.Format)
198+
out := map[string]any{
199+
"job_id": r.JobID,
200+
"created_at": r.CreatedAt,
201+
"result": resultValue,
202+
"metadata": map[string]any{
203+
"group_name": r.GroupName,
204+
},
145205
}
206+
207+
data, err := json.MarshalIndent(out, "", " ")
208+
if err != nil {
209+
return err
210+
}
211+
212+
return os.WriteFile(path, data, 0644)
146213
}
147214

215+
148216
func scanExistingExports(dir, format string) map[int64]struct{} {
149217
result := make(map[int64]struct{})
150218

‎main.go‎

Lines changed: 47 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -311,49 +311,59 @@ func main() {
311311

312312

313313

314-
// ---- Export subcommand ----
315-
var (
316-
exportFormat string
317-
exportOutDir string
318-
exportFromDate string
319-
exportDryRun bool
320-
exportOverwrite bool
321-
)
322-
323-
exportCmd := &cobra.Command{
324-
Use: "export",
325-
Short: "Export dprompts results to files",
326-
Run: func(cmd *cobra.Command, args []string) {
327-
ctx := context.Background()
328-
329-
dbPool, err := NewDBPool(ctx, configPath)
330-
if err != nil {
331-
log.Fatal().Err(err).Msg("Failed to connect to database")
332-
}
333-
defer dbPool.Close()
334-
335-
count, err := ExportResults(ctx, dbPool, ExportOptions{
336-
Format: exportFormat,
337-
OutDir: exportOutDir,
338-
FromDate: exportFromDate,
339-
DryRun: exportDryRun,
340-
Overwrite: exportOverwrite,
341-
})
342-
if err != nil {
343-
log.Fatal().Err(err).Msg("Export failed")
344-
}
345-
346-
log.Info().
347-
Int("exported_count", count).
348-
Msg("Export completed successfully")
349-
},
350-
}
314+
// ---- Export subcommand ----
315+
var (
316+
exportFormat string
317+
exportOutDir string
318+
exportFromDate string
319+
exportDryRun bool
320+
exportOverwrite bool
321+
exportFullExport bool
322+
)
323+
324+
exportCmd := &cobra.Command{
325+
Use: "export",
326+
Short: "Export dprompts results to files",
327+
Run: func(cmd *cobra.Command, args []string) {
328+
ctx := context.Background()
329+
330+
dbPool, err := NewDBPool(ctx, configPath)
331+
if err != nil {
332+
log.Fatal().Err(err).Msg("Failed to connect to database")
333+
}
334+
defer dbPool.Close()
335+
336+
count, err := ExportResults(ctx, dbPool, ExportOptions{
337+
Format: exportFormat,
338+
OutDir: exportOutDir,
339+
FromDate: exportFromDate,
340+
FullExport: exportFullExport,
341+
DryRun: exportDryRun,
342+
Overwrite: exportOverwrite,
343+
})
344+
345+
if err != nil {
346+
log.Fatal().Err(err).Msg("Export failed")
347+
}
348+
349+
log.Info().
350+
Int("exported_count", count).
351+
Msg("Export completed successfully")
352+
},
353+
}
351354

352355
exportCmd.Flags().StringVar(&exportFormat, "format", "json", "Export format: json | txt")
353356
exportCmd.Flags().StringVar(&exportOutDir, "out", "./dprompts_exports", "Output directory")
354357
exportCmd.Flags().StringVar(&exportFromDate, "from-date", "", "Export results created after this date (YYYY-MM-DD)")
355358
exportCmd.Flags().BoolVar(&exportDryRun, "dry-run", false, "Show what would be exported without writing files")
356359
exportCmd.Flags().BoolVar(&exportOverwrite, "overwrite", false, "Overwrite existing exported files")
360+
exportCmd.Flags().BoolVar(
361+
&exportFullExport,
362+
"full-export",
363+
false,
364+
"Export all results (ignores --from-date)",
365+
)
366+
357367
// Add subcommands
358368
rootCmd.AddCommand(clientCmd, workerCmd, viewCmd, deleteGroupCmd, queueCmd,exportCmd)
359369

0 commit comments

Comments
 (0)