feat: SelectKeyed scans rows into a map by column key

sqlx.go

@@ -320,6 +320,11 @@ func (db *DB) Select(dest interface{}, query string, args ...interface{}) error
 	return Select(db, dest, query, args...)
 }
 
+// SelectKeyed using this DB. See SelectKeyed.
+func (db *DB) SelectKeyed(dest interface{}, key string, query string, args ...interface{}) error {
+	return SelectKeyed(db, dest, key, query, args...)
+}
+
 // Get using this DB.
 // Any placeholder parameters are replaced with supplied args.
 // An error is returned if the result set is empty.
@@ -434,6 +439,11 @@ func (tx *Tx) Select(dest interface{}, query string, args ...interface{}) error
 	return Select(tx, dest, query, args...)
 }
 
+// SelectKeyed within a transaction. See SelectKeyed.
+func (tx *Tx) SelectKeyed(dest interface{}, key string, query string, args ...interface{}) error {
+	return SelectKeyed(tx, dest, key, query, args...)
+}
+
 // Queryx within a transaction.
 // Any placeholder parameters are replaced with supplied args.
 func (tx *Tx) Queryx(query string, args ...interface{}) (*Rows, error) {
@@ -519,6 +529,12 @@ func (s *Stmt) Select(dest interface{}, args ...interface{}) error {
 	return Select(&qStmt{s}, dest, "", args...)
 }
 
+// SelectKeyed using the prepared statement. See SelectKeyed.
+// Note that the query portion of errors will be blank; Stmt does not expose its query.
+func (s *Stmt) SelectKeyed(dest interface{}, key string, args ...interface{}) error {
+	return SelectKeyed(&qStmt{s}, dest, key, "", args...)
+}
+
 // Get using the prepared statement.
 // Any placeholder parameters are replaced with supplied args.
 // An error is returned if the result set is empty.
@@ -683,6 +699,30 @@ func Select(q Queryer, dest interface{}, query string, args ...interface{}) erro
 	return scanAll(rows, dest, false)
 }
 
+// SelectKeyed executes a query and StructScans each row into dest, which must
+// be a non-nil pointer to a map. The map is keyed by the column (or db-tagged
+// field) named by key. Map values must be structs or pointers to structs that
+// are StructScan-compatible. Scannable (non-struct) map values are not
+// supported.
+//
+// The destination map is replaced with a new empty map before scanning. When
+// multiple rows share the same key, later rows overwrite earlier ones.
+//
+// Example:
+//
+//	people := map[string]Person{}
+//	err := sqlx.SelectKeyed(db, &people, "email", `SELECT * FROM person`)
+//
+// See also SelectKeyedContext.
+func SelectKeyed(q Queryer, dest interface{}, key string, query string, args ...interface{}) error {
+	rows, err := q.Queryx(query, args...)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+	return scanAllKeyed(rows, dest, key)
+}
+
 // Get does a QueryRow using the provided Queryer, and scans the resulting row
 // to dest.  If dest is scannable, the result must only have one column.  Otherwise,
 // StructScan is used.  Get will return sql.ErrNoRows like row.Scan would.
@@ -1007,6 +1047,91 @@ func StructScan(rows rowsi, dest interface{}) error {
 
 }
 
+// scanAllKeyed scans all rows into a map destination keyed by column/field name key.
+func scanAllKeyed(rows rowsi, dest interface{}, key string) error {
+	if key == "" {
+		return errors.New("SelectKeyed: key must be non-empty")
+	}
+
+	value := reflect.ValueOf(dest)
+	if value.Kind() != reflect.Ptr {
+		return errors.New("must pass a pointer, not a value, to SelectKeyed destination")
+	}
+	if value.IsNil() {
+		return errors.New("nil pointer passed to SelectKeyed destination")
+	}
+	direct := reflect.Indirect(value)
+	if direct.Kind() != reflect.Map {
+		return fmt.Errorf("expected a map, got %s", direct.Kind())
+	}
+
+	mapType := direct.Type()
+	keyType := mapType.Key()
+	elemType := mapType.Elem()
+	isPtr := elemType.Kind() == reflect.Ptr
+	base := reflectx.Deref(elemType)
+
+	if isScannable(base) {
+		return fmt.Errorf("SelectKeyed: map value type %s must be a struct (or *struct), not scannable", base)
+	}
+
+	columns, err := rows.Columns()
+	if err != nil {
+		return err
+	}
+
+	var m *reflectx.Mapper
+	switch rows := rows.(type) {
+	case *Rows:
+		m = rows.Mapper
+	default:
+		m = mapper()
+	}
+
+	fields := m.TraversalsByName(base, columns)
+	if f, err := missingFields(fields); err != nil && !isUnsafe(rows) {
+		return fmt.Errorf("missing destination name %s in %s", columns[f], base)
+	}
+
+	keyTravs := m.TraversalsByName(base, []string{key})
+	if len(keyTravs) == 0 || len(keyTravs[0]) == 0 {
+		return fmt.Errorf("SelectKeyed: key %q not found in %s", key, base)
+	}
+	keyTraversal := keyTravs[0]
+
+	// Replace with a fresh map (mirrors Select resetting slice length).
+	direct.Set(reflect.MakeMap(mapType))
+
+	values := make([]interface{}, len(columns))
+	for rows.Next() {
+		vp := reflect.New(base)
+		v := reflect.Indirect(vp)
+
+		if err := fieldsByTraversal(v, fields, values, true); err != nil {
+			return err
+		}
+		if err := rows.Scan(values...); err != nil {
+			return err
+		}
+
+		keyField := reflectx.FieldByIndexes(v, keyTraversal)
+		mapKey := keyField
+		if !mapKey.Type().AssignableTo(keyType) {
+			if !mapKey.Type().ConvertibleTo(keyType) {
+				return fmt.Errorf("SelectKeyed: cannot use field type %s as map key type %s", mapKey.Type(), keyType)
+			}
+			mapKey = mapKey.Convert(keyType)
+		}
+
+		if isPtr {
+			direct.SetMapIndex(mapKey, vp)
+		} else {
+			direct.SetMapIndex(mapKey, v)
+		}
+	}
+	return rows.Err()
+}
+
 // reflect helpers
 
 func baseType(t reflect.Type, expected reflect.Kind) (reflect.Type, error) {

sqlx_context.go

@@ -62,6 +62,16 @@ func SelectContext(ctx context.Context, q QueryerContext, dest interface{}, quer
 	return scanAll(rows, dest, false)
 }
 
+// SelectKeyedContext is the context variant of SelectKeyed.
+func SelectKeyedContext(ctx context.Context, q QueryerContext, dest interface{}, key string, query string, args ...interface{}) error {
+	rows, err := q.QueryxContext(ctx, query, args...)
+	if err != nil {
+		return err
+	}
+	defer rows.Close()
+	return scanAllKeyed(rows, dest, key)
+}
+
 // PreparexContext prepares a statement.
 //
 // The provided context is used for the preparation of the statement, not for
@@ -141,6 +151,11 @@ func (db *DB) SelectContext(ctx context.Context, dest interface{}, query string,
 	return SelectContext(ctx, db, dest, query, args...)
 }
 
+// SelectKeyedContext using this DB. See SelectKeyed.
+func (db *DB) SelectKeyedContext(ctx context.Context, dest interface{}, key string, query string, args ...interface{}) error {
+	return SelectKeyedContext(ctx, db, dest, key, query, args...)
+}
+
 // GetContext using this DB.
 // Any placeholder parameters are replaced with supplied args.
 // An error is returned if the result set is empty.
@@ -240,6 +255,11 @@ func (c *Conn) SelectContext(ctx context.Context, dest interface{}, query string
 	return SelectContext(ctx, c, dest, query, args...)
 }
 
+// SelectKeyedContext using this Conn. See SelectKeyed.
+func (c *Conn) SelectKeyedContext(ctx context.Context, dest interface{}, key string, query string, args ...interface{}) error {
+	return SelectKeyedContext(ctx, c, dest, key, query, args...)
+}
+
 // GetContext using this Conn.
 // Any placeholder parameters are replaced with supplied args.
 // An error is returned if the result set is empty.
@@ -339,6 +359,11 @@ func (tx *Tx) SelectContext(ctx context.Context, dest interface{}, query string,
 	return SelectContext(ctx, tx, dest, query, args...)
 }
 
+// SelectKeyedContext within a transaction and context. See SelectKeyed.
+func (tx *Tx) SelectKeyedContext(ctx context.Context, dest interface{}, key string, query string, args ...interface{}) error {
+	return SelectKeyedContext(ctx, tx, dest, key, query, args...)
+}
+
 // GetContext within a transaction and context.
 // Any placeholder parameters are replaced with supplied args.
 // An error is returned if the result set is empty.
@@ -365,6 +390,11 @@ func (s *Stmt) SelectContext(ctx context.Context, dest interface{}, args ...inte
 	return SelectContext(ctx, &qStmt{s}, dest, "", args...)
 }
 
+// SelectKeyedContext using the prepared statement. See SelectKeyed.
+func (s *Stmt) SelectKeyedContext(ctx context.Context, dest interface{}, key string, args ...interface{}) error {
+	return SelectKeyedContext(ctx, &qStmt{s}, dest, key, "", args...)
+}
+
 // GetContext using the prepared statement.
 // Any placeholder parameters are replaced with supplied args.
 // An error is returned if the result set is empty.

sqlx_test.go

@@ -1923,3 +1923,74 @@ func TestSelectReset(t *testing.T) {
 		}
 	})
 }
+
+func TestSelectKeyed(t *testing.T) {
+	RunWithSchema(defaultSchema, t, func(db *DB, t *testing.T, now string) {
+		loadDefaultFixture(db, t)
+
+		// map[string]Person keyed by email
+		var byEmail map[string]Person
+		err := db.SelectKeyed(&byEmail, "email", "SELECT first_name, last_name, email FROM person")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(byEmail) != 2 {
+			t.Fatalf("expected 2 people, got %d", len(byEmail))
+		}
+		if p, ok := byEmail["jmoiron@jmoiron.net"]; !ok || p.FirstName != "Jason" {
+			t.Fatalf("Jason by email: ok=%v person=%+v", ok, p)
+		}
+		if p, ok := byEmail["johndoeDNE@gmail.net"]; !ok || p.FirstName != "John" {
+			t.Fatalf("John by email: ok=%v person=%+v", ok, p)
+		}
+
+		// map[string]*Person
+		var byEmailPtr map[string]*Person
+		err = db.SelectKeyed(&byEmailPtr, "email", "SELECT first_name, last_name, email FROM person")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if p := byEmailPtr["jmoiron@jmoiron.net"]; p == nil || p.LastName != "Moiron" {
+			t.Fatalf("ptr map: %+v", p)
+		}
+
+		// map[int]Place keyed by telcode
+		var byCode map[int]Place
+		err = db.SelectKeyed(&byCode, "telcode", "SELECT country, city, telcode FROM place")
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(byCode) != 3 {
+			t.Fatalf("expected 3 places, got %d", len(byCode))
+		}
+		if byCode[1].Country != "United States" {
+			t.Fatalf("telcode 1: %+v", byCode[1])
+		}
+
+		// missing key field
+		err = db.SelectKeyed(&byEmail, "no_such_col", "SELECT first_name, last_name, email FROM person")
+		if err == nil {
+			t.Fatal("expected error for missing key field")
+		}
+
+		// empty key name
+		err = db.SelectKeyed(&byEmail, "", "SELECT first_name FROM person")
+		if err == nil {
+			t.Fatal("expected error for empty key")
+		}
+
+		// non-map dest
+		var people []Person
+		err = db.SelectKeyed(&people, "email", "SELECT first_name, last_name, email FROM person")
+		if err == nil {
+			t.Fatal("expected error for non-map dest")
+		}
+
+		// package-level helper
+		var again map[string]Person
+		err = SelectKeyed(db, &again, "email", "SELECT first_name, last_name, email FROM person")
+		if err != nil || len(again) != 2 {
+			t.Fatalf("package SelectKeyed: err=%v len=%d", err, len(again))
+		}
+	})
+}