...
|
...
|
@@ -2,40 +2,79 @@ package es |
|
|
|
|
|
import (
|
|
|
"context"
|
|
|
"fmt"
|
|
|
"strings"
|
|
|
"sync"
|
|
|
"time"
|
|
|
|
|
|
"github.com/olivere/elastic"
|
|
|
"github.com/tal-tech/go-zero/core/fx"
|
|
|
"github.com/tal-tech/go-zero/core/lang"
|
|
|
"github.com/tal-tech/go-zero/core/logx"
|
|
|
"github.com/tal-tech/go-zero/core/syncx"
|
|
|
)
|
|
|
|
|
|
const sharedCallsKey = "ensureIndex"
|
|
|
const (
|
|
|
sharedCallsKey = "ensureIndex"
|
|
|
timestampFormat = "2006-01-02T15:04:05.000Z"
|
|
|
timestampKey = "@timestamp"
|
|
|
)
|
|
|
|
|
|
const (
|
|
|
stateNormal = iota
|
|
|
stateWrap
|
|
|
stateDot
|
|
|
)
|
|
|
|
|
|
type (
|
|
|
IndexFormat func(time.Time) string
|
|
|
IndexFormat func(m map[string]interface{}) string
|
|
|
IndexFunc func() string
|
|
|
|
|
|
Index struct {
|
|
|
client *elastic.Client
|
|
|
indexFormat IndexFormat
|
|
|
index string
|
|
|
indices map[string]lang.PlaceholderType
|
|
|
lock sync.RWMutex
|
|
|
sharedCalls syncx.SharedCalls
|
|
|
}
|
|
|
)
|
|
|
|
|
|
func NewIndex(client *elastic.Client, indexFormat IndexFormat) *Index {
|
|
|
func NewIndex(client *elastic.Client, indexFormat string, loc *time.Location) *Index {
|
|
|
var formatter func(map[string]interface{}) string
|
|
|
format, attrs := getFormat(indexFormat)
|
|
|
if len(attrs) > 0 {
|
|
|
formatter = func(m map[string]interface{}) string {
|
|
|
var vals []interface{}
|
|
|
for _, attr := range attrs {
|
|
|
if val, ok := m[attr]; ok {
|
|
|
vals = append(vals, val)
|
|
|
}
|
|
|
}
|
|
|
return getTime(m).In(loc).Format(fmt.Sprintf(format, vals...))
|
|
|
}
|
|
|
} else {
|
|
|
formatter = func(m map[string]interface{}) string {
|
|
|
return getTime(m).In(loc).Format(format)
|
|
|
}
|
|
|
}
|
|
|
|
|
|
return &Index{
|
|
|
client: client,
|
|
|
indexFormat: indexFormat,
|
|
|
indexFormat: formatter,
|
|
|
indices: make(map[string]lang.PlaceholderType),
|
|
|
sharedCalls: syncx.NewSharedCalls(),
|
|
|
}
|
|
|
}
|
|
|
|
|
|
func (idx *Index) GetIndex(t time.Time) string {
|
|
|
index := idx.indexFormat(t)
|
|
|
func (idx *Index) GetIndex(m map[string]interface{}) string {
|
|
|
index := idx.indexFormat(m)
|
|
|
idx.lock.RLock()
|
|
|
if _, ok := idx.indices[index]; ok {
|
|
|
idx.lock.RUnlock()
|
|
|
return index
|
|
|
}
|
|
|
|
|
|
idx.lock.RUnlock()
|
|
|
if err := idx.ensureIndex(index); err != nil {
|
|
|
logx.Error(err)
|
|
|
}
|
...
|
...
|
@@ -43,17 +82,14 @@ func (idx *Index) GetIndex(t time.Time) string { |
|
|
}
|
|
|
|
|
|
func (idx *Index) ensureIndex(index string) error {
|
|
|
idx.lock.RLock()
|
|
|
if index == idx.index {
|
|
|
idx.lock.RUnlock()
|
|
|
return nil
|
|
|
}
|
|
|
idx.lock.RUnlock()
|
|
|
|
|
|
_, err := idx.sharedCalls.Do(sharedCallsKey, func() (i interface{}, err error) {
|
|
|
idx.lock.Lock()
|
|
|
defer idx.lock.Unlock()
|
|
|
|
|
|
if _, ok := idx.indices[index]; ok {
|
|
|
return nil, nil
|
|
|
}
|
|
|
|
|
|
existsService := elastic.NewIndicesExistsService(idx.client)
|
|
|
existsService.Index([]string{index})
|
|
|
exist, err := existsService.Do(context.Background())
|
...
|
...
|
@@ -61,7 +97,6 @@ func (idx *Index) ensureIndex(index string) error { |
|
|
return nil, err
|
|
|
}
|
|
|
if exist {
|
|
|
idx.index = index
|
|
|
return nil, nil
|
|
|
}
|
|
|
|
...
|
...
|
@@ -74,8 +109,53 @@ func (idx *Index) ensureIndex(index string) error { |
|
|
return nil, err
|
|
|
}
|
|
|
|
|
|
idx.index = index
|
|
|
idx.indices[index] = lang.Placeholder
|
|
|
return nil, nil
|
|
|
})
|
|
|
return err
|
|
|
}
|
|
|
|
|
|
func getTime(m map[string]interface{}) time.Time {
|
|
|
if ti, ok := m[timestampKey]; ok {
|
|
|
if ts, ok := ti.(string); ok {
|
|
|
if t, err := time.Parse(timestampFormat, ts); err == nil {
|
|
|
return t
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
return time.Now()
|
|
|
}
|
|
|
|
|
|
func getFormat(indexFormat string) (format string, attrs []string) {
|
|
|
var state = stateNormal
|
|
|
var builder strings.Builder
|
|
|
var keyBuf strings.Builder
|
|
|
for _, ch := range indexFormat {
|
|
|
switch ch {
|
|
|
case '{':
|
|
|
state = stateWrap
|
|
|
case '.':
|
|
|
if state == stateWrap {
|
|
|
state = stateDot
|
|
|
} else {
|
|
|
builder.WriteRune(ch)
|
|
|
}
|
|
|
case '}':
|
|
|
state = stateNormal
|
|
|
if keyBuf.Len() > 0 {
|
|
|
attrs = append(attrs, keyBuf.String())
|
|
|
builder.WriteString("%s")
|
|
|
}
|
|
|
default:
|
|
|
if state == stateDot {
|
|
|
keyBuf.WriteRune(ch)
|
|
|
} else {
|
|
|
builder.WriteRune(ch)
|
|
|
}
|
|
|
}
|
|
|
}
|
|
|
|
|
|
format = builder.String()
|
|
|
return
|
|
|
} |
...
|
...
|
|