DB query performance improvements

internal/structset/structset.go

@@ -3,8 +3,6 @@ package structset
 
 import (
 	"database/sql"
-	"errors"
-	"fmt"
 	"reflect"
 	"strings"
 	"sync"
@@ -81,34 +79,24 @@ func GetStructFieldTagMap(Struct interface{}, keyTag string, valueTag string) ma
 // ScanRow is a cut-down version of the proposed Rows.ScanRow method. It
 // currently only handles dest being a (pointer to) struct, and does not
 // handle embedded fields. See https://github.com/golang/go/issues/61637
-func ScanRow(rows *sql.Rows, dest any) error {
-	rv := reflect.ValueOf(dest)
-	if rv.Kind() != reflect.Pointer || rv.IsNil() {
-		return errors.New("dest must be a non-nil pointer")
-	}
+func ScanRow(rows *sql.Rows, columns []string, indexes map[string]int, dest any) error {
+	// elem := reflect.ValueOf(dest).Elem()
+	// if rv.Kind() != reflect.Pointer || rv.IsNil() {
+	// 	return errors.New("dest must be a non-nil pointer")
+	// }
 
-	elem := rv.Elem()
-	if elem.Kind() != reflect.Struct {
-		return errors.New("dest must point to a struct")
-	}
-	indexes := cachedFieldIndexes(reflect.TypeOf(dest).Elem())
-
-	columns, err := rows.Columns()
-	if err != nil {
-		return fmt.Errorf("cannot fetch columns: %w", err)
-	}
+	// elem := rv.Elem()
+	// if elem.Kind() != reflect.Struct {
+	// 	return errors.New("dest must point to a struct")
+	// }
 
 	var scanArgs []any
 	for _, column := range columns {
 		index, ok := indexes[column]
 		if ok {
 			// We have a column to field mapping, scan the value.
-			field := elem.Field(index)
+			field := reflect.ValueOf(dest).Elem().Field(index)
 			scanArgs = append(scanArgs, field.Addr().Interface())
-		} else {
-			// Unassigned column, throw away the scanned value.
-			var throwAway any
-			scanArgs = append(scanArgs, &throwAway)
 		}
 	}
 	return rows.Scan(scanArgs...)
@@ -131,8 +119,8 @@ func fieldIndexes(structType reflect.Type) map[string]int {
 	return indexes
 }
 
-// cachedFieldIndexes is like fieldIndexes, but cached per struct type.
-func cachedFieldIndexes(structType reflect.Type) map[string]int {
+// CachedFieldIndexes is like fieldIndexes, but cached per struct type.
+func CachedFieldIndexes(structType reflect.Type) map[string]int {
 	if f, ok := fieldIndexesCache.Load(structType); ok {
 		return f.(map[string]int)
 	}

pkg/api/cli/cli_test.go

@@ -106,6 +106,14 @@ ceems_api_server:
 	configFile := fmt.Sprintf(configFileTmpl, dataDir)
 	configFilePath := makeConfigFile(configFile, tmpDir)
 
+	// Create sample DB file
+	os.MkdirAll(dataDir, os.ModePerm)
+	f, err := os.Create(filepath.Join(dataDir, base.CEEMSDBName))
+	if err != nil {
+		require.NoError(t, err)
+	}
+	f.Close()
+
 	// Remove test related args
 	os.Args = append([]string{os.Args[0]}, fmt.Sprintf("--config.file=%s", configFilePath))
 	os.Args = append(os.Args, "--log.level=debug")

pkg/api/db/migrations/000006_create_index_units_table.down.sql

@@ -0,0 +1 @@
+DROP INDEX IF EXISTS idx_cluster_id_project_user_ended;

pkg/api/db/migrations/000006_create_index_units_table.up.sql

@@ -0,0 +1 @@
+CREATE INDEX idx_cluster_id_project_user_ended ON units (cluster_id,project,username,ended_at);

pkg/api/http/docs/docs.go

@@ -214,7 +214,7 @@ const docTemplate = `{
                 "tags": [
                     "projects"
                 ],
-                "summary": "Admin ednpoint to fetch project details",
+                "summary": "Admin endpoint to fetch project details",
                 "parameters": [
                     {
                         "type": "string",
@@ -706,11 +706,21 @@ const docTemplate = `{
                         "items": {
                             "type": "string"
                         },
-                        "collectionFormat": "multi",
+                        "collectionFormat": "csv",
                         "description": "Project",
                         "name": "project",
                         "in": "query"
                     },
+                    {
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "collectionFormat": "multi",
+                        "description": "Username",
+                        "name": "user",
+                        "in": "query"
+                    },
                     {
                         "type": "string",
                         "description": "From timestamp",

pkg/api/http/docs/swagger.json

@@ -206,7 +206,7 @@
                 "tags": [
                     "projects"
                 ],
-                "summary": "Admin ednpoint to fetch project details",
+                "summary": "Admin endpoint to fetch project details",
                 "parameters": [
                     {
                         "type": "string",
@@ -698,11 +698,21 @@
                         "items": {
                             "type": "string"
                         },
-                        "collectionFormat": "multi",
+                        "collectionFormat": "csv",
                         "description": "Project",
                         "name": "project",
                         "in": "query"
                     },
+                    {
+                        "type": "array",
+                        "items": {
+                            "type": "string"
+                        },
+                        "collectionFormat": "multi",
+                        "description": "Username",
+                        "name": "user",
+                        "in": "query"
+                    },
                     {
                         "type": "string",
                         "description": "From timestamp",

pkg/api/http/docs/swagger.yaml

@@ -608,7 +608,7 @@ paths:
             $ref: '#/definitions/http.Response-any'
       security:
       - BasicAuth: []
-      summary: Admin ednpoint to fetch project details
+      summary: Admin endpoint to fetch project details
       tags:
       - projects
   /units:
@@ -995,13 +995,20 @@ paths:
           type: string
         name: cluster_id
         type: array
-      - collectionFormat: multi
+      - collectionFormat: csv
         description: Project
         in: query
         items:
           type: string
         name: project
         type: array
+      - collectionFormat: multi
+        description: Username
+        in: query
+        items:
+          type: string
+        name: user
+        type: array
       - description: From timestamp
         in: query
         name: from

pkg/api/http/querier.go

@@ -3,6 +3,7 @@ package http
 import (
 	"database/sql"
 	"fmt"
+	"reflect"
 	"regexp"
 	"strings"
 
@@ -76,13 +77,26 @@ func projectsSubQuery(users []string) Query {
 // Ref: https://oilbeater.com/en/2024/03/04/golang-slice-performance/
 // For the rest of queries, they should return fewer rows and hence, we can live with
 // dynamic allocation
-func scanRows[T any](rows *sql.Rows, numRows int) []T {
+func scanRows[T any](rows *sql.Rows, numRows int) ([]T, error) {
+	var columns []string
 	var values = make([]T, numRows)
 	var value T
+	var err error
+	scanErrs := 0
 	rowIdx := 0
+
+	// Get indexes
+	indexes := structset.CachedFieldIndexes(reflect.TypeOf(&value).Elem())
+
+	// Get columns
+	if columns, err = rows.Columns(); err != nil {
+		return nil, fmt.Errorf("cannot fetch columns: %w", err)
+	}
+
+	// Scan each row
 	for rows.Next() {
-		if err := structset.ScanRow(rows, &value); err != nil {
-			continue
+		if err := structset.ScanRow(rows, columns, indexes, &value); err != nil {
+			scanErrs++
 		}
 		if numRows > 0 {
 			values[rowIdx] = value
@@ -91,7 +105,13 @@ func scanRows[T any](rows *sql.Rows, numRows int) []T {
 		}
 		rowIdx++
 	}
-	return values
+
+	// If we failed to scan any rows, return error which will be included in warnings
+	// in the response
+	if scanErrs > 0 {
+		err = fmt.Errorf("failed to scan %d rows", scanErrs)
+	}
+	return values, err
 }
 
 func countRows(dbConn *sql.DB, query Query) (int, error) {
@@ -181,10 +201,9 @@ func Querier[T any](dbConn *sql.DB, query Query, logger log.Logger) ([]T, error)
 	defer rows.Close()
 
 	// Loop through rows, using Scan to assign column data to struct fields.
-	var values = scanRows[T](rows, numRows)
 	level.Debug(logger).Log(
 		"msg", "Rows", "query", queryString, "queryParams", strings.Join(queryParams, ","),
-		"num_rows", numRows,
+		"num_rows", numRows, "error", err,
 	)
-	return values, nil
+	return scanRows[T](rows, numRows)
 }