Added support for Badger as an alternative storage backend

This commit is contained in:
Pijus Kamandulis
2025-03-12 21:06:10 +02:00
parent e526b2269e
commit 813b9faeaa
25 changed files with 1574 additions and 458 deletions

View File

@@ -15,6 +15,11 @@ const (
ExplorerBaseUrlLocation = "/_explorer"
)
const (
DataStoreMap = "map"
DataStoreBadger = "badger"
)
func ParseFlags() ServerConfig {
host := flag.String("Host", "localhost", "Hostname")
port := flag.Int("Port", 8081, "Listen port")
@@ -28,6 +33,8 @@ func ParseFlags() ServerConfig {
persistDataPath := flag.String("Persist", "", "Saves data to given path on application exit")
logLevel := NewEnumValue("info", []string{"debug", "info", "error", "silent"})
flag.Var(logLevel, "LogLevel", fmt.Sprintf("Sets the logging level %s", logLevel.AllowedValuesList()))
dataStore := NewEnumValue("map", []string{DataStoreMap, DataStoreBadger})
flag.Var(dataStore, "DataStore", fmt.Sprintf("Sets the data store %s, (badger is currently in the experimental phase)", dataStore.AllowedValuesList()))
flag.Parse()
setFlagsFromEnvironment()
@@ -44,6 +51,7 @@ func ParseFlags() ServerConfig {
config.DisableTls = *disableTls
config.AccountKey = *accountKey
config.LogLevel = logLevel.value
config.DataStore = dataStore.value
config.PopulateCalculatedFields()
@@ -68,6 +76,13 @@ func (c *ServerConfig) PopulateCalculatedFields() {
default:
logger.SetLogLevel(logger.LogLevelInfo)
}
if c.DataStore == DataStoreBadger &&
(c.InitialDataFilePath != "" || c.PersistDataFilePath != "") {
logger.ErrorLn("InitialData and Persist options are currently not supported with Badger data store")
c.InitialDataFilePath = ""
c.PersistDataFilePath = ""
}
}
func (c *ServerConfig) ApplyDefaultsToEmptyFields() {

View File

@@ -17,4 +17,6 @@ type ServerConfig struct {
DisableTls bool `json:"disableTls"`
LogLevel string `json:"logLevel"`
ExplorerBaseUrlLocation string `json:"explorerBaseUrlLocation"`
DataStore string `json:"dataStore"`
}

View File

@@ -383,6 +383,7 @@ func (h *Handlers) executeQueryDocuments(databaseId string, collectionId string,
if status != datastore.StatusOk {
return nil, status
}
defer allDocumentsIterator.Close()
rowsIterator := converters.NewDocumentToRowTypeIterator(allDocumentsIterator)

View File

@@ -3,32 +3,29 @@ package tests_test
import (
"context"
"errors"
"fmt"
"net/http"
"testing"
"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
"github.com/pikami/cosmium/api/config"
"github.com/pikami/cosmium/internal/datastore"
"github.com/stretchr/testify/assert"
)
func Test_Collections(t *testing.T) {
ts := runTestServer()
defer ts.Server.Close()
presets := []testPreset{PresetMapStore, PresetBadgerStore}
client, err := azcosmos.NewClientFromConnectionString(
fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
&azcosmos.ClientOptions{},
)
assert.Nil(t, err)
setUp := func(ts *TestServer, client *azcosmos.Client) *azcosmos.DatabaseClient {
ts.DataStore.CreateDatabase(datastore.Database{ID: testDatabaseName})
databaseClient, err := client.NewDatabase(testDatabaseName)
assert.Nil(t, err)
ts.DataStore.CreateDatabase(datastore.Database{ID: testDatabaseName})
databaseClient, err := client.NewDatabase(testDatabaseName)
assert.Nil(t, err)
return databaseClient
}
runTestsWithPresets(t, "Collection Create", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
databaseClient := setUp(ts, client)
t.Run("Collection Create", func(t *testing.T) {
t.Run("Should create collection", func(t *testing.T) {
createResponse, err := databaseClient.CreateContainer(context.TODO(), azcosmos.ContainerProperties{
ID: testCollectionName,
@@ -57,7 +54,9 @@ func Test_Collections(t *testing.T) {
})
})
t.Run("Collection Read", func(t *testing.T) {
runTestsWithPresets(t, "Collection Read", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
databaseClient := setUp(ts, client)
t.Run("Should read collection", func(t *testing.T) {
ts.DataStore.CreateCollection(testDatabaseName, datastore.Collection{
ID: testCollectionName,
@@ -90,7 +89,9 @@ func Test_Collections(t *testing.T) {
})
})
t.Run("Collection Delete", func(t *testing.T) {
runTestsWithPresets(t, "Collection Delete", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
databaseClient := setUp(ts, client)
t.Run("Should delete collection", func(t *testing.T) {
ts.DataStore.CreateCollection(testDatabaseName, datastore.Collection{
ID: testCollectionName,

View File

@@ -1,13 +1,18 @@
package tests_test
import (
"fmt"
"net/http/httptest"
"testing"
"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
"github.com/pikami/cosmium/api"
"github.com/pikami/cosmium/api/config"
"github.com/pikami/cosmium/internal/datastore"
badgerdatastore "github.com/pikami/cosmium/internal/datastore/badger_datastore"
mapdatastore "github.com/pikami/cosmium/internal/datastore/map_datastore"
"github.com/pikami/cosmium/internal/logger"
"github.com/stretchr/testify/assert"
)
type TestServer struct {
@@ -16,14 +21,29 @@ type TestServer struct {
URL string
}
func runTestServerCustomConfig(config *config.ServerConfig) *TestServer {
dataStore := mapdatastore.NewMapDataStore(mapdatastore.MapDataStoreOptions{})
func getDefaultTestServerConfig() *config.ServerConfig {
return &config.ServerConfig{
AccountKey: config.DefaultAccountKey,
ExplorerPath: "/tmp/nothing",
ExplorerBaseUrlLocation: config.ExplorerBaseUrlLocation,
DataStore: "map",
}
}
api := api.NewApiServer(dataStore, config)
func runTestServerCustomConfig(configuration *config.ServerConfig) *TestServer {
var dataStore datastore.DataStore
switch configuration.DataStore {
case config.DataStoreBadger:
dataStore = badgerdatastore.NewBadgerDataStore()
default:
dataStore = mapdatastore.NewMapDataStore(mapdatastore.MapDataStoreOptions{})
}
api := api.NewApiServer(dataStore, configuration)
server := httptest.NewServer(api.GetRouter())
config.DatabaseEndpoint = server.URL
configuration.DatabaseEndpoint = server.URL
return &TestServer{
Server: server,
@@ -33,11 +53,7 @@ func runTestServerCustomConfig(config *config.ServerConfig) *TestServer {
}
func runTestServer() *TestServer {
config := &config.ServerConfig{
AccountKey: config.DefaultAccountKey,
ExplorerPath: "/tmp/nothing",
ExplorerBaseUrlLocation: config.ExplorerBaseUrlLocation,
}
config := getDefaultTestServerConfig()
config.LogLevel = "debug"
logger.SetLogLevel(logger.LogLevelDebug)
@@ -50,3 +66,46 @@ const (
testDatabaseName = "test-db"
testCollectionName = "test-coll"
)
type testFunc func(t *testing.T, ts *TestServer, cosmosClient *azcosmos.Client)
type testPreset string
const (
PresetMapStore testPreset = "MapDS"
PresetBadgerStore testPreset = "BadgerDS"
)
func runTestsWithPreset(t *testing.T, name string, testPreset testPreset, f testFunc) {
serverConfig := getDefaultTestServerConfig()
serverConfig.LogLevel = "debug"
logger.SetLogLevel(logger.LogLevelDebug)
switch testPreset {
case PresetBadgerStore:
serverConfig.DataStore = config.DataStoreBadger
case PresetMapStore:
serverConfig.DataStore = config.DataStoreMap
}
ts := runTestServerCustomConfig(serverConfig)
defer ts.Server.Close()
client, err := azcosmos.NewClientFromConnectionString(
fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
&azcosmos.ClientOptions{},
)
assert.Nil(t, err)
testName := fmt.Sprintf("%s_%s", testPreset, name)
t.Run(testName, func(t *testing.T) {
f(t, ts, client)
})
}
func runTestsWithPresets(t *testing.T, name string, testPresets []testPreset, f testFunc) {
for _, testPreset := range testPresets {
runTestsWithPreset(t, name, testPreset, f)
}
}

View File

@@ -3,28 +3,19 @@ package tests_test
import (
"context"
"errors"
"fmt"
"net/http"
"testing"
"github.com/Azure/azure-sdk-for-go/sdk/azcore"
"github.com/Azure/azure-sdk-for-go/sdk/data/azcosmos"
"github.com/pikami/cosmium/api/config"
"github.com/pikami/cosmium/internal/datastore"
"github.com/stretchr/testify/assert"
)
func Test_Databases(t *testing.T) {
ts := runTestServer()
defer ts.Server.Close()
presets := []testPreset{PresetMapStore, PresetBadgerStore}
client, err := azcosmos.NewClientFromConnectionString(
fmt.Sprintf("AccountEndpoint=%s;AccountKey=%s", ts.URL, config.DefaultAccountKey),
&azcosmos.ClientOptions{},
)
assert.Nil(t, err)
t.Run("Database Create", func(t *testing.T) {
runTestsWithPresets(t, "Database Create", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
t.Run("Should create database", func(t *testing.T) {
ts.DataStore.DeleteDatabase(testDatabaseName)
@@ -55,7 +46,7 @@ func Test_Databases(t *testing.T) {
})
})
t.Run("Database Read", func(t *testing.T) {
runTestsWithPresets(t, "Database Read", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
t.Run("Should read database", func(t *testing.T) {
ts.DataStore.CreateDatabase(datastore.Database{
ID: testDatabaseName,
@@ -88,7 +79,7 @@ func Test_Databases(t *testing.T) {
})
})
t.Run("Database Delete", func(t *testing.T) {
runTestsWithPresets(t, "Database Delete", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
t.Run("Should delete database", func(t *testing.T) {
ts.DataStore.CreateDatabase(datastore.Database{
ID: testDatabaseName,

View File

@@ -53,9 +53,7 @@ func testCosmosQuery(t *testing.T,
}
}
func documents_InitializeDb(t *testing.T) (*TestServer, *azcosmos.ContainerClient) {
ts := runTestServer()
func documents_InitializeDb(t *testing.T, ts *TestServer) *azcosmos.ContainerClient {
ts.DataStore.CreateDatabase(datastore.Database{ID: testDatabaseName})
ts.DataStore.CreateCollection(testDatabaseName, datastore.Collection{
ID: testCollectionName,
@@ -79,438 +77,439 @@ func documents_InitializeDb(t *testing.T) (*TestServer, *azcosmos.ContainerClien
collectionClient, err := client.NewContainer(testDatabaseName, testCollectionName)
assert.Nil(t, err)
return ts, collectionClient
return collectionClient
}
func Test_Documents(t *testing.T) {
ts, collectionClient := documents_InitializeDb(t)
defer ts.Server.Close()
presets := []testPreset{PresetMapStore, PresetBadgerStore}
t.Run("Should query document", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
"SELECT c.id, c[\"pk\"] FROM c ORDER BY c.id",
nil,
[]interface{}{
map[string]interface{}{"id": "12345", "pk": "123"},
map[string]interface{}{"id": "67890", "pk": "456"},
},
)
runTestsWithPresets(t, "Test_Documents", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
collectionClient := documents_InitializeDb(t, ts)
t.Run("Should query document", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
"SELECT c.id, c[\"pk\"] FROM c ORDER BY c.id",
nil,
[]interface{}{
map[string]interface{}{"id": "12345", "pk": "123"},
map[string]interface{}{"id": "67890", "pk": "456"},
},
)
})
t.Run("Should query VALUE array", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
"SELECT VALUE [c.id, c[\"pk\"]] FROM c ORDER BY c.id",
nil,
[]interface{}{
[]interface{}{"12345", "123"},
[]interface{}{"67890", "456"},
},
)
})
t.Run("Should query VALUE object", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
"SELECT VALUE { id: c.id, _pk: c.pk } FROM c ORDER BY c.id",
nil,
[]interface{}{
map[string]interface{}{"id": "12345", "_pk": "123"},
map[string]interface{}{"id": "67890", "_pk": "456"},
},
)
})
t.Run("Should query document with single WHERE condition", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`select c.id
FROM c
WHERE c.isCool=true
ORDER BY c.id`,
nil,
[]interface{}{
map[string]interface{}{"id": "67890"},
},
)
})
t.Run("Should query document with query parameters", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`select c.id
FROM c
WHERE c.id=@param_id
ORDER BY c.id`,
[]azcosmos.QueryParameter{
{Name: "@param_id", Value: "67890"},
},
[]interface{}{
map[string]interface{}{"id": "67890"},
},
)
})
t.Run("Should query document with query parameters as accessor", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`select c.id
FROM c
WHERE c[@param]="67890"
ORDER BY c.id`,
[]azcosmos.QueryParameter{
{Name: "@param", Value: "id"},
},
[]interface{}{
map[string]interface{}{"id": "67890"},
},
)
})
t.Run("Should query array accessor", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`SELECT c.id,
c["arr"][0] AS arr0,
c["arr"][1] AS arr1,
c["arr"][2] AS arr2,
c["arr"][3] AS arr3
FROM c ORDER BY c.id`,
nil,
[]interface{}{
map[string]interface{}{"id": "12345", "arr0": 1.0, "arr1": 2.0, "arr2": 3.0, "arr3": nil},
map[string]interface{}{"id": "67890", "arr0": 6.0, "arr1": 7.0, "arr2": 8.0, "arr3": nil},
},
)
})
t.Run("Should handle parallel writes", func(t *testing.T) {
var wg sync.WaitGroup
rutineCount := 100
results := make(chan error, rutineCount)
createCall := func(i int) {
defer wg.Done()
item := map[string]interface{}{
"id": fmt.Sprintf("id-%d", i),
"pk": fmt.Sprintf("pk-%d", i),
"val": i,
}
bytes, err := json.Marshal(item)
if err != nil {
results <- err
return
}
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
_, err = collectionClient.CreateItem(
ctx,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
results <- err
collectionClient.ReadItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
collectionClient.DeleteItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
}
for i := 0; i < rutineCount; i++ {
wg.Add(1)
go createCall(i)
}
wg.Wait()
close(results)
for err := range results {
if err != nil {
t.Errorf("Error creating item: %v", err)
}
}
})
})
t.Run("Should query VALUE array", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
"SELECT VALUE [c.id, c[\"pk\"]] FROM c ORDER BY c.id",
nil,
[]interface{}{
[]interface{}{"12345", "123"},
[]interface{}{"67890", "456"},
},
)
})
runTestsWithPresets(t, "Test_Documents_Patch", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
collectionClient := documents_InitializeDb(t, ts)
t.Run("Should query VALUE object", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
"SELECT VALUE { id: c.id, _pk: c.pk } FROM c ORDER BY c.id",
nil,
[]interface{}{
map[string]interface{}{"id": "12345", "_pk": "123"},
map[string]interface{}{"id": "67890", "_pk": "456"},
},
)
})
t.Run("Should PATCH document", func(t *testing.T) {
context := context.TODO()
expectedData := map[string]interface{}{"id": "67890", "pk": "666", "newField": "newValue", "incr": 15., "setted": "isSet"}
t.Run("Should query document with single WHERE condition", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`select c.id
FROM c
WHERE c.isCool=true
ORDER BY c.id`,
nil,
[]interface{}{
map[string]interface{}{"id": "67890"},
},
)
})
patch := azcosmos.PatchOperations{}
patch.AppendAdd("/newField", "newValue")
patch.AppendIncrement("/incr", 15)
patch.AppendRemove("/isCool")
patch.AppendReplace("/pk", "666")
patch.AppendSet("/setted", "isSet")
t.Run("Should query document with query parameters", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`select c.id
FROM c
WHERE c.id=@param_id
ORDER BY c.id`,
[]azcosmos.QueryParameter{
{Name: "@param_id", Value: "67890"},
},
[]interface{}{
map[string]interface{}{"id": "67890"},
},
)
})
itemResponse, err := collectionClient.PatchItem(
context,
azcosmos.PartitionKey{},
"67890",
patch,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.Nil(t, err)
t.Run("Should query document with query parameters as accessor", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`select c.id
FROM c
WHERE c[@param]="67890"
ORDER BY c.id`,
[]azcosmos.QueryParameter{
{Name: "@param", Value: "id"},
},
[]interface{}{
map[string]interface{}{"id": "67890"},
},
)
})
var itemResponseBody map[string]interface{}
json.Unmarshal(itemResponse.Value, &itemResponseBody)
t.Run("Should query array accessor", func(t *testing.T) {
testCosmosQuery(t, collectionClient,
`SELECT c.id,
c["arr"][0] AS arr0,
c["arr"][1] AS arr1,
c["arr"][2] AS arr2,
c["arr"][3] AS arr3
FROM c ORDER BY c.id`,
nil,
[]interface{}{
map[string]interface{}{"id": "12345", "arr0": 1.0, "arr1": 2.0, "arr2": 3.0, "arr3": nil},
map[string]interface{}{"id": "67890", "arr0": 6.0, "arr1": 7.0, "arr2": 8.0, "arr3": nil},
},
)
})
assert.Equal(t, expectedData["id"], itemResponseBody["id"])
assert.Equal(t, expectedData["pk"], itemResponseBody["pk"])
assert.Empty(t, itemResponseBody["isCool"])
assert.Equal(t, expectedData["newField"], itemResponseBody["newField"])
assert.Equal(t, expectedData["incr"], itemResponseBody["incr"])
assert.Equal(t, expectedData["setted"], itemResponseBody["setted"])
})
t.Run("Should handle parallel writes", func(t *testing.T) {
var wg sync.WaitGroup
rutineCount := 100
results := make(chan error, rutineCount)
t.Run("Should not allow to PATCH document ID", func(t *testing.T) {
context := context.TODO()
patch := azcosmos.PatchOperations{}
patch.AppendReplace("/id", "newValue")
_, err := collectionClient.PatchItem(
context,
azcosmos.PartitionKey{},
"67890",
patch,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, err)
var respErr *azcore.ResponseError
if errors.As(err, &respErr) {
assert.Equal(t, http.StatusUnprocessableEntity, respErr.StatusCode)
} else {
panic(err)
}
})
t.Run("CreateItem", func(t *testing.T) {
context := context.TODO()
createCall := func(i int) {
defer wg.Done()
item := map[string]interface{}{
"id": fmt.Sprintf("id-%d", i),
"pk": fmt.Sprintf("pk-%d", i),
"val": i,
"Id": "6789011",
"pk": "456",
"newField": "newValue2",
}
bytes, err := json.Marshal(item)
if err != nil {
results <- err
return
}
assert.Nil(t, err)
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
_, err = collectionClient.CreateItem(
ctx,
r, err2 := collectionClient.CreateItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
results <- err
assert.NotNil(t, r)
assert.Nil(t, err2)
})
collectionClient.ReadItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
collectionClient.DeleteItem(ctx, azcosmos.PartitionKey{}, fmt.Sprintf("id-%d", i), nil)
}
t.Run("CreateItem that already exists", func(t *testing.T) {
context := context.TODO()
for i := 0; i < rutineCount; i++ {
wg.Add(1)
go createCall(i)
}
item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3}}
bytes, err := json.Marshal(item)
assert.Nil(t, err)
wg.Wait()
close(results)
r, err := collectionClient.CreateItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, r)
assert.NotNil(t, err)
for err := range results {
if err != nil {
t.Errorf("Error creating item: %v", err)
var respErr *azcore.ResponseError
if errors.As(err, &respErr) {
assert.Equal(t, http.StatusConflict, respErr.StatusCode)
} else {
panic(err)
}
}
})
}
func Test_Documents_Patch(t *testing.T) {
ts, collectionClient := documents_InitializeDb(t)
defer ts.Server.Close()
t.Run("Should PATCH document", func(t *testing.T) {
context := context.TODO()
expectedData := map[string]interface{}{"id": "67890", "pk": "666", "newField": "newValue", "incr": 15., "setted": "isSet"}
patch := azcosmos.PatchOperations{}
patch.AppendAdd("/newField", "newValue")
patch.AppendIncrement("/incr", 15)
patch.AppendRemove("/isCool")
patch.AppendReplace("/pk", "666")
patch.AppendSet("/setted", "isSet")
itemResponse, err := collectionClient.PatchItem(
context,
azcosmos.PartitionKey{},
"67890",
patch,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.Nil(t, err)
var itemResponseBody map[string]interface{}
json.Unmarshal(itemResponse.Value, &itemResponseBody)
assert.Equal(t, expectedData["id"], itemResponseBody["id"])
assert.Equal(t, expectedData["pk"], itemResponseBody["pk"])
assert.Empty(t, itemResponseBody["isCool"])
assert.Equal(t, expectedData["newField"], itemResponseBody["newField"])
assert.Equal(t, expectedData["incr"], itemResponseBody["incr"])
assert.Equal(t, expectedData["setted"], itemResponseBody["setted"])
})
t.Run("Should not allow to PATCH document ID", func(t *testing.T) {
context := context.TODO()
patch := azcosmos.PatchOperations{}
patch.AppendReplace("/id", "newValue")
_, err := collectionClient.PatchItem(
context,
azcosmos.PartitionKey{},
"67890",
patch,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, err)
var respErr *azcore.ResponseError
if errors.As(err, &respErr) {
assert.Equal(t, http.StatusUnprocessableEntity, respErr.StatusCode)
} else {
panic(err)
}
})
t.Run("CreateItem", func(t *testing.T) {
context := context.TODO()
item := map[string]interface{}{
"Id": "6789011",
"pk": "456",
"newField": "newValue2",
}
bytes, err := json.Marshal(item)
assert.Nil(t, err)
r, err2 := collectionClient.CreateItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, r)
assert.Nil(t, err2)
})
t.Run("CreateItem that already exists", func(t *testing.T) {
context := context.TODO()
item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3}}
bytes, err := json.Marshal(item)
assert.Nil(t, err)
r, err := collectionClient.CreateItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, r)
assert.NotNil(t, err)
var respErr *azcore.ResponseError
if errors.As(err, &respErr) {
assert.Equal(t, http.StatusConflict, respErr.StatusCode)
} else {
panic(err)
}
})
t.Run("UpsertItem new", func(t *testing.T) {
context := context.TODO()
item := map[string]interface{}{"id": "123456", "pk": "1234", "isCool": false, "arr": []int{1, 2, 3}}
bytes, err := json.Marshal(item)
assert.Nil(t, err)
r, err2 := collectionClient.UpsertItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, r)
assert.Nil(t, err2)
})
t.Run("UpsertItem that already exists", func(t *testing.T) {
context := context.TODO()
item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3, 4}}
bytes, err := json.Marshal(item)
assert.Nil(t, err)
r, err2 := collectionClient.UpsertItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, r)
assert.Nil(t, err2)
})
}
func Test_Documents_TransactionalBatch(t *testing.T) {
ts, collectionClient := documents_InitializeDb(t)
defer ts.Server.Close()
t.Run("Should execute CREATE transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
newItem := map[string]interface{}{
"id": "678901",
}
bytes, err := json.Marshal(newItem)
assert.Nil(t, err)
batch.CreateItem(bytes, nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, newItem["id"], itemResponseBody["id"])
createdDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
assert.Equal(t, newItem["id"], createdDoc["id"])
})
t.Run("Should execute DELETE transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
batch.DeleteItem("12345", nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.Equal(t, int32(http.StatusNoContent), operationResponse.StatusCode)
_, status := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, "12345")
assert.Equal(t, datastore.StatusNotFound, int(status))
})
t.Run("Should execute REPLACE transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
newItem := map[string]interface{}{
"id": "67890",
"pk": "666",
}
bytes, err := json.Marshal(newItem)
assert.Nil(t, err)
batch.ReplaceItem("67890", bytes, nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, newItem["id"], itemResponseBody["id"])
assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
assert.Equal(t, newItem["id"], updatedDoc["id"])
assert.Equal(t, newItem["pk"], updatedDoc["pk"])
})
t.Run("Should execute UPSERT transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
newItem := map[string]interface{}{
"id": "678901",
"pk": "666",
}
bytes, err := json.Marshal(newItem)
assert.Nil(t, err)
batch.UpsertItem(bytes, nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, newItem["id"], itemResponseBody["id"])
assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
assert.Equal(t, newItem["id"], updatedDoc["id"])
assert.Equal(t, newItem["pk"], updatedDoc["pk"])
})
t.Run("Should execute READ transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
batch.ReadItem("67890", nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusOK), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, "67890", itemResponseBody["id"])
})
t.Run("UpsertItem new", func(t *testing.T) {
context := context.TODO()
item := map[string]interface{}{"id": "123456", "pk": "1234", "isCool": false, "arr": []int{1, 2, 3}}
bytes, err := json.Marshal(item)
assert.Nil(t, err)
r, err2 := collectionClient.UpsertItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, r)
assert.Nil(t, err2)
})
t.Run("UpsertItem that already exists", func(t *testing.T) {
context := context.TODO()
item := map[string]interface{}{"id": "12345", "pk": "123", "isCool": false, "arr": []int{1, 2, 3, 4}}
bytes, err := json.Marshal(item)
assert.Nil(t, err)
r, err2 := collectionClient.UpsertItem(
context,
azcosmos.PartitionKey{},
bytes,
&azcosmos.ItemOptions{
EnableContentResponseOnWrite: false,
},
)
assert.NotNil(t, r)
assert.Nil(t, err2)
})
})
runTestsWithPresets(t, "Test_Documents_TransactionalBatch", presets, func(t *testing.T, ts *TestServer, client *azcosmos.Client) {
collectionClient := documents_InitializeDb(t, ts)
t.Run("Should execute CREATE transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
newItem := map[string]interface{}{
"id": "678901",
}
bytes, err := json.Marshal(newItem)
assert.Nil(t, err)
batch.CreateItem(bytes, nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, newItem["id"], itemResponseBody["id"])
createdDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
assert.Equal(t, newItem["id"], createdDoc["id"])
})
t.Run("Should execute DELETE transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
batch.DeleteItem("12345", nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.Equal(t, int32(http.StatusNoContent), operationResponse.StatusCode)
_, status := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, "12345")
assert.Equal(t, datastore.StatusNotFound, int(status))
})
t.Run("Should execute REPLACE transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
newItem := map[string]interface{}{
"id": "67890",
"pk": "666",
}
bytes, err := json.Marshal(newItem)
assert.Nil(t, err)
batch.ReplaceItem("67890", bytes, nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, newItem["id"], itemResponseBody["id"])
assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
assert.Equal(t, newItem["id"], updatedDoc["id"])
assert.Equal(t, newItem["pk"], updatedDoc["pk"])
})
t.Run("Should execute UPSERT transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
newItem := map[string]interface{}{
"id": "678901",
"pk": "666",
}
bytes, err := json.Marshal(newItem)
assert.Nil(t, err)
batch.UpsertItem(bytes, nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusCreated), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, newItem["id"], itemResponseBody["id"])
assert.Equal(t, newItem["pk"], itemResponseBody["pk"])
updatedDoc, _ := ts.DataStore.GetDocument(testDatabaseName, testCollectionName, newItem["id"].(string))
assert.Equal(t, newItem["id"], updatedDoc["id"])
assert.Equal(t, newItem["pk"], updatedDoc["pk"])
})
t.Run("Should execute READ transactional batch", func(t *testing.T) {
context := context.TODO()
batch := collectionClient.NewTransactionalBatch(azcosmos.NewPartitionKeyString("pk"))
batch.ReadItem("67890", nil)
response, err := collectionClient.ExecuteTransactionalBatch(context, batch, &azcosmos.TransactionalBatchOptions{})
assert.Nil(t, err)
assert.True(t, response.Success)
assert.Equal(t, 1, len(response.OperationResults))
operationResponse := response.OperationResults[0]
assert.NotNil(t, operationResponse)
assert.NotNil(t, operationResponse.ResourceBody)
assert.Equal(t, int32(http.StatusOK), operationResponse.StatusCode)
var itemResponseBody map[string]interface{}
json.Unmarshal(operationResponse.ResourceBody, &itemResponseBody)
assert.Equal(t, "67890", itemResponseBody["id"])
})
})
}

View File

@@ -14,7 +14,8 @@ import (
// Request document with trailing slash like python cosmosdb client does.
func Test_Documents_Read_Trailing_Slash(t *testing.T) {
ts, _ := documents_InitializeDb(t)
ts := runTestServer()
documents_InitializeDb(t, ts)
defer ts.Server.Close()
t.Run("Read doc with client that appends slash to path", func(t *testing.T) {