| 1 | // SPDX-License-Identifier: GPL-3.0-or-later |
| 2 | |
| 3 | package sql |
| 4 | |
| 5 | import ( |
| 6 | "context" |
| 7 | "fmt" |
| 8 | "maps" |
| 9 | "time" |
| 10 | ) |
| 11 | |
| 12 | // queryRowsCache maps query_id -> slice of row maps (col -> string value) |
| 13 | type queryRowsCache map[string][]map[string]string |
| 14 | |
| 15 | // execReusableQueries runs all queries declared in Config.Queries (new schema) |
| 16 | // and returns their full rowsets in memory. It also returns per-query durations (ms). |
| 17 | func (c *Collector) execReusableQueries(ctx context.Context) (queryRowsCache, map[string]int64, error) { |
| 18 | cache := make(queryRowsCache, len(c.Queries)) |
| 19 | durations := make(map[string]int64, len(c.Queries)) |
| 20 | |
| 21 | for i, q := range c.Queries { |
| 22 | if q.ID == "" { |
| 23 | return nil, nil, fmt.Errorf("queries[%d] missing id", i+1) |
| 24 | } |
| 25 | if q.Query == "" { |
| 26 | return nil, nil, fmt.Errorf("queries[%d] missing query", i+1) |
| 27 | } |
| 28 | |
| 29 | rows, dur, err := c.runSQL(ctx, q.Query) |
| 30 | if err != nil { |
| 31 | return nil, nil, fmt.Errorf("query %q failed: %w", q.ID, err) |
| 32 | } |
| 33 | cache[q.ID] = rows |
| 34 | durations[q.ID] = dur |
| 35 | } |
| 36 | |
| 37 | return cache, durations, nil |
| 38 | } |
| 39 | |
| 40 | // execMetricQueries resolves and executes the query for each metric block. |
| 41 | // If a metric uses query_ref, it reuses rows from qcache (no re-query). |
| 42 | // If a metric has inline query, it executes it and stores the rows. |
| 43 | func (c *Collector) execMetricQueries(ctx context.Context, qcache queryRowsCache) (queryRowsCache, map[string]int64, error) { |
| 44 | cache := make(queryRowsCache, len(c.Metrics)) |
| 45 | durations := make(map[string]int64, len(c.Metrics)) |
| 46 | |
| 47 | for i, m := range c.Metrics { |
| 48 | if m.ID == "" { |
| 49 | return nil, nil, fmt.Errorf("metrics[%d] missing id", i+1) |
| 50 | } |
| 51 | |
| 52 | switch { |
| 53 | case m.QueryRef != "": |
| 54 | // reuse pre-fetched rows; duration is 0 because we didn't re-run |
| 55 | rows, ok := qcache[m.QueryRef] |
| 56 | if !ok { |
| 57 | return nil, nil, fmt.Errorf("metrics[%d] query_ref %q not found in queries cache", i+1, m.QueryRef) |
| 58 | } |
| 59 | cache[m.ID] = rows |
| 60 | case m.Query != "": |
| 61 | rows, dur, err := c.runSQL(ctx, m.Query) |
| 62 | if err != nil { |
| 63 | return nil, nil, fmt.Errorf("metrics[%d] (%q) query failed: %w", i+1, m.ID, err) |
| 64 | } |
| 65 | cache[m.ID] = rows |
| 66 | durations[m.ID] = dur |
| 67 | default: |
| 68 | return nil, nil, fmt.Errorf("metrics[%d] must set one of query_ref or query", i+1) |
| 69 | } |
| 70 | } |
| 71 | |
| 72 | return cache, durations, nil |
| 73 | } |
| 74 | |
| 75 | // runSQL executes a SQL statement with c.Timeout and returns the rowset as []map[col]value. |
| 76 | // Duration is milliseconds from QueryContext() start to first successful return. |
| 77 | func (c *Collector) runSQL(ctx context.Context, query string) ([]map[string]string, int64, error) { |
| 78 | qctx := ctx |
| 79 | cancel := func() {} |
| 80 | if d := c.Timeout.Duration(); d > 0 { |
| 81 | qctx, cancel = context.WithTimeout(ctx, d) |
| 82 | } |
| 83 | defer cancel() |
| 84 | |
| 85 | start := time.Now() |
| 86 | rows, err := c.db.QueryContext(qctx, query) |
| 87 | if err != nil { |
| 88 | return nil, 0, err |
| 89 | } |
| 90 | defer func() { _ = rows.Close() }() |
| 91 | |
| 92 | duration := time.Since(start).Milliseconds() |
| 93 | |
| 94 | columns, err := rows.Columns() |
| 95 | if err != nil { |
| 96 | return nil, duration, err |
| 97 | } |
| 98 | |
| 99 | scan := makeRawBytesSlice(len(columns)) |
| 100 | out := make([]map[string]string, 0, 64) |
| 101 | row := make(map[string]string, len(columns)) |
| 102 | |
| 103 | for rows.Next() { |
| 104 | if err := rows.Scan(scan...); err != nil { |
| 105 | return nil, duration, err |
| 106 | } |
| 107 | clear(row) |
| 108 | for i := range columns { |
| 109 | row[columns[i]] = rawBytesToString(scan[i]) |
| 110 | } |
| 111 | |
| 112 | out = append(out, maps.Clone(row)) |
| 113 | } |
| 114 | |
| 115 | return out, duration, rows.Err() |
| 116 | } |