github.com/siglens/siglens@v0.0.0-20240328180423-f7ce9ae441ed/pkg/es/writer/esBulkHandler_test.go (about)

     1  /*
     2  Copyright 2023.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package writer
    18  
    19  import (
    20  	"bytes"
    21  	"encoding/json"
    22  	"fmt"
    23  	"os"
    24  	"testing"
    25  
    26  	jp "github.com/buger/jsonparser"
    27  	"github.com/siglens/siglens/pkg/config"
    28  	"github.com/siglens/siglens/pkg/utils"
    29  	vtable "github.com/siglens/siglens/pkg/virtualtable"
    30  	"github.com/stretchr/testify/assert"
    31  )
    32  
    33  // Test ingesting multiple types of values into one column.
    34  // Currently the only test is that it doesn't crash.
    35  func Test_IngestMultipleTypesIntoOneColumn(t *testing.T) {
    36  	// Setup ingestion parameters.
    37  	now := utils.GetCurrentTimeInMs()
    38  	indexName := "traces"
    39  	shouldFlush := false
    40  	localIndexMap := make(map[string]string)
    41  	orgId := uint64(0)
    42  
    43  	flush := func() {
    44  		jsonBytes := []byte(`{"hello": "world"}`)
    45  		err := ProcessIndexRequest(jsonBytes, now, indexName, uint64(len(jsonBytes)), true, localIndexMap, orgId)
    46  		assert.Nil(t, err)
    47  	}
    48  
    49  	config.InitializeTestingConfig()
    50  	_ = vtable.InitVTable()
    51  
    52  	// Ingest some data that can all be converted to numbers.
    53  	jsons := [][]byte{
    54  		[]byte(`{"age": "171"}`),
    55  		[]byte(`{"age": 103}`),
    56  		[]byte(`{"age": 5.123}`),
    57  		[]byte(`{"age": "181"}`),
    58  		[]byte(`{"age": 30}`),
    59  		[]byte(`{"age": 6.321}`),
    60  	}
    61  
    62  	for _, jsonBytes := range jsons {
    63  		err := ProcessIndexRequest(jsonBytes, now, indexName, uint64(len(jsonBytes)), shouldFlush, localIndexMap, orgId)
    64  		assert.Nil(t, err)
    65  	}
    66  	flush()
    67  
    68  	// Ingest some data that will need to be converted to strings.
    69  	jsons = [][]byte{
    70  		[]byte(`{"age": "171"}`),
    71  		[]byte(`{"age": 103}`),
    72  		[]byte(`{"age": 5.123}`),
    73  		[]byte(`{"age": true}`),
    74  		[]byte(`{"age": "181"}`),
    75  		[]byte(`{"age": 30}`),
    76  		[]byte(`{"age": 6.321}`),
    77  		[]byte(`{"age": false}`),
    78  		[]byte(`{"age": "hello"}`),
    79  	}
    80  
    81  	for _, jsonBytes := range jsons {
    82  		err := ProcessIndexRequest(jsonBytes, now, indexName, uint64(len(jsonBytes)), shouldFlush, localIndexMap, orgId)
    83  		assert.Nil(t, err)
    84  	}
    85  	flush()
    86  
    87  	// Cleanup
    88  	os.RemoveAll(config.GetDataPath())
    89  }
    90  
    91  func flattenJson(currKey string, data []byte) error {
    92  	handler := func(key []byte, value []byte, valueType jp.ValueType, off int) error {
    93  		// Maybe push some state onto a stack here?
    94  		var finalKey string
    95  		if currKey == "" {
    96  			finalKey = string(key)
    97  		} else {
    98  			finalKey = fmt.Sprintf("%s.%s", currKey, key)
    99  		}
   100  		if valueType == jp.Object {
   101  			return flattenJson(finalKey, value)
   102  		} else if valueType == jp.Array {
   103  			return flattenArray(finalKey, value)
   104  		}
   105  		return fmt.Errorf("unknown value %+v", value)
   106  	}
   107  	return jp.ObjectEach(data, handler)
   108  
   109  }
   110  
   111  func flattenArray(currKey string, data []byte) error {
   112  	i := 0
   113  	_, _ = jp.ArrayEach(data, func(value []byte, valueType jp.ValueType, offset int, err error) {
   114  		var finalKey string
   115  		if currKey == "" {
   116  			finalKey = fmt.Sprintf("%d", i)
   117  		} else {
   118  			finalKey = fmt.Sprintf("%s.%d", currKey, i)
   119  		}
   120  		if valueType == jp.Object {
   121  			_ = flattenJson(finalKey, value)
   122  			return
   123  		}
   124  		// log.Infof("key: '%s', value: '%s', value type: %s", finalKey, string(value), valueType)
   125  		i++
   126  	})
   127  	return nil
   128  }
   129  
   130  func Benchmark_OriginalJson(b *testing.B) {
   131  	data := []byte(`{
   132  		"person": {
   133  		  "name": {
   134  			"first": "Leonid",
   135  			"last": "Bugaev",
   136  			"fullName": "Leonid Bugaev"
   137  		  },
   138  		  "github": {
   139  			"handle": "buger",
   140  			"followers": 109
   141  		  },
   142  		  "avatars": [
   143  			{ "url": "https://avatars1.githubusercontent.com/u/14009?v=3&s=460", "type": "thumbnail" }
   144  		  ]
   145  		},
   146  		"company": {
   147  		  "name": "Acme"
   148  		}
   149  	  }`)
   150  	jsonAction := make(map[string]interface{})
   151  	for i := 0; i < b.N; i++ {
   152  		decoder := json.NewDecoder(bytes.NewReader(data))
   153  		decoder.UseNumber()
   154  		_ = decoder.Decode(&jsonAction)
   155  	}
   156  }
   157  
   158  func Benchmark_bugerJsonParse(b *testing.B) {
   159  	data := []byte(`{
   160  		"person": {
   161  		  "name": {
   162  			"first": "Leonid",
   163  			"last": "Bugaev",
   164  			"fullName": "Leonid Bugaev"
   165  		  },
   166  		  "github": {
   167  			"handle": "buger",
   168  			"followers": 109
   169  		  },
   170  		  "avatars": [
   171  			{ "url": "https://avatars1.githubusercontent.com/u/14009?v=3&s=460", "type": "thumbnail" }
   172  		  ]
   173  		},
   174  		"company": {
   175  		  "name": "Acme"
   176  		}
   177  	  }`)
   178  	for i := 0; i < b.N; i++ {
   179  		_ = flattenJson("", data)
   180  	}
   181  }