Skip to content

Commit

Permalink
Almost
Browse files Browse the repository at this point in the history
  • Loading branch information
trzysiek committed Nov 11, 2024
1 parent 56245b8 commit 61e367a
Show file tree
Hide file tree
Showing 9 changed files with 251 additions and 169 deletions.
13 changes: 13 additions & 0 deletions quesma/elasticsearch/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,19 @@ func (es *SimpleClient) RequestWithHeaders(ctx context.Context, method, endpoint
return es.doRequest(ctx, method, endpoint, body, headers)
}

func (es *SimpleClient) DoRequestCheckResponseStatus(ctx context.Context, method, endpoint string, body []byte) (resp *http.Response, err error) {
resp, err = es.doRequest(ctx, "GET", endpoint, body, nil)
if err != nil {
return
}
defer resp.Body.Close()

if resp.StatusCode != http.StatusOK {
return resp, fmt.Errorf("response code from Elastic is not 200 OK, but %s", resp.Status)
}
return resp, nil
}

func (es *SimpleClient) Authenticate(ctx context.Context, authHeader string) bool {
resp, err := es.doRequest(ctx, "GET", "_security/_authenticate", nil, http.Header{"Authorization": {authHeader}})
if err != nil {
Expand Down
221 changes: 88 additions & 133 deletions quesma/persistence/elastic_with_eviction.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ package persistence
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
Expand All @@ -14,8 +15,8 @@ import (
"time"
)

const MAX_DOC_COUNT = 10000 // TODO: fix/make configurable/idk/etc
const defaultSizeInBytesLimit = int64(1_000_000_000) // 1GB
const MAX_DOC_COUNT = 10000 // TODO: fix/make configurable/idk/etc
const defaultHugeSizeInBytesLimit = int64(500_000_000_000) // 500GB

// so far I serialize entire struct and keep only 1 string in ES
type ElasticDatabaseWithEviction struct {
Expand All @@ -34,32 +35,28 @@ func NewElasticDatabaseWithEviction(cfg config.ElasticsearchConfiguration, index
}
}

// mutexy? or what
func (db *ElasticDatabaseWithEviction) Put(ctx context.Context, doc *document) bool {
dbSize, success := db.SizeInBytes()
if !success {
return false
func (db *ElasticDatabaseWithEviction) Put(doc *document) error {
dbSize, err := db.SizeInBytes()
if err != nil {
return err
}
fmt.Println("kk dbg Put() dbSize:", dbSize)
bytesNeeded := dbSize + doc.SizeInBytes
if bytesNeeded > db.SizeInBytesLimit() {
logger.InfoWithCtx(ctx).Msgf("Database is full, need %d bytes more. Evicting documents", bytesNeeded-db.SizeInBytesLimit())
allDocs, ok := db.getAll()
if !ok {
logger.WarnWithCtx(ctx).Msg("Error getting all documents")
return false
logger.Info().Msgf("elastic database: is full, need %d bytes more. Evicting documents", bytesNeeded-db.SizeInBytesLimit())
allDocs, err := db.getAll()
if err != nil {
return err
}
indexesToEvict, bytesEvicted := db.SelectToEvict(allDocs, bytesNeeded-db.SizeInBytesLimit())
logger.InfoWithCtx(ctx).Msgf("Evicting %v indexes, %d bytes", indexesToEvict, bytesEvicted)
logger.Info().Msgf("elastic database: evicting %v indexes, %d bytes", indexesToEvict, bytesEvicted)
db.evict(indexesToEvict)
bytesNeeded -= bytesEvicted
}
if bytesNeeded > db.SizeInBytesLimit() {
// put document
return false
return errors.New("elastic database: is full, cannot put document")
}

//elasticsearchURL := fmt.Sprintf("%s/_update/%s", db.fullIndexName(), doc.Id)
elasticsearchURL := fmt.Sprintf("%s/_update/%s", db.indexName, doc.Id)
fmt.Println("kk dbg Put() elasticsearchURL:", elasticsearchURL)

Expand All @@ -69,40 +66,23 @@ func (db *ElasticDatabaseWithEviction) Put(ctx context.Context, doc *document) b

jsonData, err := json.Marshal(updateContent)
if err != nil {
logger.WarnWithCtx(ctx).Msgf("Error marshalling document: %v", err)
return false
}

resp, err := db.httpClient.Request(context.Background(), "POST", elasticsearchURL, jsonData)
if err != nil {
logger.WarnWithCtx(ctx).Msgf("Error sending request to elastic: %v", err)
return false
return err
}
defer resp.Body.Close()

switch resp.StatusCode {
case http.StatusCreated, http.StatusOK:
return true
default:
respBody, err := io.ReadAll(resp.Body)
if err != nil {
logger.WarnWithCtx(ctx).Msgf("Error reading response body: %v, respBody: %v", err, respBody)
}
return false
resp, err := db.httpClient.DoRequestCheckResponseStatus(context.Background(), http.MethodPost, elasticsearchURL, jsonData)
if err != nil && resp.StatusCode != http.StatusCreated {
return err
}
return nil
}

// co zwraca? zrobić switch na oba typy jakie teraz mamy?
func (db *ElasticDatabaseWithEviction) Get(ctx context.Context, id string) (string, bool) { // probably change return type to *Sizeable
value, success, err := db.ElasticJSONDatabase.Get(id)
if err != nil {
logger.WarnWithCtx(ctx).Msgf("Error getting document, id: %s, error: %v", id, err)
return "", false
}
return value, success
func (db *ElasticDatabaseWithEviction) Get(id string) (string, error) { // probably change return type to *Sizeable
value, _, err := db.ElasticJSONDatabase.Get(id)
return value, err
}

func (db *ElasticDatabaseWithEviction) Delete(id string) bool {
func (db *ElasticDatabaseWithEviction) Delete(id string) error {
// mark as deleted, don't actually delete
// (single document deletion is hard in ES, it's done by evictor for entire index)

Expand All @@ -115,30 +95,21 @@ func (db *ElasticDatabaseWithEviction) Delete(id string) bool {

jsonData, err := json.Marshal(updateContent)
if err != nil {
logger.WarnWithCtx(db.ctx).Msgf("Error marshalling document: %v", err)
return false
return err
}

resp, err := db.httpClient.Request(context.Background(), "POST", elasticsearchURL, jsonData)
if err != nil {
logger.WarnWithCtx(db.ctx).Msgf("Error sending request to elastic: %v", err)
return false
resp, err := db.httpClient.DoRequestCheckResponseStatus(context.Background(), http.MethodPost, elasticsearchURL, jsonData)
if err != nil && resp.StatusCode != http.StatusCreated {
return err
}
defer resp.Body.Close()
return nil
}

switch resp.StatusCode {
case http.StatusCreated, http.StatusOK:
return true
default:
respBody, err := io.ReadAll(resp.Body)
if err != nil {
logger.WarnWithCtx(db.ctx).Msgf("Error reading response body: %v, respBody: %v", err, respBody)
}
return false
}
func (db *ElasticDatabaseWithEviction) DeleteOld(deleteOlderThan time.Duration) error {
return nil
}

func (db *ElasticDatabaseWithEviction) DocCount() (count int, success bool) {
func (db *ElasticDatabaseWithEviction) DocCount() (docCount int, err error) {
elasticsearchURL := fmt.Sprintf("%s/_search", db.indexName)
query := `{
"_source": false,
Expand All @@ -153,77 +124,57 @@ func (db *ElasticDatabaseWithEviction) DocCount() (count int, success bool) {
}
}`

resp, err := db.httpClient.Request(context.Background(), "GET", elasticsearchURL, []byte(query))
var resp *http.Response
resp, err = db.httpClient.DoRequestCheckResponseStatus(context.Background(), http.MethodGet, elasticsearchURL, []byte(query))
if err != nil {
return
if resp.StatusCode == http.StatusNoContent || resp.StatusCode == http.StatusNotFound {
return 0, nil
}
return -1, err
}
defer resp.Body.Close()

jsonAsBytes, err := io.ReadAll(resp.Body)
var jsonAsBytes []byte
jsonAsBytes, err = io.ReadAll(resp.Body)
if err != nil {
return
}

fmt.Println("kk dbg DocCount() resp.StatusCode:", resp.StatusCode)

switch resp.StatusCode {
case http.StatusOK:
break
case http.StatusNoContent, http.StatusNotFound:
return 0, true
default:
logger.WarnWithCtx(db.ctx).Msgf("failed to get from elastic: %s, response status code: %v", string(jsonAsBytes), resp.StatusCode)
return
}

// Unmarshal the JSON response
var result map[string]interface{}
if err = json.Unmarshal(jsonAsBytes, &result); err != nil {
logger.WarnWithCtx(db.ctx).Msgf("Error parsing the response JSON: %s", err)
return
}

fmt.Println("kk dbg DocCount() result:", result)

count = int(result["hits"].(map[string]interface{})["total"].(map[string]interface{})["value"].(float64)) // TODO: add some checks... to prevent panic
return count, true
return int(result["hits"].(map[string]interface{})["total"].(map[string]interface{})["value"].(float64)), nil // TODO: add some checks... to prevent panic
}

func (db *ElasticDatabaseWithEviction) SizeInBytes() (sizeInBytes int64, success bool) {
func (db *ElasticDatabaseWithEviction) SizeInBytes() (sizeInBytes int64, err error) {
elasticsearchURL := fmt.Sprintf("%s/_search", db.indexName)
query := `{
"_source": ["sizeInBytes"],
"size": 10000,
"track_total_hits": true
}`

resp, err := db.httpClient.Request(context.Background(), "GET", elasticsearchURL, []byte(query))
var resp *http.Response
resp, err = db.httpClient.DoRequestCheckResponseStatus(context.Background(), http.MethodGet, elasticsearchURL, []byte(query))
if err != nil {
return
}
defer resp.Body.Close()

jsonAsBytes, err := io.ReadAll(resp.Body)
var jsonAsBytes []byte
jsonAsBytes, err = io.ReadAll(resp.Body)
if err != nil {
return
}

fmt.Println("kk dbg SizeInBytes() resp.StatusCode:", resp.StatusCode)

switch resp.StatusCode {
case http.StatusOK:
break
case http.StatusNoContent, http.StatusNotFound:
return 0, true
default:
logger.WarnWithCtx(db.ctx).Msgf("failed to get from elastic: %s, response status code: %v", string(jsonAsBytes), resp.StatusCode)
return
}

// Unmarshal the JSON response
var result map[string]interface{}
if err = json.Unmarshal(jsonAsBytes, &result); err != nil {
logger.WarnWithCtx(db.ctx).Msgf("Error parsing the response JSON: %s", err)
return
}

Expand All @@ -234,64 +185,68 @@ func (db *ElasticDatabaseWithEviction) SizeInBytes() (sizeInBytes int64, success
a = append(a, sizeInBytes-b)
}
fmt.Println("kk dbg SizeInBytes() sizes in storage:", a)
return sizeInBytes, true
return sizeInBytes, nil
}

func (db *ElasticDatabaseWithEviction) SizeInBytesLimit() int64 {
return db.sizeInBytesLimit
}

func (db *ElasticDatabaseWithEviction) getAll() (documents []*document, success bool) {
elasticsearchURL := fmt.Sprintf("%s*/_search", db.indexName)
query := `{
func (db *ElasticDatabaseWithEviction) getAll() (documents []*document, err error) {
_ = fmt.Sprintf("%s*/_search", db.indexName)
_ = `{
"_source": {
"excludes": "data"
},
"size": 10000,
"track_total_hits": true
}`
/*
db.httpClient.
resp, err := db.httpClient.Request(context.Background(), "GET", elasticsearchURL, []byte(query))
if err != nil {
return
}
defer resp.Body.Close()
resp, err := db.httpClient.Request(context.Background(), "GET", elasticsearchURL, []byte(query))
if err != nil {
return
}
defer resp.Body.Close()
jsonAsBytes, err := io.ReadAll(resp.Body)
if err != nil {
return
}
jsonAsBytes, err := io.ReadAll(resp.Body)
if err != nil {
return
}
fmt.Println("kk dbg getAll() resp.StatusCode:", resp.StatusCode)
fmt.Println("kk dbg getAll() resp.StatusCode:", resp.StatusCode)
switch resp.StatusCode {
case http.StatusOK:
break
default:
logger.WarnWithCtx(db.ctx).Msgf("failed to get from elastic: %s, response status code: %v", string(jsonAsBytes), resp.StatusCode)
return
}
switch resp.StatusCode {
case http.StatusOK:
break
default:
logger.WarnWithCtx(db.ctx).Msgf("failed to get from elastic: %s, response status code: %v", string(jsonAsBytes), resp.StatusCode)
return
}
// Unmarshal the JSON response
var result map[string]interface{}
if err = json.Unmarshal(jsonAsBytes, &result); err != nil {
logger.WarnWithCtx(db.ctx).Msgf("Error parsing the response JSON: %s", err)
return
}
// Unmarshal the JSON response
var result map[string]interface{}
if err = json.Unmarshal(jsonAsBytes, &result); err != nil {
logger.WarnWithCtx(db.ctx).Msgf("Error parsing the response JSON: %s", err)
return
}
fmt.Println("kk dbg getAll() documents:")
for _, hit := range result["hits"].(map[string]interface{})["hits"].([]interface{}) {
doc := &document{
Id: hit.(map[string]interface{})["_id"].(string),
Index: hit.(map[string]interface{})["_index"].(string),
SizeInBytes: int64(hit.(map[string]interface{})["_source"].(map[string]interface{})["sizeInBytes"].(float64)), // TODO: add checks
//Timestamp: hit.(map[string]interface{})["_source"].(map[string]interface{})["timestamp"].(time.Time), // TODO: add checks
MarkedAsDeleted: hit.(map[string]interface{})["_source"].(map[string]interface{})["markedAsDeleted"].(bool), // TODO: add checks
fmt.Println("kk dbg getAll() documents:")
for _, hit := range result["hits"].(map[string]interface{})["hits"].([]interface{}) {
doc := &document{
Id: hit.(map[string]interface{})["_id"].(string),
Index: hit.(map[string]interface{})["_index"].(string),
SizeInBytes: int64(hit.(map[string]interface{})["_source"].(map[string]interface{})["sizeInBytes"].(float64)), // TODO: add checks
//Timestamp: hit.(map[string]interface{})["_source"].(map[string]interface{})["timestamp"].(time.Time), // TODO: add checks
MarkedAsDeleted: hit.(map[string]interface{})["_source"].(map[string]interface{})["markedAsDeleted"].(bool), // TODO: add checks
}
fmt.Println(doc)
documents = append(documents, doc)
}
fmt.Println(doc)
documents = append(documents, doc)
}
return documents, true
*/
return documents, nil
}

func (db *ElasticDatabaseWithEviction) evict(indexes []string) {
Expand Down
13 changes: 7 additions & 6 deletions quesma/persistence/model.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,13 @@ type JSONDatabase interface {
Put(key string, data string) error
}

type JSONDatabaseWithEviction interface { // for sure JSON? maybe not only json? check
Put(doc document) bool
Get(id string) (document, bool)
Delete(id string)
DocCount() (int, bool)
SizeInBytes() (int64, bool)
type DatabaseWithEviction interface { // for sure JSON? maybe not only json? check
Put(doc document) error
Get(id string) (document, error)
Delete(id string) error
DeleteOld(time.Duration) error
DocCount() (int, error)
SizeInBytes() (int64, error)
SizeInBytesLimit() int64
}

Expand Down
Loading

0 comments on commit 61e367a

Please sign in to comment.