helpers.go 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154
  1. package arangodb
  2. import (
  3. "context"
  4. "crypto/md5"
  5. "encoding/binary"
  6. "encoding/hex"
  7. "io"
  8. "strings"
  9. "github.com/arangodb/go-driver"
  10. "github.com/seaweedfs/seaweedfs/weed/util"
  11. )
  12. // convert a string into arango-key safe hex bytes hash
  13. func hashString(dir string) string {
  14. h := md5.New()
  15. io.WriteString(h, dir)
  16. b := h.Sum(nil)
  17. return hex.EncodeToString(b)
  18. }
  19. // convert slice of bytes into slice of uint64
  20. // the first uint64 indicates the length in bytes
  21. func bytesToArray(bs []byte) []uint64 {
  22. out := make([]uint64, 0, 2+len(bs)/8)
  23. out = append(out, uint64(len(bs)))
  24. for len(bs)%8 != 0 {
  25. bs = append(bs, 0)
  26. }
  27. for i := 0; i < len(bs); i = i + 8 {
  28. out = append(out, binary.BigEndian.Uint64(bs[i:]))
  29. }
  30. return out
  31. }
  32. // convert from slice of uint64 back to bytes
  33. // if input length is 0 or 1, will return nil
  34. func arrayToBytes(xs []uint64) []byte {
  35. if len(xs) < 2 {
  36. return nil
  37. }
  38. first := xs[0]
  39. out := make([]byte, len(xs)*8) // i think this can actually be len(xs)*8-8, but i dont think an extra 8 bytes hurts...
  40. for i := 1; i < len(xs); i = i + 1 {
  41. binary.BigEndian.PutUint64(out[((i-1)*8):], xs[i])
  42. }
  43. return out[:first]
  44. }
  45. // gets the collection the bucket points to from filepath
  46. func (store *ArangodbStore) extractBucketCollection(ctx context.Context, fullpath util.FullPath) (c driver.Collection, err error) {
  47. bucket, _ := extractBucket(fullpath)
  48. if bucket == "" {
  49. bucket = DEFAULT_COLLECTION
  50. }
  51. c, err = store.ensureBucket(ctx, bucket)
  52. if err != nil {
  53. return nil, err
  54. }
  55. return c, err
  56. }
  57. // called by extractBucketCollection
  58. func extractBucket(fullpath util.FullPath) (string, string) {
  59. if !strings.HasPrefix(string(fullpath), BUCKET_PREFIX+"/") {
  60. return "", string(fullpath)
  61. }
  62. if strings.Count(string(fullpath), "/") < 3 {
  63. return "", string(fullpath)
  64. }
  65. bucketAndObjectKey := string(fullpath)[len(BUCKET_PREFIX+"/"):]
  66. t := strings.Index(bucketAndObjectKey, "/")
  67. bucket := bucketAndObjectKey
  68. shortPath := "/"
  69. if t > 0 {
  70. bucket = bucketAndObjectKey[:t]
  71. shortPath = string(util.FullPath(bucketAndObjectKey[t:]))
  72. }
  73. return bucket, shortPath
  74. }
  75. // get bucket collection from cache. if not exist, creates the buckets collection and grab it
  76. func (store *ArangodbStore) ensureBucket(ctx context.Context, bucket string) (bc driver.Collection, err error) {
  77. var ok bool
  78. store.mu.RLock()
  79. bc, ok = store.buckets[bucket]
  80. store.mu.RUnlock()
  81. if ok && bc != nil {
  82. return bc, nil
  83. }
  84. store.mu.Lock()
  85. defer store.mu.Unlock()
  86. store.buckets[bucket], err = store.ensureCollection(ctx, bucket)
  87. if err != nil {
  88. return nil, err
  89. }
  90. return store.buckets[bucket], nil
  91. }
  92. // transform to an arango compliant name
  93. func bucketToCollectionName(s string) string {
  94. if len(s) == 0 {
  95. return ""
  96. }
  97. // replace all "." with _
  98. s = strings.ReplaceAll(s, ".", "_")
  99. // if starts with number or '.' then add a special prefix
  100. if (s[0] >= '0' && s[0] <= '9') || (s[0] == '.' || s[0] == '_' || s[0] == '-') {
  101. s = "xN--" + s
  102. }
  103. return s
  104. }
  105. // creates collection if not exist, ensures indices if not exist
  106. func (store *ArangodbStore) ensureCollection(ctx context.Context, bucket_name string) (c driver.Collection, err error) {
  107. // convert the bucket to collection name
  108. name := bucketToCollectionName(bucket_name)
  109. ok, err := store.database.CollectionExists(ctx, name)
  110. if err != nil {
  111. return
  112. }
  113. if ok {
  114. c, err = store.database.Collection(ctx, name)
  115. } else {
  116. c, err = store.database.CreateCollection(ctx, name, &driver.CreateCollectionOptions{})
  117. }
  118. if err != nil {
  119. return
  120. }
  121. // ensure indices
  122. if _, _, err = c.EnsurePersistentIndex(ctx, []string{"directory", "name"},
  123. &driver.EnsurePersistentIndexOptions{
  124. Name: "directory_name_multi", Unique: true,
  125. }); err != nil {
  126. return
  127. }
  128. if _, _, err = c.EnsurePersistentIndex(ctx, []string{"directory"},
  129. &driver.EnsurePersistentIndexOptions{Name: "IDX_directory"}); err != nil {
  130. return
  131. }
  132. if _, _, err = c.EnsureTTLIndex(ctx, "ttl", 1,
  133. &driver.EnsureTTLIndexOptions{Name: "IDX_TTL"}); err != nil {
  134. return
  135. }
  136. if _, _, err = c.EnsurePersistentIndex(ctx, []string{"name"}, &driver.EnsurePersistentIndexOptions{
  137. Name: "IDX_name",
  138. }); err != nil {
  139. return
  140. }
  141. return c, nil
  142. }