diff --git a/sqlx.go b/sqlx.go index 8259a4feb..b7a538c8e 100644 --- a/sqlx.go +++ b/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) { diff --git a/sqlx_context.go b/sqlx_context.go index 32621d56d..e7f63903d 100644 --- a/sqlx_context.go +++ b/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. diff --git a/sqlx_test.go b/sqlx_test.go index 9fac2cd4f..c471cbcf4 100644 --- a/sqlx_test.go +++ b/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)) + } + }) +}