index.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899
  1. package memdb
  2. import (
  3. "encoding/binary"
  4. "encoding/hex"
  5. "errors"
  6. "fmt"
  7. "math/bits"
  8. "reflect"
  9. "strings"
  10. )
  11. // Indexer is an interface used for defining indexes. Indexes are used
  12. // for efficient lookup of objects in a MemDB table. An Indexer must also
  13. // implement one of SingleIndexer or MultiIndexer.
  14. //
  15. // Indexers are primarily responsible for returning the lookup key as
  16. // a byte slice. The byte slice is the key data in the underlying data storage.
  17. type Indexer interface {
  18. // FromArgs is called to build the exact index key from a list of arguments.
  19. FromArgs(args ...interface{}) ([]byte, error)
  20. }
  21. // SingleIndexer is an interface used for defining indexes that generate a
  22. // single value per object
  23. type SingleIndexer interface {
  24. // FromObject extracts the index value from an object. The return values
  25. // are whether the index value was found, the index value, and any error
  26. // while extracting the index value, respectively.
  27. FromObject(raw interface{}) (bool, []byte, error)
  28. }
  29. // MultiIndexer is an interface used for defining indexes that generate
  30. // multiple values per object. Each value is stored as a seperate index
  31. // pointing to the same object.
  32. //
  33. // For example, an index that extracts the first and last name of a person
  34. // and allows lookup based on eitherd would be a MultiIndexer. The FromObject
  35. // of this example would split the first and last name and return both as
  36. // values.
  37. type MultiIndexer interface {
  38. // FromObject extracts index values from an object. The return values
  39. // are the same as a SingleIndexer except there can be multiple index
  40. // values.
  41. FromObject(raw interface{}) (bool, [][]byte, error)
  42. }
  43. // PrefixIndexer is an optional interface on top of an Indexer that allows
  44. // indexes to support prefix-based iteration.
  45. type PrefixIndexer interface {
  46. // PrefixFromArgs is the same as FromArgs for an Indexer except that
  47. // the index value returned should return all prefix-matched values.
  48. PrefixFromArgs(args ...interface{}) ([]byte, error)
  49. }
  50. // StringFieldIndex is used to extract a field from an object
  51. // using reflection and builds an index on that field.
  52. type StringFieldIndex struct {
  53. Field string
  54. Lowercase bool
  55. }
  56. func (s *StringFieldIndex) FromObject(obj interface{}) (bool, []byte, error) {
  57. v := reflect.ValueOf(obj)
  58. v = reflect.Indirect(v) // Dereference the pointer if any
  59. fv := v.FieldByName(s.Field)
  60. isPtr := fv.Kind() == reflect.Ptr
  61. fv = reflect.Indirect(fv)
  62. if !isPtr && !fv.IsValid() {
  63. return false, nil,
  64. fmt.Errorf("field '%s' for %#v is invalid %v ", s.Field, obj, isPtr)
  65. }
  66. if isPtr && !fv.IsValid() {
  67. val := ""
  68. return false, []byte(val), nil
  69. }
  70. val := fv.String()
  71. if val == "" {
  72. return false, nil, nil
  73. }
  74. if s.Lowercase {
  75. val = strings.ToLower(val)
  76. }
  77. // Add the null character as a terminator
  78. val += "\x00"
  79. return true, []byte(val), nil
  80. }
  81. func (s *StringFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  82. if len(args) != 1 {
  83. return nil, fmt.Errorf("must provide only a single argument")
  84. }
  85. arg, ok := args[0].(string)
  86. if !ok {
  87. return nil, fmt.Errorf("argument must be a string: %#v", args[0])
  88. }
  89. if s.Lowercase {
  90. arg = strings.ToLower(arg)
  91. }
  92. // Add the null character as a terminator
  93. arg += "\x00"
  94. return []byte(arg), nil
  95. }
  96. func (s *StringFieldIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  97. val, err := s.FromArgs(args...)
  98. if err != nil {
  99. return nil, err
  100. }
  101. // Strip the null terminator, the rest is a prefix
  102. n := len(val)
  103. if n > 0 {
  104. return val[:n-1], nil
  105. }
  106. return val, nil
  107. }
  108. // StringSliceFieldIndex builds an index from a field on an object that is a
  109. // string slice ([]string). Each value within the string slice can be used for
  110. // lookup.
  111. type StringSliceFieldIndex struct {
  112. Field string
  113. Lowercase bool
  114. }
  115. func (s *StringSliceFieldIndex) FromObject(obj interface{}) (bool, [][]byte, error) {
  116. v := reflect.ValueOf(obj)
  117. v = reflect.Indirect(v) // Dereference the pointer if any
  118. fv := v.FieldByName(s.Field)
  119. if !fv.IsValid() {
  120. return false, nil,
  121. fmt.Errorf("field '%s' for %#v is invalid", s.Field, obj)
  122. }
  123. if fv.Kind() != reflect.Slice || fv.Type().Elem().Kind() != reflect.String {
  124. return false, nil, fmt.Errorf("field '%s' is not a string slice", s.Field)
  125. }
  126. length := fv.Len()
  127. vals := make([][]byte, 0, length)
  128. for i := 0; i < fv.Len(); i++ {
  129. val := fv.Index(i).String()
  130. if val == "" {
  131. continue
  132. }
  133. if s.Lowercase {
  134. val = strings.ToLower(val)
  135. }
  136. // Add the null character as a terminator
  137. val += "\x00"
  138. vals = append(vals, []byte(val))
  139. }
  140. if len(vals) == 0 {
  141. return false, nil, nil
  142. }
  143. return true, vals, nil
  144. }
  145. func (s *StringSliceFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  146. if len(args) != 1 {
  147. return nil, fmt.Errorf("must provide only a single argument")
  148. }
  149. arg, ok := args[0].(string)
  150. if !ok {
  151. return nil, fmt.Errorf("argument must be a string: %#v", args[0])
  152. }
  153. if s.Lowercase {
  154. arg = strings.ToLower(arg)
  155. }
  156. // Add the null character as a terminator
  157. arg += "\x00"
  158. return []byte(arg), nil
  159. }
  160. func (s *StringSliceFieldIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  161. val, err := s.FromArgs(args...)
  162. if err != nil {
  163. return nil, err
  164. }
  165. // Strip the null terminator, the rest is a prefix
  166. n := len(val)
  167. if n > 0 {
  168. return val[:n-1], nil
  169. }
  170. return val, nil
  171. }
  172. // StringMapFieldIndex is used to extract a field of type map[string]string
  173. // from an object using reflection and builds an index on that field.
  174. //
  175. // Note that although FromArgs in theory supports using either one or
  176. // two arguments, there is a bug: FromObject only creates an index
  177. // using key/value, and does not also create an index using key. This
  178. // means a lookup using one argument will never actually work.
  179. //
  180. // It is currently left as-is to prevent backwards compatibility
  181. // issues.
  182. //
  183. // TODO: Fix this in the next major bump.
  184. type StringMapFieldIndex struct {
  185. Field string
  186. Lowercase bool
  187. }
  188. var MapType = reflect.MapOf(reflect.TypeOf(""), reflect.TypeOf("")).Kind()
  189. func (s *StringMapFieldIndex) FromObject(obj interface{}) (bool, [][]byte, error) {
  190. v := reflect.ValueOf(obj)
  191. v = reflect.Indirect(v) // Dereference the pointer if any
  192. fv := v.FieldByName(s.Field)
  193. if !fv.IsValid() {
  194. return false, nil, fmt.Errorf("field '%s' for %#v is invalid", s.Field, obj)
  195. }
  196. if fv.Kind() != MapType {
  197. return false, nil, fmt.Errorf("field '%s' is not a map[string]string", s.Field)
  198. }
  199. length := fv.Len()
  200. vals := make([][]byte, 0, length)
  201. for _, key := range fv.MapKeys() {
  202. k := key.String()
  203. if k == "" {
  204. continue
  205. }
  206. val := fv.MapIndex(key).String()
  207. if s.Lowercase {
  208. k = strings.ToLower(k)
  209. val = strings.ToLower(val)
  210. }
  211. // Add the null character as a terminator
  212. k += "\x00" + val + "\x00"
  213. vals = append(vals, []byte(k))
  214. }
  215. if len(vals) == 0 {
  216. return false, nil, nil
  217. }
  218. return true, vals, nil
  219. }
  220. // WARNING: Because of a bug in FromObject, this function will never return
  221. // a value when using the single-argument version.
  222. func (s *StringMapFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  223. if len(args) > 2 || len(args) == 0 {
  224. return nil, fmt.Errorf("must provide one or two arguments")
  225. }
  226. key, ok := args[0].(string)
  227. if !ok {
  228. return nil, fmt.Errorf("argument must be a string: %#v", args[0])
  229. }
  230. if s.Lowercase {
  231. key = strings.ToLower(key)
  232. }
  233. // Add the null character as a terminator
  234. key += "\x00"
  235. if len(args) == 2 {
  236. val, ok := args[1].(string)
  237. if !ok {
  238. return nil, fmt.Errorf("argument must be a string: %#v", args[1])
  239. }
  240. if s.Lowercase {
  241. val = strings.ToLower(val)
  242. }
  243. // Add the null character as a terminator
  244. key += val + "\x00"
  245. }
  246. return []byte(key), nil
  247. }
  248. // IntFieldIndex is used to extract an int field from an object using
  249. // reflection and builds an index on that field.
  250. type IntFieldIndex struct {
  251. Field string
  252. }
  253. func (i *IntFieldIndex) FromObject(obj interface{}) (bool, []byte, error) {
  254. v := reflect.ValueOf(obj)
  255. v = reflect.Indirect(v) // Dereference the pointer if any
  256. fv := v.FieldByName(i.Field)
  257. if !fv.IsValid() {
  258. return false, nil,
  259. fmt.Errorf("field '%s' for %#v is invalid", i.Field, obj)
  260. }
  261. // Check the type
  262. k := fv.Kind()
  263. size, ok := IsIntType(k)
  264. if !ok {
  265. return false, nil, fmt.Errorf("field %q is of type %v; want an int", i.Field, k)
  266. }
  267. // Get the value and encode it
  268. val := fv.Int()
  269. buf := make([]byte, size)
  270. binary.PutVarint(buf, val)
  271. return true, buf, nil
  272. }
  273. func (i *IntFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  274. if len(args) != 1 {
  275. return nil, fmt.Errorf("must provide only a single argument")
  276. }
  277. v := reflect.ValueOf(args[0])
  278. if !v.IsValid() {
  279. return nil, fmt.Errorf("%#v is invalid", args[0])
  280. }
  281. k := v.Kind()
  282. size, ok := IsIntType(k)
  283. if !ok {
  284. return nil, fmt.Errorf("arg is of type %v; want a int", k)
  285. }
  286. val := v.Int()
  287. buf := make([]byte, size)
  288. binary.PutVarint(buf, val)
  289. return buf, nil
  290. }
  291. // IsIntType returns whether the passed type is a type of int and the number
  292. // of bytes needed to encode the type.
  293. func IsIntType(k reflect.Kind) (size int, okay bool) {
  294. switch k {
  295. case reflect.Int:
  296. return binary.MaxVarintLen64, true
  297. case reflect.Int8:
  298. return 2, true
  299. case reflect.Int16:
  300. return binary.MaxVarintLen16, true
  301. case reflect.Int32:
  302. return binary.MaxVarintLen32, true
  303. case reflect.Int64:
  304. return binary.MaxVarintLen64, true
  305. default:
  306. return 0, false
  307. }
  308. }
  309. // UintFieldIndex is used to extract a uint field from an object using
  310. // reflection and builds an index on that field.
  311. type UintFieldIndex struct {
  312. Field string
  313. }
  314. func (u *UintFieldIndex) FromObject(obj interface{}) (bool, []byte, error) {
  315. v := reflect.ValueOf(obj)
  316. v = reflect.Indirect(v) // Dereference the pointer if any
  317. fv := v.FieldByName(u.Field)
  318. if !fv.IsValid() {
  319. return false, nil,
  320. fmt.Errorf("field '%s' for %#v is invalid", u.Field, obj)
  321. }
  322. // Check the type
  323. k := fv.Kind()
  324. size, ok := IsUintType(k)
  325. if !ok {
  326. return false, nil, fmt.Errorf("field %q is of type %v; want a uint", u.Field, k)
  327. }
  328. // Get the value and encode it
  329. val := fv.Uint()
  330. buf := encodeUInt(val, size)
  331. return true, buf, nil
  332. }
  333. func (u *UintFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  334. if len(args) != 1 {
  335. return nil, fmt.Errorf("must provide only a single argument")
  336. }
  337. v := reflect.ValueOf(args[0])
  338. if !v.IsValid() {
  339. return nil, fmt.Errorf("%#v is invalid", args[0])
  340. }
  341. k := v.Kind()
  342. size, ok := IsUintType(k)
  343. if !ok {
  344. return nil, fmt.Errorf("arg is of type %v; want a uint", k)
  345. }
  346. val := v.Uint()
  347. buf := encodeUInt(val, size)
  348. return buf, nil
  349. }
  350. func encodeUInt(val uint64, size int) []byte {
  351. buf := make([]byte, size)
  352. switch size {
  353. case 1:
  354. buf[0] = uint8(val)
  355. case 2:
  356. binary.BigEndian.PutUint16(buf, uint16(val))
  357. case 4:
  358. binary.BigEndian.PutUint32(buf, uint32(val))
  359. case 8:
  360. binary.BigEndian.PutUint64(buf, val)
  361. }
  362. return buf
  363. }
  364. // IsUintType returns whether the passed type is a type of uint and the number
  365. // of bytes needed to encode the type.
  366. func IsUintType(k reflect.Kind) (size int, okay bool) {
  367. switch k {
  368. case reflect.Uint:
  369. return bits.UintSize / 8, true
  370. case reflect.Uint8:
  371. return 1, true
  372. case reflect.Uint16:
  373. return 2, true
  374. case reflect.Uint32:
  375. return 4, true
  376. case reflect.Uint64:
  377. return 8, true
  378. default:
  379. return 0, false
  380. }
  381. }
  382. // BoolFieldIndex is used to extract an boolean field from an object using
  383. // reflection and builds an index on that field.
  384. type BoolFieldIndex struct {
  385. Field string
  386. }
  387. func (i *BoolFieldIndex) FromObject(obj interface{}) (bool, []byte, error) {
  388. v := reflect.ValueOf(obj)
  389. v = reflect.Indirect(v) // Dereference the pointer if any
  390. fv := v.FieldByName(i.Field)
  391. if !fv.IsValid() {
  392. return false, nil,
  393. fmt.Errorf("field '%s' for %#v is invalid", i.Field, obj)
  394. }
  395. // Check the type
  396. k := fv.Kind()
  397. if k != reflect.Bool {
  398. return false, nil, fmt.Errorf("field %q is of type %v; want a bool", i.Field, k)
  399. }
  400. // Get the value and encode it
  401. buf := make([]byte, 1)
  402. if fv.Bool() {
  403. buf[0] = 1
  404. }
  405. return true, buf, nil
  406. }
  407. func (i *BoolFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  408. return fromBoolArgs(args)
  409. }
  410. // UUIDFieldIndex is used to extract a field from an object
  411. // using reflection and builds an index on that field by treating
  412. // it as a UUID. This is an optimization to using a StringFieldIndex
  413. // as the UUID can be more compactly represented in byte form.
  414. type UUIDFieldIndex struct {
  415. Field string
  416. }
  417. func (u *UUIDFieldIndex) FromObject(obj interface{}) (bool, []byte, error) {
  418. v := reflect.ValueOf(obj)
  419. v = reflect.Indirect(v) // Dereference the pointer if any
  420. fv := v.FieldByName(u.Field)
  421. if !fv.IsValid() {
  422. return false, nil,
  423. fmt.Errorf("field '%s' for %#v is invalid", u.Field, obj)
  424. }
  425. val := fv.String()
  426. if val == "" {
  427. return false, nil, nil
  428. }
  429. buf, err := u.parseString(val, true)
  430. return true, buf, err
  431. }
  432. func (u *UUIDFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  433. if len(args) != 1 {
  434. return nil, fmt.Errorf("must provide only a single argument")
  435. }
  436. switch arg := args[0].(type) {
  437. case string:
  438. return u.parseString(arg, true)
  439. case []byte:
  440. if len(arg) != 16 {
  441. return nil, fmt.Errorf("byte slice must be 16 characters")
  442. }
  443. return arg, nil
  444. default:
  445. return nil,
  446. fmt.Errorf("argument must be a string or byte slice: %#v", args[0])
  447. }
  448. }
  449. func (u *UUIDFieldIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  450. if len(args) != 1 {
  451. return nil, fmt.Errorf("must provide only a single argument")
  452. }
  453. switch arg := args[0].(type) {
  454. case string:
  455. return u.parseString(arg, false)
  456. case []byte:
  457. return arg, nil
  458. default:
  459. return nil,
  460. fmt.Errorf("argument must be a string or byte slice: %#v", args[0])
  461. }
  462. }
  463. // parseString parses a UUID from the string. If enforceLength is false, it will
  464. // parse a partial UUID. An error is returned if the input, stripped of hyphens,
  465. // is not even length.
  466. func (u *UUIDFieldIndex) parseString(s string, enforceLength bool) ([]byte, error) {
  467. // Verify the length
  468. l := len(s)
  469. if enforceLength && l != 36 {
  470. return nil, fmt.Errorf("UUID must be 36 characters")
  471. } else if l > 36 {
  472. return nil, fmt.Errorf("Invalid UUID length. UUID have 36 characters; got %d", l)
  473. }
  474. hyphens := strings.Count(s, "-")
  475. if hyphens > 4 {
  476. return nil, fmt.Errorf(`UUID should have maximum of 4 "-"; got %d`, hyphens)
  477. }
  478. // The sanitized length is the length of the original string without the "-".
  479. sanitized := strings.Replace(s, "-", "", -1)
  480. sanitizedLength := len(sanitized)
  481. if sanitizedLength%2 != 0 {
  482. return nil, fmt.Errorf("Input (without hyphens) must be even length")
  483. }
  484. dec, err := hex.DecodeString(sanitized)
  485. if err != nil {
  486. return nil, fmt.Errorf("Invalid UUID: %v", err)
  487. }
  488. return dec, nil
  489. }
  490. // FieldSetIndex is used to extract a field from an object using reflection and
  491. // builds an index on whether the field is set by comparing it against its
  492. // type's nil value.
  493. type FieldSetIndex struct {
  494. Field string
  495. }
  496. func (f *FieldSetIndex) FromObject(obj interface{}) (bool, []byte, error) {
  497. v := reflect.ValueOf(obj)
  498. v = reflect.Indirect(v) // Dereference the pointer if any
  499. fv := v.FieldByName(f.Field)
  500. if !fv.IsValid() {
  501. return false, nil,
  502. fmt.Errorf("field '%s' for %#v is invalid", f.Field, obj)
  503. }
  504. if fv.Interface() == reflect.Zero(fv.Type()).Interface() {
  505. return true, []byte{0}, nil
  506. }
  507. return true, []byte{1}, nil
  508. }
  509. func (f *FieldSetIndex) FromArgs(args ...interface{}) ([]byte, error) {
  510. return fromBoolArgs(args)
  511. }
  512. // ConditionalIndex builds an index based on a condition specified by a passed
  513. // user function. This function may examine the passed object and return a
  514. // boolean to encapsulate an arbitrarily complex conditional.
  515. type ConditionalIndex struct {
  516. Conditional ConditionalIndexFunc
  517. }
  518. // ConditionalIndexFunc is the required function interface for a
  519. // ConditionalIndex.
  520. type ConditionalIndexFunc func(obj interface{}) (bool, error)
  521. func (c *ConditionalIndex) FromObject(obj interface{}) (bool, []byte, error) {
  522. // Call the user's function
  523. res, err := c.Conditional(obj)
  524. if err != nil {
  525. return false, nil, fmt.Errorf("ConditionalIndexFunc(%#v) failed: %v", obj, err)
  526. }
  527. if res {
  528. return true, []byte{1}, nil
  529. }
  530. return true, []byte{0}, nil
  531. }
  532. func (c *ConditionalIndex) FromArgs(args ...interface{}) ([]byte, error) {
  533. return fromBoolArgs(args)
  534. }
  535. // fromBoolArgs is a helper that expects only a single boolean argument and
  536. // returns a single length byte array containing either a one or zero depending
  537. // on whether the passed input is true or false respectively.
  538. func fromBoolArgs(args []interface{}) ([]byte, error) {
  539. if len(args) != 1 {
  540. return nil, fmt.Errorf("must provide only a single argument")
  541. }
  542. if val, ok := args[0].(bool); !ok {
  543. return nil, fmt.Errorf("argument must be a boolean type: %#v", args[0])
  544. } else if val {
  545. return []byte{1}, nil
  546. }
  547. return []byte{0}, nil
  548. }
  549. // CompoundIndex is used to build an index using multiple sub-indexes
  550. // Prefix based iteration is supported as long as the appropriate prefix
  551. // of indexers support it. All sub-indexers are only assumed to expect
  552. // a single argument.
  553. type CompoundIndex struct {
  554. Indexes []Indexer
  555. // AllowMissing results in an index based on only the indexers
  556. // that return data. If true, you may end up with 2/3 columns
  557. // indexed which might be useful for an index scan. Otherwise,
  558. // the CompoundIndex requires all indexers to be satisfied.
  559. AllowMissing bool
  560. }
  561. func (c *CompoundIndex) FromObject(raw interface{}) (bool, []byte, error) {
  562. var out []byte
  563. for i, idxRaw := range c.Indexes {
  564. idx, ok := idxRaw.(SingleIndexer)
  565. if !ok {
  566. return false, nil, fmt.Errorf("sub-index %d error: %s", i, "sub-index must be a SingleIndexer")
  567. }
  568. ok, val, err := idx.FromObject(raw)
  569. if err != nil {
  570. return false, nil, fmt.Errorf("sub-index %d error: %v", i, err)
  571. }
  572. if !ok {
  573. if c.AllowMissing {
  574. break
  575. } else {
  576. return false, nil, nil
  577. }
  578. }
  579. out = append(out, val...)
  580. }
  581. return true, out, nil
  582. }
  583. func (c *CompoundIndex) FromArgs(args ...interface{}) ([]byte, error) {
  584. if len(args) != len(c.Indexes) {
  585. return nil, fmt.Errorf("non-equivalent argument count and index fields")
  586. }
  587. var out []byte
  588. for i, arg := range args {
  589. val, err := c.Indexes[i].FromArgs(arg)
  590. if err != nil {
  591. return nil, fmt.Errorf("sub-index %d error: %v", i, err)
  592. }
  593. out = append(out, val...)
  594. }
  595. return out, nil
  596. }
  597. func (c *CompoundIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  598. if len(args) > len(c.Indexes) {
  599. return nil, fmt.Errorf("more arguments than index fields")
  600. }
  601. var out []byte
  602. for i, arg := range args {
  603. if i+1 < len(args) {
  604. val, err := c.Indexes[i].FromArgs(arg)
  605. if err != nil {
  606. return nil, fmt.Errorf("sub-index %d error: %v", i, err)
  607. }
  608. out = append(out, val...)
  609. } else {
  610. prefixIndexer, ok := c.Indexes[i].(PrefixIndexer)
  611. if !ok {
  612. return nil, fmt.Errorf("sub-index %d does not support prefix scanning", i)
  613. }
  614. val, err := prefixIndexer.PrefixFromArgs(arg)
  615. if err != nil {
  616. return nil, fmt.Errorf("sub-index %d error: %v", i, err)
  617. }
  618. out = append(out, val...)
  619. }
  620. }
  621. return out, nil
  622. }
  623. // CompoundMultiIndex is used to build an index using multiple
  624. // sub-indexes.
  625. //
  626. // Unlike CompoundIndex, CompoundMultiIndex can have both
  627. // SingleIndexer and MultiIndexer sub-indexers. However, each
  628. // MultiIndexer adds considerable overhead/complexity in terms of
  629. // the number of indexes created under-the-hood. It is not suggested
  630. // to use more than one or two, if possible.
  631. //
  632. // Another change from CompoundIndexer is that if AllowMissing is
  633. // set, not only is it valid to have empty index fields, but it will
  634. // still create index values up to the first empty index. This means
  635. // that if you have a value with an empty field, rather than using a
  636. // prefix for lookup, you can simply pass in less arguments. As an
  637. // example, if {Foo, Bar} is indexed but Bar is missing for a value
  638. // and AllowMissing is set, an index will still be created for {Foo}
  639. // and it is valid to do a lookup passing in only Foo as an argument.
  640. // Note that the ordering isn't guaranteed -- it's last-insert wins,
  641. // but this is true if you have two objects that have the same
  642. // indexes not using AllowMissing anyways.
  643. //
  644. // Because StringMapFieldIndexers can take a varying number of args,
  645. // it is currently a requirement that whenever it is used, two
  646. // arguments must _always_ be provided for it. In theory we only
  647. // need one, except a bug in that indexer means the single-argument
  648. // version will never work. You can leave the second argument nil,
  649. // but it will never produce a value. We support this for whenever
  650. // that bug is fixed, likely in a next major version bump.
  651. //
  652. // Prefix-based indexing is not currently supported.
  653. type CompoundMultiIndex struct {
  654. Indexes []Indexer
  655. // AllowMissing results in an index based on only the indexers
  656. // that return data. If true, you may end up with 2/3 columns
  657. // indexed which might be useful for an index scan. Otherwise,
  658. // CompoundMultiIndex requires all indexers to be satisfied.
  659. AllowMissing bool
  660. }
  661. func (c *CompoundMultiIndex) FromObject(raw interface{}) (bool, [][]byte, error) {
  662. // At each entry, builder is storing the results from the next index
  663. builder := make([][][]byte, 0, len(c.Indexes))
  664. // Start with something higher to avoid resizing if possible
  665. out := make([][]byte, 0, len(c.Indexes)^3)
  666. forloop:
  667. // This loop goes through each indexer and adds the value(s) provided to the next
  668. // entry in the slice. We can then later walk it like a tree to construct the indices.
  669. for i, idxRaw := range c.Indexes {
  670. switch idx := idxRaw.(type) {
  671. case SingleIndexer:
  672. ok, val, err := idx.FromObject(raw)
  673. if err != nil {
  674. return false, nil, fmt.Errorf("single sub-index %d error: %v", i, err)
  675. }
  676. if !ok {
  677. if c.AllowMissing {
  678. break forloop
  679. } else {
  680. return false, nil, nil
  681. }
  682. }
  683. builder = append(builder, [][]byte{val})
  684. case MultiIndexer:
  685. ok, vals, err := idx.FromObject(raw)
  686. if err != nil {
  687. return false, nil, fmt.Errorf("multi sub-index %d error: %v", i, err)
  688. }
  689. if !ok {
  690. if c.AllowMissing {
  691. break forloop
  692. } else {
  693. return false, nil, nil
  694. }
  695. }
  696. // Add each of the new values to each of the old values
  697. builder = append(builder, vals)
  698. default:
  699. return false, nil, fmt.Errorf("sub-index %d does not satisfy either SingleIndexer or MultiIndexer", i)
  700. }
  701. }
  702. // We are walking through the builder slice essentially in a depth-first fashion,
  703. // building the prefix and leaves as we go. If AllowMissing is false, we only insert
  704. // these full paths to leaves. Otherwise, we also insert each prefix along the way.
  705. // This allows for lookup in FromArgs when AllowMissing is true that does not contain
  706. // the full set of arguments. e.g. for {Foo, Bar} where an object has only the Foo
  707. // field specified as "abc", it is valid to call FromArgs with just "abc".
  708. var walkVals func([]byte, int)
  709. walkVals = func(currPrefix []byte, depth int) {
  710. if depth == len(builder)-1 {
  711. // These are the "leaves", so append directly
  712. for _, v := range builder[depth] {
  713. out = append(out, append(currPrefix, v...))
  714. }
  715. return
  716. }
  717. for _, v := range builder[depth] {
  718. nextPrefix := append(currPrefix, v...)
  719. if c.AllowMissing {
  720. out = append(out, nextPrefix)
  721. }
  722. walkVals(nextPrefix, depth+1)
  723. }
  724. }
  725. walkVals(nil, 0)
  726. return true, out, nil
  727. }
  728. func (c *CompoundMultiIndex) FromArgs(args ...interface{}) ([]byte, error) {
  729. var stringMapCount int
  730. var argCount int
  731. for _, index := range c.Indexes {
  732. if argCount >= len(args) {
  733. break
  734. }
  735. if _, ok := index.(*StringMapFieldIndex); ok {
  736. // We require pairs for StringMapFieldIndex, but only got one
  737. if argCount+1 >= len(args) {
  738. return nil, errors.New("invalid number of arguments")
  739. }
  740. stringMapCount++
  741. argCount += 2
  742. } else {
  743. argCount++
  744. }
  745. }
  746. argCount = 0
  747. switch c.AllowMissing {
  748. case true:
  749. if len(args) > len(c.Indexes)+stringMapCount {
  750. return nil, errors.New("too many arguments")
  751. }
  752. default:
  753. if len(args) != len(c.Indexes)+stringMapCount {
  754. return nil, errors.New("number of arguments does not equal number of indexers")
  755. }
  756. }
  757. var out []byte
  758. var val []byte
  759. var err error
  760. for i, idx := range c.Indexes {
  761. if argCount >= len(args) {
  762. // We're done; should only hit this if AllowMissing
  763. break
  764. }
  765. if _, ok := idx.(*StringMapFieldIndex); ok {
  766. if args[argCount+1] == nil {
  767. val, err = idx.FromArgs(args[argCount])
  768. } else {
  769. val, err = idx.FromArgs(args[argCount : argCount+2]...)
  770. }
  771. argCount += 2
  772. } else {
  773. val, err = idx.FromArgs(args[argCount])
  774. argCount++
  775. }
  776. if err != nil {
  777. return nil, fmt.Errorf("sub-index %d error: %v", i, err)
  778. }
  779. out = append(out, val...)
  780. }
  781. return out, nil
  782. }