Skip to content

Commit 1e91f45

Browse files
RijulTPRijulTP
committed
Added commands
1 parent 3e52a85 commit 1e91f45

6 files changed

Lines changed: 358 additions & 50 deletions

File tree

‎.dprompts.toml‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,7 @@ port = "5432"
1010
api-endpoint = "http://localhost:11434/api/chat"
1111
model = "gemma2:2b"
1212
temperature = 0.7
13-
topP = 0.9
13+
topP = 0.9
14+
15+
[worker]
16+
max_workers = 1

‎client.go‎

Lines changed: 142 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
package main
22

33
import (
4+
"bufio"
45
"context"
56
"encoding/json"
7+
"errors"
8+
"fmt"
9+
"io"
610
"os"
711

812
"github.com/jackc/pgx/v5"
@@ -73,48 +77,164 @@ func enqueueBulkJobsFromFile(ctx context.Context, riverClient *river.Client[pgx.
7377
}
7478
defer file.Close()
7579

76-
var jobs []BulkJob
77-
if err := json.NewDecoder(file).Decode(&jobs); err != nil {
78-
return err
79-
}
80+
decoder := json.NewDecoder(bufio.NewReader(file))
8081

81-
tx, err := dbPool.Begin(ctx)
82+
// Peek first non-whitespace token
83+
tok, err := nextNonSpaceToken(decoder)
8284
if err != nil {
83-
return err
85+
return fmt.Errorf("cannot read file: %w", err)
8486
}
85-
defer tx.Rollback(ctx)
8687

87-
var jobsToInsert []river.InsertManyParams
88+
// NDJSON format (each line = JSON object)
89+
if tok != json.Delim('[') {
90+
return processNDJSON(ctx, decoder, riverClient, dbPool)
91+
}
92+
93+
return processJSONArray(ctx, decoder, riverClient, dbPool)
94+
}
95+
96+
// ------ JSON ARRAY VERSION ------
97+
func processJSONArray(ctx context.Context, decoder *json.Decoder, riverClient *river.Client[pgx.Tx], dbPool *pgxpool.Pool) error {
98+
const batchSize = 500
99+
100+
batch := make([]river.InsertManyParams, 0, batchSize)
101+
total := 0
102+
count := 0
103+
104+
for decoder.More() {
105+
var job BulkJob
106+
if err := decoder.Decode(&job); err != nil {
107+
return fmt.Errorf("decode error at item %d: %w", total, err)
108+
}
109+
110+
params, err := toInsertParams(job)
111+
if err != nil {
112+
return err
113+
}
88114

89-
for i := range jobs {
115+
batch = append(batch, params)
116+
total++
117+
count++
90118

91-
var insertOpts *river.InsertOpts
92-
if jobs[i].Metadata != nil {
93-
metadataBytes, err := json.Marshal(jobs[i].Metadata)
94-
if err != nil {
119+
if total%100 == 0 {
120+
log.Info().Msgf("Loaded %d jobs...", total)
121+
}
122+
123+
if count == batchSize {
124+
if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
95125
return err
96126
}
97-
insertOpts = &river.InsertOpts{Metadata: metadataBytes}
127+
log.Info().Msgf("Inserted batch of %d jobs (total: %d)", batchSize, total)
128+
batch = batch[:0]
129+
count = 0
130+
}
131+
}
132+
133+
if len(batch) > 0 {
134+
if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
135+
return err
98136
}
99-
jobsToInsert = append(jobsToInsert, river.InsertManyParams{
100-
Args: jobs[i].Args,
101-
InsertOpts: insertOpts,
102-
})
137+
log.Info().Msgf("Inserted final batch of %d jobs (total: %d)", len(batch), total)
103138
}
104139

105-
results, err := riverClient.InsertManyTx(ctx, tx, jobsToInsert)
140+
log.Info().Msgf("Bulk insert complete. Total jobs inserted: %d", total)
141+
return nil
142+
}
143+
144+
// ------ NDJSON VERSION ------
145+
func processNDJSON(ctx context.Context, decoder *json.Decoder, riverClient *river.Client[pgx.Tx], dbPool *pgxpool.Pool) error {
146+
const batchSize = 500
147+
148+
batch := make([]river.InsertManyParams, 0, batchSize)
149+
total := 0
150+
count := 0
151+
152+
for {
153+
var job BulkJob
154+
if err := decoder.Decode(&job); err != nil {
155+
if errors.Is(err, io.EOF) {
156+
break
157+
}
158+
return err
159+
}
160+
161+
params, err := toInsertParams(job)
162+
if err != nil {
163+
return err
164+
}
165+
166+
batch = append(batch, params)
167+
total++
168+
count++
169+
170+
if total%100 == 0 {
171+
log.Info().Msgf("Loaded %d jobs...", total)
172+
}
173+
174+
if count == batchSize {
175+
if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
176+
return err
177+
}
178+
log.Info().Msgf("Inserted batch of %d jobs (total: %d)", batchSize, total)
179+
batch = batch[:0]
180+
count = 0
181+
}
182+
}
183+
184+
if len(batch) > 0 {
185+
if err := insertBatch(ctx, riverClient, dbPool, batch); err != nil {
186+
return err
187+
}
188+
log.Info().Msgf("Inserted final batch of %d jobs (total: %d)", len(batch), total)
189+
}
190+
191+
log.Info().Msgf("Bulk insert complete. Total jobs inserted: %d", total)
192+
return nil
193+
}
194+
195+
// Helper: Read first non-space token
196+
func nextNonSpaceToken(dec *json.Decoder) (json.Token, error) {
197+
for {
198+
t, err := dec.Token()
199+
if err != nil {
200+
return nil, err
201+
}
202+
if _, ok := t.(json.Delim); ok || t != nil {
203+
return t, nil
204+
}
205+
}
206+
}
207+
208+
func toInsertParams(job BulkJob) (river.InsertManyParams, error) {
209+
var insertOpts *river.InsertOpts
210+
if job.Metadata != nil {
211+
metaBytes, err := json.Marshal(job.Metadata)
212+
if err != nil {
213+
return river.InsertManyParams{}, err
214+
}
215+
insertOpts = &river.InsertOpts{Metadata: metaBytes}
216+
}
217+
return river.InsertManyParams{
218+
Args: job.Args,
219+
InsertOpts: insertOpts,
220+
}, nil
221+
}
222+
223+
func insertBatch(ctx context.Context, riverClient *river.Client[pgx.Tx], dbPool *pgxpool.Pool, batch []river.InsertManyParams) error {
224+
tx, err := dbPool.Begin(ctx)
106225
if err != nil {
107226
return err
108227
}
228+
defer tx.Rollback(ctx)
109229

110-
if err := tx.Commit(ctx); err != nil {
230+
if _, err := riverClient.InsertManyTx(ctx, tx, batch); err != nil {
111231
return err
112232
}
113233

114-
log.Info().Msgf("Successfully enqueued %d jobs", len(results))
115-
return nil
234+
return tx.Commit(ctx)
116235
}
117236

237+
118238
func newRiverClient(driver *riverpgxv5.Driver) (*river.Client[pgx.Tx], error) {
119239
return river.NewClient[pgx.Tx](driver, &river.Config{})
120240
}

‎main.go‎

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,13 @@ func main() {
2727
n := flag.Int("n", 10, "Number of results to display (view mode)")
2828

2929
queueN := flag.Int("queue-n", 10, "Number of queued jobs to display (for view action)")
30-
queueAction := flag.String("action", "", "Queue action: 'view' to list queued jobs, 'clear' to delete all queued jobs")
30+
queueAction := flag.String(
31+
"action",
32+
"",
33+
"Queue action: 'view', 'clear', 'count', 'completed-count', 'completed-first', 'completed-last'",
34+
)
35+
36+
3137

3238
flag.Parse()
3339

@@ -80,9 +86,28 @@ func main() {
8086
if err := ClearQueuedJobs(ctx, dbPool); err != nil {
8187
log.Fatal().Err(err).Msg("Failed to clear queued jobs")
8288
}
89+
case "count":
90+
if err := CountQueuedJobs(ctx, dbPool); err != nil {
91+
log.Fatal().Err(err).Msg("Failed to count queued jobs")
92+
}
93+
case "completed-count":
94+
if err := CountCompletedJobs(ctx, dbPool); err != nil {
95+
log.Fatal().Err(err).Msg("Failed to count completed jobs")
96+
}
97+
case "completed-first":
98+
if err := ViewFirstCompletedJobs(ctx, dbPool, *queueN); err != nil {
99+
log.Fatal().Err(err).Msg("Failed to view first completed jobs")
100+
}
101+
102+
case "completed-last":
103+
if err := ViewLastCompletedJobs(ctx, dbPool, *queueN); err != nil {
104+
log.Fatal().Err(err).Msg("Failed to view last completed jobs")
105+
}
106+
83107
default:
84108
log.Fatal().Str("action", *queueAction).Msg("Unknown queue action")
85109
}
110+
86111
default:
87112
log.Fatal().Str("mode", *mode).Msg("Unknown mode")
88113
}

‎managequeuedjobs.go‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,24 @@ import (
88
"github.com/jackc/pgx/v5/pgxpool"
99
)
1010

11+
func CountQueuedJobs(ctx context.Context, db *pgxpool.Pool) error {
12+
var count int64
13+
14+
err := db.QueryRow(ctx, `
15+
SELECT COUNT(*)
16+
FROM river_job
17+
WHERE state IN ('available', 'scheduled')
18+
`).Scan(&count)
19+
20+
if err != nil {
21+
return err
22+
}
23+
24+
fmt.Printf("Total queued jobs remaining: %d\n", count)
25+
return nil
26+
}
27+
28+
1129
func ViewQueuedJobs(ctx context.Context, db *pgxpool.Pool, n int) error {
1230
rows, err := db.Query(ctx, `
1331
SELECT id, state, created_at, scheduled_at
@@ -46,3 +64,100 @@ func ClearQueuedJobs(ctx context.Context, db *pgxpool.Pool) error {
4664
fmt.Printf("Deleted %d queued jobs\n", res.RowsAffected())
4765
return nil
4866
}
67+
68+
func CountCompletedJobs(ctx context.Context, db *pgxpool.Pool) error {
69+
var count int64
70+
71+
err := db.QueryRow(ctx, `
72+
SELECT COUNT(*)
73+
FROM river_job
74+
WHERE state = 'completed'
75+
`).Scan(&count)
76+
77+
if err != nil {
78+
return err
79+
}
80+
81+
fmt.Printf("Total completed jobs: %d\n", count)
82+
return nil
83+
}
84+
85+
func ViewFirstCompletedJobs(ctx context.Context, db *pgxpool.Pool, n int) error {
86+
rows, err := db.Query(ctx, `
87+
SELECT id, created_at, finalized_at, args
88+
FROM river_job
89+
WHERE state = 'completed'
90+
ORDER BY finalized_at ASC
91+
LIMIT $1
92+
`, n)
93+
if err != nil {
94+
return err
95+
}
96+
defer rows.Close()
97+
98+
fmt.Printf("First %d completed jobs:\n", n)
99+
for rows.Next() {
100+
var (
101+
id int64
102+
createdAt time.Time
103+
finalizedAt time.Time
104+
args []byte
105+
)
106+
107+
if err := rows.Scan(&id, &createdAt, &finalizedAt, &args); err != nil {
108+
return err
109+
}
110+
111+
fmt.Printf(
112+
"ID: %d | CreatedAt: %s | CompletedAt: %s | Args: %s\n",
113+
id,
114+
createdAt.Format(time.RFC3339),
115+
finalizedAt.Format(time.RFC3339),
116+
string(args),
117+
)
118+
}
119+
120+
return rows.Err()
121+
}
122+
123+
124+
125+
func ViewLastCompletedJobs(ctx context.Context, db *pgxpool.Pool, n int) error {
126+
rows, err := db.Query(ctx, `
127+
SELECT id, created_at, finalized_at, args
128+
FROM river_job
129+
WHERE state = 'completed'
130+
ORDER BY finalized_at DESC
131+
LIMIT $1
132+
`, n)
133+
if err != nil {
134+
return err
135+
}
136+
defer rows.Close()
137+
138+
fmt.Printf("Last %d completed jobs:\n", n)
139+
for rows.Next() {
140+
var (
141+
id int64
142+
createdAt time.Time
143+
finalizedAt time.Time
144+
args []byte
145+
)
146+
147+
if err := rows.Scan(&id, &createdAt, &finalizedAt, &args); err != nil {
148+
return err
149+
}
150+
151+
fmt.Printf(
152+
"ID: %d | CreatedAt: %s | CompletedAt: %s | Args: %s\n",
153+
id,
154+
createdAt.Format(time.RFC3339),
155+
finalizedAt.Format(time.RFC3339),
156+
string(args),
157+
)
158+
}
159+
160+
return rows.Err()
161+
}
162+
163+

0 commit comments

Comments
 (0)