101 lines
1.7 KiB
Go
101 lines
1.7 KiB
Go
package elasticstream
|
|
|
|
import (
|
|
"fmt"
|
|
"log"
|
|
|
|
"github.com/boltdb/bolt"
|
|
"github.com/elastic/go-elasticsearch/v8"
|
|
)
|
|
|
|
type Client struct {
|
|
es *elasticsearch.Client
|
|
config *Config
|
|
ch chan Data
|
|
db *bolt.DB
|
|
offsets map[string]int
|
|
}
|
|
|
|
func NewClient(config *Config) (*Client, error) {
|
|
client := &Client{
|
|
config: config,
|
|
offsets: make(map[string]int),
|
|
}
|
|
return client, nil
|
|
}
|
|
|
|
func (c *Client) Open() error {
|
|
// Open a connection with ElasticSearch
|
|
cfg := elasticsearch.Config{
|
|
Addresses: []string{
|
|
c.config.Host,
|
|
},
|
|
}
|
|
|
|
var err error
|
|
c.es, err = elasticsearch.NewClient(cfg)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// create a buffer channel
|
|
c.ch = make(chan Data, 25)
|
|
|
|
// open bolt db
|
|
c.db, err = bolt.Open(c.config.DBPath, 0666, nil)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
// defer db.Close()
|
|
|
|
for _, index := range c.config.Indexes {
|
|
offset, err := getOffset(c.db, index)
|
|
if err != nil {
|
|
log.Printf("getOffset(%s) err:%s\n", index, err)
|
|
continue
|
|
}
|
|
|
|
c.offsets[index] = offset
|
|
|
|
NewWorker(c, index, offset, c.config.BatchSize)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (c *Client) Read() (*Data, error) {
|
|
|
|
data, ok := <-c.ch
|
|
if !ok {
|
|
return nil, fmt.Errorf("error reading data from channel")
|
|
}
|
|
|
|
index, ok := data.Header["index"]
|
|
if !ok {
|
|
return nil, fmt.Errorf("error index not found in data header")
|
|
}
|
|
|
|
offset, ok := c.offsets[index]
|
|
if !ok {
|
|
return nil, fmt.Errorf("error offset index not found")
|
|
}
|
|
|
|
c.offsets[index] = offset + 1
|
|
|
|
err := setOffset(c.db, index, c.offsets[index])
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &data, nil
|
|
}
|
|
|
|
// func (c *Client) Ack(ctx context.Context, position int) error {
|
|
// return nil
|
|
// }
|
|
|
|
// close the client
|
|
func (c *Client) Teardown() error {
|
|
return nil
|
|
}
|