Skip to content

Commit 368c78d

Browse files
Thanus-KumaarThanus-Kumaar
committed
Refactoring of notebook module into repo
1 parent c053d20 commit 368c78d

File tree

3 files changed

+361
-313
lines changed

3 files changed

+361
-313
lines changed

‎db/repository/notebook_repository.go‎

Lines changed: 333 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,333 @@
1+
package repository
2+
3+
import (
4+
"context"
5+
"database/sql"
6+
"errors"
7+
"strconv"
8+
"time"
9+
10+
"github.com/Thanus-Kumaar/controller_microservice_v2/pkg/models"
11+
"github.com/google/uuid"
12+
"github.com/jackc/pgx/v4/pgxpool"
13+
)
14+
15+
type NotebookRepository interface {
16+
CreateNotebook(ctx context.Context, req *models.CreateNotebookRequest) (*models.Notebook, error)
17+
ListNotebooks(ctx context.Context, filters map[string]string) ([]models.Notebook, error)
18+
GetNotebookByID(ctx context.Context, id string) (*models.Notebook, error)
19+
UpdateNotebook(ctx context.Context, id string, req *models.UpdateNotebookRequest) (*models.Notebook, error)
20+
DeleteNotebook(ctx context.Context, id string) error
21+
}
22+
23+
type notebookRepository struct {
24+
pool *pgxpool.Pool
25+
}
26+
27+
func NewNotebookRepository(pool *pgxpool.Pool) NotebookRepository {
28+
return &notebookRepository{
29+
pool: pool,
30+
}
31+
}
32+
33+
// Implementations will be moved here.
34+
func (r *notebookRepository) CreateNotebook(ctx context.Context, req *models.CreateNotebookRequest) (*models.Notebook, error) {
35+
id := uuid.New().String()
36+
now := time.Now().UTC()
37+
38+
query := `
39+
INSERT INTO notebooks (id, title, context_minio_url, problem_statement_id, created_at, last_modified_at)
40+
VALUES ($1, $2, $3, $4, $5, $6)
41+
RETURNING id, title, context_minio_url, problem_statement_id, created_at, last_modified_at;
42+
`
43+
44+
row := r.pool.QueryRow(ctx, query,
45+
id,
46+
req.Title,
47+
nil, // TODO: Should include logic for context minIO url
48+
req.ProblemStatementID,
49+
now,
50+
now,
51+
)
52+
53+
var nb models.Notebook
54+
if err := row.Scan(
55+
&nb.ID,
56+
&nb.Title,
57+
&nb.ContextMinioURL,
58+
&nb.ProblemStatementID,
59+
&nb.CreatedAt,
60+
&nb.LastModifiedAt,
61+
); err != nil {
62+
return nil, err
63+
}
64+
65+
return &nb, nil
66+
}
67+
68+
func (r *notebookRepository) ListNotebooks(ctx context.Context, filters map[string]string) ([]models.Notebook, error) {
69+
query := `SELECT id, title, context_minio_url, problem_statement_id, created_at, last_modified_at FROM notebooks`
70+
args := []any{}
71+
where := ""
72+
73+
if filters != nil {
74+
i := 1
75+
for key, val := range filters {
76+
if i == 1 {
77+
where += " WHERE "
78+
} else {
79+
where += " AND "
80+
}
81+
where += key + " = $" + strconv.Itoa(i)
82+
args = append(args, val)
83+
i++
84+
}
85+
}
86+
87+
rows, err := r.pool.Query(ctx, query+where, args...)
88+
if err != nil {
89+
return nil, err
90+
}
91+
defer rows.Close()
92+
93+
var notebooks []models.Notebook
94+
for rows.Next() {
95+
var nb models.Notebook
96+
if err := rows.Scan(
97+
&nb.ID,
98+
&nb.Title,
99+
&nb.ContextMinioURL,
100+
&nb.ProblemStatementID,
101+
&nb.CreatedAt,
102+
&nb.LastModifiedAt,
103+
); err != nil {
104+
return nil, err
105+
}
106+
notebooks = append(notebooks, nb)
107+
}
108+
109+
return notebooks, nil
110+
}
111+
112+
func (r *notebookRepository) GetNotebookByID(ctx context.Context, id string) (*models.Notebook, error) {
113+
notebookUUID, err := uuid.Parse(id)
114+
if err != nil {
115+
return nil, errors.New("invalid notebook ID format")
116+
}
117+
118+
query := `
119+
SELECT
120+
n.id, n.title, n.context_minio_url, n.problem_statement_id, n.created_at, n.last_modified_at,
121+
c.id, c.notebook_id, c.cell_index, c.cell_type, c.source, c.execution_count,
122+
co.id, co.cell_id, co.output_index, co.type, co.data_json, co.minio_url, co.execution_count,
123+
er.id, er.source_cell_id, er.start_time, er.end_time, er.status,
124+
cv.id, cv.evolution_run_id, cv.code, cv.metric, cv.is_best, cv.generation, cv.parent_variant_id
125+
FROM
126+
notebooks n
127+
LEFT JOIN
128+
cells c ON n.id = c.notebook_id
129+
LEFT JOIN
130+
cell_outputs co ON c.id = co.cell_id
131+
LEFT JOIN
132+
evolution_runs er ON c.id = er.source_cell_id
133+
LEFT JOIN
134+
cell_variations cv ON er.id = cv.evolution_run_id
135+
WHERE
136+
n.id = $1
137+
ORDER BY
138+
c.cell_index ASC,
139+
co.output_index ASC,
140+
er.start_time ASC,
141+
cv.generation ASC;
142+
`
143+
144+
rows, err := r.pool.Query(ctx, query, notebookUUID)
145+
if err != nil {
146+
return nil, err
147+
}
148+
defer rows.Close()
149+
150+
var notebook models.Notebook
151+
cellMap := make(map[uuid.UUID]*models.Cell)
152+
runMap := make(map[uuid.UUID]*models.EvolutionRun)
153+
var orderedCellIDs []uuid.UUID
154+
155+
for rows.Next() {
156+
var (
157+
cellID uuid.NullUUID
158+
cellNotebookID uuid.NullUUID
159+
cellIndex sql.NullInt32
160+
cellType sql.NullString
161+
cellSource sql.NullString
162+
cellExecCount sql.NullInt32
163+
outputID uuid.NullUUID
164+
outputCellID uuid.NullUUID
165+
outputIndex sql.NullInt32
166+
outputType sql.NullString
167+
outputDataJSON []byte
168+
outputMinioURL sql.NullString
169+
outputExecCount sql.NullInt32
170+
erID uuid.NullUUID
171+
erSourceCellID uuid.NullUUID
172+
erStartTime sql.NullTime
173+
erEndTime sql.NullTime
174+
erStatus sql.NullString
175+
cvID uuid.NullUUID
176+
cvEvolutionRunID uuid.NullUUID
177+
cvCode sql.NullString
178+
cvMetric sql.NullFloat64
179+
cvIsBest sql.NullBool
180+
cvGeneration sql.NullInt32
181+
cvParentVariantID uuid.NullUUID
182+
)
183+
184+
if err := rows.Scan(
185+
&notebook.ID, &notebook.Title, &notebook.ContextMinioURL, &notebook.ProblemStatementID, &notebook.CreatedAt, &notebook.LastModifiedAt,
186+
&cellID, &cellNotebookID, &cellIndex, &cellType, &cellSource, &cellExecCount,
187+
&outputID, &outputCellID, &outputIndex, &outputType, &outputDataJSON, &outputMinioURL, &outputExecCount,
188+
&erID, &erSourceCellID, &erStartTime, &erEndTime, &erStatus,
189+
&cvID, &cvEvolutionRunID, &cvCode, &cvMetric, &cvIsBest, &cvGeneration, &cvParentVariantID,
190+
); err != nil {
191+
return nil, err
192+
}
193+
194+
if cellID.Valid {
195+
if _, exists := cellMap[cellID.UUID]; !exists {
196+
cellMap[cellID.UUID] = &models.Cell{
197+
ID: cellID.UUID,
198+
NotebookID: cellNotebookID.UUID,
199+
CellIndex: int(cellIndex.Int32),
200+
CellType: cellType.String,
201+
Source: cellSource.String,
202+
ExecutionCount: int(cellExecCount.Int32),
203+
Outputs: []models.CellOutput{},
204+
EvolutionRuns: []models.EvolutionRun{},
205+
}
206+
orderedCellIDs = append(orderedCellIDs, cellID.UUID)
207+
}
208+
}
209+
210+
if outputID.Valid {
211+
if cell, exists := cellMap[outputCellID.UUID]; exists {
212+
cell.Outputs = append(cell.Outputs, models.CellOutput{
213+
ID: outputID.UUID,
214+
CellID: outputCellID.UUID,
215+
OutputIndex: int(outputIndex.Int32),
216+
Type: outputType.String,
217+
DataJSON: outputDataJSON,
218+
MinioURL: outputMinioURL.String,
219+
ExecutionCount: int(outputExecCount.Int32),
220+
})
221+
}
222+
}
223+
224+
if erID.Valid {
225+
if _, exists := runMap[erID.UUID]; !exists {
226+
run := models.EvolutionRun{
227+
ID: erID.UUID,
228+
SourceCellID: erSourceCellID.UUID,
229+
StartTime: erStartTime.Time,
230+
Status: erStatus.String,
231+
Variations: []models.CellVariation{},
232+
}
233+
if erEndTime.Valid {
234+
run.EndTime = &erEndTime.Time
235+
}
236+
runMap[erID.UUID] = &run
237+
}
238+
}
239+
240+
if cvID.Valid {
241+
if run, exists := runMap[cvEvolutionRunID.UUID]; exists {
242+
variation := models.CellVariation{
243+
ID: cvID.UUID,
244+
EvolutionRunID: cvEvolutionRunID.UUID,
245+
Code: cvCode.String,
246+
Metric: cvMetric.Float64,
247+
IsBest: cvIsBest.Bool,
248+
Generation: int(cvGeneration.Int32),
249+
}
250+
if cvParentVariantID.Valid {
251+
variation.ParentVariantID = &cvParentVariantID.UUID
252+
}
253+
run.Variations = append(run.Variations, variation)
254+
}
255+
}
256+
}
257+
258+
if notebook.ID == "" {
259+
return nil, errors.New("notebook not found")
260+
}
261+
262+
for _, run := range runMap {
263+
if cell, exists := cellMap[run.SourceCellID]; exists {
264+
cell.EvolutionRuns = append(cell.EvolutionRuns, *run)
265+
}
266+
}
267+
268+
notebook.Cells = make([]models.Cell, len(orderedCellIDs))
269+
for i, cellID := range orderedCellIDs {
270+
notebook.Cells[i] = *cellMap[cellID]
271+
}
272+
273+
return &notebook, nil
274+
}
275+
276+
func (r *notebookRepository) UpdateNotebook(ctx context.Context, id string, req *models.UpdateNotebookRequest) (*models.Notebook, error) {
277+
setClause := ""
278+
args := []any{}
279+
argIndex := 1
280+
281+
if req.Title != nil {
282+
setClause += "title = $" + strconv.Itoa(argIndex)
283+
args = append(args, *req.Title)
284+
argIndex++
285+
}
286+
287+
if setClause == "" {
288+
// This business logic should be in the module.
289+
// For now, let's just return the notebook as is.
290+
// A better approach is for the module to call GetNotebookByID if req is empty.
291+
return r.GetNotebookByID(ctx, id)
292+
}
293+
294+
setClause += ", last_modified_at = $" + strconv.Itoa(argIndex)
295+
args = append(args, time.Now().UTC())
296+
argIndex++
297+
298+
args = append(args, id)
299+
300+
query := `
301+
UPDATE notebooks
302+
SET ` + setClause + `
303+
WHERE id = $` + strconv.Itoa(argIndex) + `
304+
RETURNING id, title, context_minio_url, problem_statement_id, created_at, last_modified_at;
305+
`
306+
307+
row := r.pool.QueryRow(ctx, query, args...)
308+
309+
var nb models.Notebook
310+
if err := row.Scan(
311+
&nb.ID,
312+
&nb.Title,
313+
&nb.ContextMinioURL,
314+
&nb.ProblemStatementID,
315+
&nb.CreatedAt,
316+
&nb.LastModifiedAt,
317+
); err != nil {
318+
return nil, err
319+
}
320+
321+
return &nb, nil
322+
}
323+
324+
func (r *notebookRepository) DeleteNotebook(ctx context.Context, id string) error {
325+
cmd, err := r.pool.Exec(ctx, `DELETE FROM notebooks WHERE id = $1;`, id)
326+
if err != nil {
327+
return err
328+
}
329+
if cmd.RowsAffected() == 0 {
330+
return errors.New("notebook not found")
331+
}
332+
return nil
333+
}

0 commit comments

Comments
 (0)