Remove todo contexts

4 files changed, 14 insertions, 17 deletions

diff --git a/pkg/data/db.go b/pkg/data/db.go
index cf58caf..63e8ff0 100644
--- a/pkg/data/db.go
+++ b/pkg/data/db.go
@@ -78,8 +78,7 @@ func NewMemDB() *sql.DB {
 }
 
 func createSchema(db *sql.DB) error {
-	ctx := context.TODO()
-	tx, err := db.BeginTx(ctx, nil)
+	tx, err := db.Begin()
 	if err != nil {
 		return err
 	}
@@ -293,9 +292,7 @@ func (q Query) Close() error {
 }
 
 // Create an index
-func (q Query) Get(indexRoot string) (*index.Index, error) {
-	ctx := context.TODO()
-
+func (q Query) Get(ctx context.Context, indexRoot string) (*index.Index, error) {
 	f := FillMany{Db: q.db}
 	docs, err := f.Get(ctx)
 	if err != nil {
@@ -312,9 +309,7 @@ func (q Query) Get(indexRoot string) (*index.Index, error) {
 }
 
 // Write from index to database
-func (q Query) Put(idx index.Index) error {
-	ctx := context.TODO()
-
+func (q Query) Put(ctx context.Context, idx index.Index) error {
 	p, err := NewPutMany(ctx, q.db, idx.Documents)
 	if err != nil {
 		return err
@@ -324,14 +319,12 @@ func (q Query) Put(idx index.Index) error {
 }
 
 // Update database with values from index, removes entries for deleted files
-func (q Query) Update(idx index.Index) error {
-	ctx := context.TODO()
+func (q Query) Update(ctx context.Context, idx index.Index) error {
 	u := UpdateMany{Db: q.db, PathDocs: idx.Documents}
 	return u.Update(ctx)
 }
 
-func (q Query) GetDocument(path string) (*index.Document, error) {
-	ctx := context.TODO()
+func (q Query) GetDocument(ctx context.Context, path string) (*index.Document, error) {
 	f := Fill{Path: path, Db: q.db}
 	return f.Get(ctx)
 }
@@ -385,8 +378,7 @@ func (q Query) PeriodicOptimize(ctx context.Context, d time.Duration) {
 	}
 }
 
-func (q Query) Execute(artifact query.CompilationArtifact) (map[string]*index.Document, error) {
-	ctx := context.TODO()
+func (q Query) Execute(ctx context.Context, artifact query.CompilationArtifact) (map[string]*index.Document, error) {
 	f := FillMany{
 		Db:   q.db,
 		docs: make(map[string]*index.Document),
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 68750a2..2352f21 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -63,7 +63,7 @@ func NewMux(db *data.Query) *http.ServeMux {
 			return
 		}
 
-		pathDocs, err := db.Execute(artifact)
+		pathDocs, err := db.Execute(r.Context(), artifact)
 		if err != nil {
 			w.WriteHeader(http.StatusInternalServerError)
 			w.Write([]byte("Error executing query"))
diff --git a/pkg/server/unix.go b/pkg/server/unix.go
index a46320d..9d8c266 100644
--- a/pkg/server/unix.go
+++ b/pkg/server/unix.go
@@ -6,6 +6,7 @@ import (
 	"log/slog"
 	"net"
 	"sync"
+	"time"
 
 	"github.com/jpappel/atlas/pkg/data"
 	"github.com/jpappel/atlas/pkg/index"
@@ -128,15 +129,18 @@ func (s *UnixServer) handleConn(conn *net.UnixConn, id uint64) {
 			break
 		}
 
-		docs, err := s.Db.Execute(artifact)
+		ctx, cancel := context.WithTimeout(context.Background(), 400*time.Millisecond)
+		docs, err := s.Db.Execute(ctx, artifact)
 		if err != nil {
 			slog.Warn("Failed to execute query",
 				slog.String("query", queryTxt),
 				slog.String("err", err.Error()),
 			)
 			s.writeError(conn, "query execution error")
+			cancel()
 			break
 		}
+		cancel()
 
 		slog.Debug("Sending results")
 		s.writeResults(conn, docs)
diff --git a/pkg/shell/interpreter.go b/pkg/shell/interpreter.go
index 02d5aaf..287c28f 100644
--- a/pkg/shell/interpreter.go
+++ b/pkg/shell/interpreter.go
@@ -1,6 +1,7 @@
 package shell
 
 import (
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -514,7 +515,7 @@ out:
 				return true, errors.New("Type corruption during compilation, expected query.CompilationArtifact")
 			}
 
-			results, err := inter.querier.Execute(artifact)
+			results, err := inter.querier.Execute(context.Background(), artifact)
 			if err != nil {
 				return false, fmt.Errorf("Error occured while excuting query: %s", err)
 			}