Skip to content

Support forward/backward direction in loki_query and fix tests #19

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+43 −15
Open

Support forward/backward direction in loki_query and fix tests #19

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ The `loki_query` tool allows you to query Grafana Loki log data:
- `end`: End time for the query (default: now)
- `limit`: Maximum number of entries to return (default: 100)
- `org`: Organization ID for the query (sent as X-Scope-OrgID header)
- `direction`: Sort order of logs: forward (oldest first) or backward (newest first, default: backward)

#### Environment Variables

Expand Down
22 changes: 17 additions & 5 deletions cmd/client/main.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -83,15 +83,15 @@ func main() {
switch os.Args[1] {
case "loki_query":
if len(os.Args) < 3 {
fmt.Println("Usage: client loki_query [url] <query> [start] [end] [limit]")
fmt.Println("Usage: client loki_query [url] <query> [start] [end] [limit] [org] [direction]")
fmt.Println("Examples:")
fmt.Println(" client loki_query \"{job=\\\"varlogs\\\"}\"")
fmt.Println(" client loki_query http://localhost:3100 \"{job=\\\"varlogs\\\"}\"")
fmt.Println(" client loki_query \"{job=\\\"varlogs\\\"}\" \"-1h\" \"now\" 100")
os.Exit(1)
}

var lokiURL, query, start, end, org string
var lokiURL, query, start, end, org, direction string
var limit float64
// Check if the first argument is a URL or a query
if strings.HasPrefix(os.Args[2], "http") {
Expand Down Expand Up @@ -124,6 +124,10 @@ func main() {
if len(os.Args) > argOffset+3 {
org = os.Args[argOffset+3]
}

if len(os.Args) > argOffset+4 {
direction = os.Args[argOffset+4]
}
} else {
// First arg is the query (URL comes from environment)
query = os.Args[2]
Expand All @@ -149,10 +153,14 @@ func main() {
if len(os.Args) > argOffset+3 {
org = os.Args[argOffset+3]
}

if len(os.Args) > argOffset+4 {
direction = os.Args[argOffset+4]
}
}

// Create the Loki query request
req = createLokiQueryRequest(lokiURL, query, start, end, limit, org)
req = createLokiQueryRequest(lokiURL, query, start, end, limit, org, direction)

default:
showUsage()
Expand Down Expand Up @@ -216,15 +224,15 @@ func main() {

func showUsage() {
fmt.Println("Usage:")
fmt.Println(" client loki_query [url] <query> [start] [end] [limit]")
fmt.Println(" client loki_query [url] <query> [start] [end] [limit] [org] [direction]")
fmt.Println(" Examples:")
fmt.Println(" client loki_query \"{job=\\\"varlogs\\\"}\"")
fmt.Println(" client loki_query http://localhost:3100 \"{job=\\\"varlogs\\\"}\"")
fmt.Println(" client loki_query \"{job=\\\"varlogs\\\"}\" \"-1h\" \"now\" 100")
fmt.Println(" client loki_query \"{job=\\\"varlogs\\\"}\" \"-1h\" \"now\" 100 \"tenant-123\"")
}

func createLokiQueryRequest(url, query, start, end string, limit float64, org string) Request {
func createLokiQueryRequest(url, query, start, end string, limit float64, org, direction string) Request {
// Create arguments map
args := map[string]any{
"query": query,
Expand Down Expand Up @@ -252,6 +260,10 @@ func createLokiQueryRequest(url, query, start, end string, limit float64, org st
args["org"] = org
}

if direction != "" {
args["direction"] = direction
}

return Request{
JSONRPC: "2.0",
ID: "1",
Expand Down
19 changes: 17 additions & 2 deletions internal/handlers/loki.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,10 @@ func NewLokiQueryTool() mcp.Tool {
mcp.Description("Output format: raw, json, or text (default: raw)"),
mcp.DefaultString("raw"),
),
mcp.WithString("direction",
mcp.Description("Sort order of logs: forward (oldest first) or backward (newest first, default: backward)"),
mcp.DefaultString("backward"),
),
)
}

Expand Down Expand Up @@ -205,8 +209,18 @@ func HandleLokiQuery(ctx context.Context, request mcp.CallToolRequest) (*mcp.Cal
format = formatArg
}

// Extract direction parameter
direction := "backward" // default
if directionArg, ok := args["direction"].(string); ok && directionArg != "" {
// Validate direction parameter
if directionArg != "forward" && directionArg != "backward" {
return nil, fmt.Errorf("invalid direction: %s. Must be 'forward' or 'backward'", directionArg)
}
direction = directionArg
}

// Build query URL
queryURL, err := buildLokiQueryURL(lokiURL, queryString, start, end, limit)
queryURL, err := buildLokiQueryURL(lokiURL, queryString, start, end, limit, direction)
if err != nil {
return nil, fmt.Errorf("failed to build query URL: %v", err)
}
Expand Down Expand Up @@ -277,7 +291,7 @@ func parseTime(timeStr string) (time.Time, error) {
}

// buildLokiQueryURL constructs the Loki query URL
func buildLokiQueryURL(baseURL, query string, start, end int64, limit int) (string, error) {
func buildLokiQueryURL(baseURL, query string, start, end int64, limit int, direction string) (string, error) {
u, err := url.Parse(baseURL)
if err != nil {
return "", err
Expand All @@ -303,6 +317,7 @@ func buildLokiQueryURL(baseURL, query string, start, end int64, limit int) (stri
q.Set("start", fmt.Sprintf("%d", start))
q.Set("end", fmt.Sprintf("%d", end))
q.Set("limit", fmt.Sprintf("%d", limit))
q.Set("direction", direction)
u.RawQuery = q.Encode()

return u.String(), nil
Expand Down
16 changes: 8 additions & 8 deletions internal/handlers/loki_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ func TestFormatLokiResults_TimestampParsing(t *testing.T) {
}

// Format the results
output, err := formatLokiResults(result)
output, err := formatLokiResults(result, "raw")
if err != nil {
t.Fatalf("formatLokiResults failed: %v", err)
}
Expand Down Expand Up @@ -83,7 +83,7 @@ func TestFormatLokiResults_MultipleTimestamps(t *testing.T) {
},
}

output, err := formatLokiResults(result)
output, err := formatLokiResults(result, "raw")
if err != nil {
t.Fatalf("formatLokiResults failed: %v", err)
}
Expand Down Expand Up @@ -119,14 +119,14 @@ func TestFormatLokiResults_InvalidTimestamp(t *testing.T) {
},
}

output, err := formatLokiResults(result)
output, err := formatLokiResults(result, "raw")
if err != nil {
t.Fatalf("formatLokiResults failed: %v", err)
}

// Should contain the original invalid timestamp as fallback
if !strings.Contains(output, "[invalid-timestamp]") {
t.Errorf("Expected output to contain '[invalid-timestamp]' as fallback, but got:\n%s", output)
if !strings.Contains(output, "invalid-timestamp") {
t.Errorf("Expected output to contain 'invalid-timestamp' as fallback, but got:\n%s", output)
}

// Should still contain the log message
Expand All @@ -145,7 +145,7 @@ func TestFormatLokiResults_EmptyResult(t *testing.T) {
},
}

output, err := formatLokiResults(result)
output, err := formatLokiResults(result, "raw")
if err != nil {
t.Fatalf("formatLokiResults failed: %v", err)
}
Expand Down Expand Up @@ -180,7 +180,7 @@ func TestFormatLokiResults_RecentTimestamp(t *testing.T) {
},
}

output, err := formatLokiResults(result)
output, err := formatLokiResults(result, "raw")
if err != nil {
t.Fatalf("formatLokiResults failed: %v", err)
}
Expand Down Expand Up @@ -241,7 +241,7 @@ func TestFormatLokiResults_NoYear2262Bug(t *testing.T) {
},
}

output, err := formatLokiResults(result)
output, err := formatLokiResults(result, "raw")
if err != nil {
t.Fatalf("formatLokiResults failed: %v", err)
}
Expand Down