index.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419
  1. package memdb
  2. import (
  3. "encoding/hex"
  4. "fmt"
  5. "reflect"
  6. "strings"
  7. )
  8. // Indexer is an interface used for defining indexes
  9. type Indexer interface {
  10. // ExactFromArgs is used to build an exact index lookup
  11. // based on arguments
  12. FromArgs(args ...interface{}) ([]byte, error)
  13. }
  14. // SingleIndexer is an interface used for defining indexes
  15. // generating a single entry per object
  16. type SingleIndexer interface {
  17. // FromObject is used to extract an index value from an
  18. // object or to indicate that the index value is missing.
  19. FromObject(raw interface{}) (bool, []byte, error)
  20. }
  21. // MultiIndexer is an interface used for defining indexes
  22. // generating multiple entries per object
  23. type MultiIndexer interface {
  24. // FromObject is used to extract index values from an
  25. // object or to indicate that the index value is missing.
  26. FromObject(raw interface{}) (bool, [][]byte, error)
  27. }
  28. // PrefixIndexer can optionally be implemented for any
  29. // indexes that support prefix based iteration. This may
  30. // not apply to all indexes.
  31. type PrefixIndexer interface {
  32. // PrefixFromArgs returns a prefix that should be used
  33. // for scanning based on the arguments
  34. PrefixFromArgs(args ...interface{}) ([]byte, error)
  35. }
  36. // StringFieldIndex is used to extract a field from an object
  37. // using reflection and builds an index on that field.
  38. type StringFieldIndex struct {
  39. Field string
  40. Lowercase bool
  41. }
  42. func (s *StringFieldIndex) FromObject(obj interface{}) (bool, []byte, error) {
  43. v := reflect.ValueOf(obj)
  44. v = reflect.Indirect(v) // Dereference the pointer if any
  45. fv := v.FieldByName(s.Field)
  46. if !fv.IsValid() {
  47. return false, nil,
  48. fmt.Errorf("field '%s' for %#v is invalid", s.Field, obj)
  49. }
  50. val := fv.String()
  51. if val == "" {
  52. return false, nil, nil
  53. }
  54. if s.Lowercase {
  55. val = strings.ToLower(val)
  56. }
  57. // Add the null character as a terminator
  58. val += "\x00"
  59. return true, []byte(val), nil
  60. }
  61. func (s *StringFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  62. if len(args) != 1 {
  63. return nil, fmt.Errorf("must provide only a single argument")
  64. }
  65. arg, ok := args[0].(string)
  66. if !ok {
  67. return nil, fmt.Errorf("argument must be a string: %#v", args[0])
  68. }
  69. if s.Lowercase {
  70. arg = strings.ToLower(arg)
  71. }
  72. // Add the null character as a terminator
  73. arg += "\x00"
  74. return []byte(arg), nil
  75. }
  76. func (s *StringFieldIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  77. val, err := s.FromArgs(args...)
  78. if err != nil {
  79. return nil, err
  80. }
  81. // Strip the null terminator, the rest is a prefix
  82. n := len(val)
  83. if n > 0 {
  84. return val[:n-1], nil
  85. }
  86. return val, nil
  87. }
  88. // StringSliceFieldIndex is used to extract a field from an object
  89. // using reflection and builds an index on that field.
  90. type StringSliceFieldIndex struct {
  91. Field string
  92. Lowercase bool
  93. }
  94. func (s *StringSliceFieldIndex) FromObject(obj interface{}) (bool, [][]byte, error) {
  95. v := reflect.ValueOf(obj)
  96. v = reflect.Indirect(v) // Dereference the pointer if any
  97. fv := v.FieldByName(s.Field)
  98. if !fv.IsValid() {
  99. return false, nil,
  100. fmt.Errorf("field '%s' for %#v is invalid", s.Field, obj)
  101. }
  102. if fv.Kind() != reflect.Slice || fv.Type().Elem().Kind() != reflect.String {
  103. return false, nil, fmt.Errorf("field '%s' is not a string slice", s.Field)
  104. }
  105. length := fv.Len()
  106. vals := make([][]byte, 0, length)
  107. for i := 0; i < fv.Len(); i++ {
  108. val := fv.Index(i).String()
  109. if val == "" {
  110. continue
  111. }
  112. if s.Lowercase {
  113. val = strings.ToLower(val)
  114. }
  115. // Add the null character as a terminator
  116. val += "\x00"
  117. vals = append(vals, []byte(val))
  118. }
  119. if len(vals) == 0 {
  120. return false, nil, nil
  121. }
  122. return true, vals, nil
  123. }
  124. func (s *StringSliceFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  125. if len(args) != 1 {
  126. return nil, fmt.Errorf("must provide only a single argument")
  127. }
  128. arg, ok := args[0].(string)
  129. if !ok {
  130. return nil, fmt.Errorf("argument must be a string: %#v", args[0])
  131. }
  132. if s.Lowercase {
  133. arg = strings.ToLower(arg)
  134. }
  135. // Add the null character as a terminator
  136. arg += "\x00"
  137. return []byte(arg), nil
  138. }
  139. func (s *StringSliceFieldIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  140. val, err := s.FromArgs(args...)
  141. if err != nil {
  142. return nil, err
  143. }
  144. // Strip the null terminator, the rest is a prefix
  145. n := len(val)
  146. if n > 0 {
  147. return val[:n-1], nil
  148. }
  149. return val, nil
  150. }
  151. // UUIDFieldIndex is used to extract a field from an object
  152. // using reflection and builds an index on that field by treating
  153. // it as a UUID. This is an optimization to using a StringFieldIndex
  154. // as the UUID can be more compactly represented in byte form.
  155. type UUIDFieldIndex struct {
  156. Field string
  157. }
  158. func (u *UUIDFieldIndex) FromObject(obj interface{}) (bool, []byte, error) {
  159. v := reflect.ValueOf(obj)
  160. v = reflect.Indirect(v) // Dereference the pointer if any
  161. fv := v.FieldByName(u.Field)
  162. if !fv.IsValid() {
  163. return false, nil,
  164. fmt.Errorf("field '%s' for %#v is invalid", u.Field, obj)
  165. }
  166. val := fv.String()
  167. if val == "" {
  168. return false, nil, nil
  169. }
  170. buf, err := u.parseString(val, true)
  171. return true, buf, err
  172. }
  173. func (u *UUIDFieldIndex) FromArgs(args ...interface{}) ([]byte, error) {
  174. if len(args) != 1 {
  175. return nil, fmt.Errorf("must provide only a single argument")
  176. }
  177. switch arg := args[0].(type) {
  178. case string:
  179. return u.parseString(arg, true)
  180. case []byte:
  181. if len(arg) != 16 {
  182. return nil, fmt.Errorf("byte slice must be 16 characters")
  183. }
  184. return arg, nil
  185. default:
  186. return nil,
  187. fmt.Errorf("argument must be a string or byte slice: %#v", args[0])
  188. }
  189. }
  190. func (u *UUIDFieldIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  191. if len(args) != 1 {
  192. return nil, fmt.Errorf("must provide only a single argument")
  193. }
  194. switch arg := args[0].(type) {
  195. case string:
  196. return u.parseString(arg, false)
  197. case []byte:
  198. return arg, nil
  199. default:
  200. return nil,
  201. fmt.Errorf("argument must be a string or byte slice: %#v", args[0])
  202. }
  203. }
  204. // parseString parses a UUID from the string. If enforceLength is false, it will
  205. // parse a partial UUID. An error is returned if the input, stripped of hyphens,
  206. // is not even length.
  207. func (u *UUIDFieldIndex) parseString(s string, enforceLength bool) ([]byte, error) {
  208. // Verify the length
  209. l := len(s)
  210. if enforceLength && l != 36 {
  211. return nil, fmt.Errorf("UUID must be 36 characters")
  212. } else if l > 36 {
  213. return nil, fmt.Errorf("Invalid UUID length. UUID have 36 characters; got %d", l)
  214. }
  215. hyphens := strings.Count(s, "-")
  216. if hyphens > 4 {
  217. return nil, fmt.Errorf(`UUID should have maximum of 4 "-"; got %d`, hyphens)
  218. }
  219. // The sanitized length is the length of the original string without the "-".
  220. sanitized := strings.Replace(s, "-", "", -1)
  221. sanitizedLength := len(sanitized)
  222. if sanitizedLength%2 != 0 {
  223. return nil, fmt.Errorf("Input (without hyphens) must be even length")
  224. }
  225. dec, err := hex.DecodeString(sanitized)
  226. if err != nil {
  227. return nil, fmt.Errorf("Invalid UUID: %v", err)
  228. }
  229. return dec, nil
  230. }
  231. // FieldSetIndex is used to extract a field from an object using reflection and
  232. // builds an index on whether the field is set by comparing it against its
  233. // type's nil value.
  234. type FieldSetIndex struct {
  235. Field string
  236. }
  237. func (f *FieldSetIndex) FromObject(obj interface{}) (bool, []byte, error) {
  238. v := reflect.ValueOf(obj)
  239. v = reflect.Indirect(v) // Dereference the pointer if any
  240. fv := v.FieldByName(f.Field)
  241. if !fv.IsValid() {
  242. return false, nil,
  243. fmt.Errorf("field '%s' for %#v is invalid", f.Field, obj)
  244. }
  245. if fv.Interface() == reflect.Zero(fv.Type()).Interface() {
  246. return true, []byte{0}, nil
  247. }
  248. return true, []byte{1}, nil
  249. }
  250. func (f *FieldSetIndex) FromArgs(args ...interface{}) ([]byte, error) {
  251. return fromBoolArgs(args)
  252. }
  253. // ConditionalIndex builds an index based on a condition specified by a passed
  254. // user function. This function may examine the passed object and return a
  255. // boolean to encapsulate an arbitrarily complex conditional.
  256. type ConditionalIndex struct {
  257. Conditional ConditionalIndexFunc
  258. }
  259. // ConditionalIndexFunc is the required function interface for a
  260. // ConditionalIndex.
  261. type ConditionalIndexFunc func(obj interface{}) (bool, error)
  262. func (c *ConditionalIndex) FromObject(obj interface{}) (bool, []byte, error) {
  263. // Call the user's function
  264. res, err := c.Conditional(obj)
  265. if err != nil {
  266. return false, nil, fmt.Errorf("ConditionalIndexFunc(%#v) failed: %v", obj, err)
  267. }
  268. if res {
  269. return true, []byte{1}, nil
  270. }
  271. return true, []byte{0}, nil
  272. }
  273. func (c *ConditionalIndex) FromArgs(args ...interface{}) ([]byte, error) {
  274. return fromBoolArgs(args)
  275. }
  276. // fromBoolArgs is a helper that expects only a single boolean argument and
  277. // returns a single length byte array containing either a one or zero depending
  278. // on whether the passed input is true or false respectively.
  279. func fromBoolArgs(args []interface{}) ([]byte, error) {
  280. if len(args) != 1 {
  281. return nil, fmt.Errorf("must provide only a single argument")
  282. }
  283. if val, ok := args[0].(bool); !ok {
  284. return nil, fmt.Errorf("argument must be a boolean type: %#v", args[0])
  285. } else if val {
  286. return []byte{1}, nil
  287. }
  288. return []byte{0}, nil
  289. }
  290. // CompoundIndex is used to build an index using multiple sub-indexes
  291. // Prefix based iteration is supported as long as the appropriate prefix
  292. // of indexers support it. All sub-indexers are only assumed to expect
  293. // a single argument.
  294. type CompoundIndex struct {
  295. Indexes []Indexer
  296. // AllowMissing results in an index based on only the indexers
  297. // that return data. If true, you may end up with 2/3 columns
  298. // indexed which might be useful for an index scan. Otherwise,
  299. // the CompoundIndex requires all indexers to be satisfied.
  300. AllowMissing bool
  301. }
  302. func (c *CompoundIndex) FromObject(raw interface{}) (bool, []byte, error) {
  303. var out []byte
  304. for i, idxRaw := range c.Indexes {
  305. idx, ok := idxRaw.(SingleIndexer)
  306. if !ok {
  307. return false, nil, fmt.Errorf("sub-index %d error: %s", i, "sub-index must be a SingleIndexer")
  308. }
  309. ok, val, err := idx.FromObject(raw)
  310. if err != nil {
  311. return false, nil, fmt.Errorf("sub-index %d error: %v", i, err)
  312. }
  313. if !ok {
  314. if c.AllowMissing {
  315. break
  316. } else {
  317. return false, nil, nil
  318. }
  319. }
  320. out = append(out, val...)
  321. }
  322. return true, out, nil
  323. }
  324. func (c *CompoundIndex) FromArgs(args ...interface{}) ([]byte, error) {
  325. if len(args) != len(c.Indexes) {
  326. return nil, fmt.Errorf("less arguments than index fields")
  327. }
  328. var out []byte
  329. for i, arg := range args {
  330. val, err := c.Indexes[i].FromArgs(arg)
  331. if err != nil {
  332. return nil, fmt.Errorf("sub-index %d error: %v", i, err)
  333. }
  334. out = append(out, val...)
  335. }
  336. return out, nil
  337. }
  338. func (c *CompoundIndex) PrefixFromArgs(args ...interface{}) ([]byte, error) {
  339. if len(args) > len(c.Indexes) {
  340. return nil, fmt.Errorf("more arguments than index fields")
  341. }
  342. var out []byte
  343. for i, arg := range args {
  344. if i+1 < len(args) {
  345. val, err := c.Indexes[i].FromArgs(arg)
  346. if err != nil {
  347. return nil, fmt.Errorf("sub-index %d error: %v", i, err)
  348. }
  349. out = append(out, val...)
  350. } else {
  351. prefixIndexer, ok := c.Indexes[i].(PrefixIndexer)
  352. if !ok {
  353. return nil, fmt.Errorf("sub-index %d does not support prefix scanning", i)
  354. }
  355. val, err := prefixIndexer.PrefixFromArgs(arg)
  356. if err != nil {
  357. return nil, fmt.Errorf("sub-index %d error: %v", i, err)
  358. }
  359. out = append(out, val...)
  360. }
  361. }
  362. return out, nil
  363. }