query-orchestration/internal/database/repository/query_test.go

package repository_test

import (
	"context"
	"testing"

	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/test"

	"github.com/google/uuid"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestQueries(t *testing.T) {
	t.Parallel()
	if testing.Short() {
		t.SkipNow()
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	net := test.GetNetwork(t, ctx)
	_ = test.CreateDB(t, ctx, cfg, net, &test.CreateDatabaseConfig{})

	queries := cfg.GetDBQueries()

	contextQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeContextFull))
	require.NoError(t, err)

	contextQuery, err := queries.GetQuery(ctx, contextQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            contextQueryID,
		Type:          repository.QuerytypeContextFull,
		Activeversion: 0,
		Latestversion: 0,
		Config:        nil,
		Requiredids:   []uuid.UUID{},
	}, contextQuery)

	ctxVersion, err := queries.AddLatestQueryVersion(ctx, contextQueryID)
	require.NoError(t, err)
	assert.Equal(t, int32(1), ctxVersion)

	contextQuery, err = queries.GetQuery(ctx, contextQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            contextQueryID,
		Type:          repository.QuerytypeContextFull,
		Activeversion: 0,
		Latestversion: 1,
		Config:        nil,
		Requiredids:   []uuid.UUID{},
	}, contextQuery)

	jsonQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeJsonExtractor))
	require.NoError(t, err)

	jsonQuery, err := queries.GetQuery(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            jsonQueryID,
		Type:          repository.QuerytypeJsonExtractor,
		Activeversion: 0,
		Latestversion: 0,
		Config:        nil,
		Requiredids:   []uuid.UUID{},
	}, jsonQuery)

	jsonVersion, err := queries.AddLatestQueryVersion(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.Equal(t, int32(1), jsonVersion)

	jsonQuery, err = queries.GetQuery(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            jsonQueryID,
		Type:          repository.QuerytypeJsonExtractor,
		Activeversion: 0,
		Latestversion: 1,
		Config:        nil,
		Requiredids:   []uuid.UUID{},
	}, jsonQuery)

	err = queries.AddActiveQueryVersion(ctx, &repository.AddActiveQueryVersionParams{
		Versionid: 1,
		Queryid:   jsonQueryID,
	})
	require.NoError(t, err)

	jsonQuery, err = queries.GetQuery(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            jsonQueryID,
		Type:          repository.QuerytypeJsonExtractor,
		Activeversion: 1,
		Latestversion: 1,
		Config:        nil,
		Requiredids:   []uuid.UUID{},
	}, jsonQuery)

	jsonConfig := []byte("{\"path\": \"example_path\"}")

	err = queries.AddRequiredQuery(ctx, &repository.AddRequiredQueryParams{
		Queryid:         jsonQueryID,
		Requiredqueryid: contextQueryID,
		Addedversion:    1,
	})
	require.NoError(t, err)

	jsonQuery, err = queries.GetQuery(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            jsonQueryID,
		Type:          repository.QuerytypeJsonExtractor,
		Activeversion: 1,
		Latestversion: 1,
		Config:        nil,
		Requiredids:   []uuid.UUID{contextQueryID},
	}, jsonQuery)

	jsonVersion, err = queries.AddLatestQueryVersion(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.Equal(t, int32(2), jsonVersion)

	jsonQuery, err = queries.GetQuery(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            jsonQueryID,
		Type:          repository.QuerytypeJsonExtractor,
		Activeversion: 1,
		Latestversion: 2,
		Config:        nil,
		Requiredids:   []uuid.UUID{contextQueryID},
	}, jsonQuery)

	removeV := int32(2)
	err = queries.RemoveRequiredQuery(ctx, &repository.RemoveRequiredQueryParams{
		Queryid:         jsonQueryID,
		Requiredqueryid: contextQueryID,
		Removedversion:  &removeV,
	})
	require.NoError(t, err)

	err = queries.SetQueryConfig(ctx, &repository.SetQueryConfigParams{
		Queryid:      jsonQueryID,
		Config:       jsonConfig,
		Addedversion: 1,
	})
	require.NoError(t, err)

	jsonQueryConfig, err := queries.GetActiveQueryConfig(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.Equal(t, jsonConfig, jsonQueryConfig)

	err = queries.SetQueryConfig(ctx, &repository.SetQueryConfigParams{
		Queryid:      jsonQueryID,
		Config:       []byte(`{"second":"key"}`),
		Addedversion: 2,
	})
	require.NoError(t, err)

	jsonQueryConfig, err = queries.GetActiveQueryConfig(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.Equal(t, jsonConfig, jsonQueryConfig)

	err = queries.AddActiveQueryVersion(ctx, &repository.AddActiveQueryVersionParams{
		Queryid:   jsonQueryID,
		Versionid: 2,
	})
	require.NoError(t, err)

	jsonQueryConfig, err = queries.GetActiveQueryConfig(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.Equal(t, []byte(`{"second": "key"}`), jsonQueryConfig)

	jsonQuery, err = queries.GetQuery(ctx, jsonQueryID)
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.Fullactivequery{
		ID:            jsonQueryID,
		Type:          repository.QuerytypeJsonExtractor,
		Activeversion: 2,
		Latestversion: 2,
		Config:        jsonQueryConfig,
		Requiredids:   []uuid.UUID{},
	}, jsonQuery)

	v := int32(1)
	versionedQuery, err := queries.GetQueryWithVersion(ctx, &repository.GetQueryWithVersionParams{
		ID:      &jsonQueryID,
		Version: &v,
	})
	require.NoError(t, err)
	assert.EqualExportedValues(t, &repository.GetQueryWithVersionRow{
		ID:            jsonQueryID,
		Type:          repository.QuerytypeJsonExtractor,
		Activeversion: 2,
		Latestversion: 2,
		Config:        jsonConfig,
		Requiredids:   []uuid.UUID{contextQueryID},
	}, versionedQuery)

	all_exist, err := queries.AllQueriesExist(ctx, []uuid.UUID{})
	require.NoError(t, err)
	assert.True(t, all_exist)

	all_exist, err = queries.AllQueriesExist(ctx, []uuid.UUID{uuid.New()})
	require.NoError(t, err)
	assert.False(t, all_exist)

	all_exist, err = queries.AllQueriesExist(ctx, []uuid.UUID{jsonQueryID})
	require.NoError(t, err)
	assert.True(t, all_exist)

	all_exist, err = queries.AllQueriesExist(ctx, []uuid.UUID{jsonQueryID, contextQueryID})
	require.NoError(t, err)
	assert.True(t, all_exist)

	all_exist, err = queries.AllQueriesExist(ctx, []uuid.UUID{jsonQueryID, uuid.New()})
	require.NoError(t, err)
	assert.False(t, all_exist)
}

func TestQueryDependencyTree(t *testing.T) {
	t.Parallel()
	if testing.Short() {
		t.SkipNow()
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	net := test.GetNetwork(t, ctx)
	_ = test.CreateDB(t, ctx, cfg, net, &test.CreateDatabaseConfig{})

	queries := cfg.GetDBQueries()

	clientID := "EXAMPLE"
	err := queries.CreateClient(ctx, &repository.CreateClientParams{
		Name: "example_client",
		ID:   clientID,
	})
	require.NoError(t, err)
	version, err := queries.AddLatestCollectorVersion(ctx, clientID)
	require.NoError(t, err)
	err = queries.SetActiveCollectorVersion(ctx, &repository.SetActiveCollectorVersionParams{
		Clientid:  clientID,
		Versionid: version,
	})
	require.NoError(t, err)
	docID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
		Clientid: clientID,
		Hash:     "sample",
	})
	require.NoError(t, err)

	contextQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeContextFull))
	require.NoError(t, err)
	_, err = queries.AddLatestQueryVersion(ctx, contextQueryID)
	require.NoError(t, err)
	err = queries.AddActiveQueryVersion(ctx, &repository.AddActiveQueryVersionParams{
		Queryid:   contextQueryID,
		Versionid: 1,
	})
	require.NoError(t, err)

	dependents, err := queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    contextQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []uuid.UUID{}, dependents)

	jsonQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeJsonExtractor))
	require.NoError(t, err)
	_, err = queries.AddLatestQueryVersion(ctx, jsonQueryID)
	require.NoError(t, err)
	err = queries.AddActiveQueryVersion(ctx, &repository.AddActiveQueryVersionParams{
		Queryid:   jsonQueryID,
		Versionid: 1,
	})
	require.NoError(t, err)

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    jsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []uuid.UUID{}, dependents)

	err = queries.AddRequiredQuery(ctx, &repository.AddRequiredQueryParams{
		Queryid:         jsonQueryID,
		Requiredqueryid: contextQueryID,
		Addedversion:    1,
	})
	require.NoError(t, err)

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    jsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []uuid.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    contextQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []uuid.UUID{}, dependents)

	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Clientid:     clientID,
		Name:         "example_name'",
		Queryid:      jsonQueryID,
		Addedversion: 1,
	})
	require.NoError(t, err)

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    jsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    contextQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{&jsonQueryID}, dependents)

	secondJsonQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeJsonExtractor))
	require.NoError(t, err)
	_, err = queries.AddLatestQueryVersion(ctx, secondJsonQueryID)
	require.NoError(t, err)
	err = queries.AddActiveQueryVersion(ctx, &repository.AddActiveQueryVersionParams{
		Queryid:   secondJsonQueryID,
		Versionid: 1,
	})
	require.NoError(t, err)

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    jsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    contextQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{&jsonQueryID}, dependents)

	err = queries.AddRequiredQuery(ctx, &repository.AddRequiredQueryParams{
		Queryid:         secondJsonQueryID,
		Requiredqueryid: jsonQueryID,
		Addedversion:    1,
	})
	require.NoError(t, err)

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    secondJsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    jsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    contextQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{&jsonQueryID}, dependents)

	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Clientid:     clientID,
		Name:         "example_name'",
		Queryid:      secondJsonQueryID,
		Addedversion: 1,
	})
	require.NoError(t, err)

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    secondJsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    jsonQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{&secondJsonQueryID}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		Documentid: docID,
		Queryid:    contextQueryID,
	})
	require.NoError(t, err)
	assert.ElementsMatch(t, []*uuid.UUID{&jsonQueryID}, dependents)

	isdependent, err := queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Queryid:          &jsonQueryID,
		Requiredqueryids: []uuid.UUID{contextQueryID},
	})
	require.NoError(t, err)
	assert.False(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Queryid:          &jsonQueryID,
		Requiredqueryids: []uuid.UUID{secondJsonQueryID},
	})
	require.NoError(t, err)
	assert.True(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Queryid:          &jsonQueryID,
		Requiredqueryids: []uuid.UUID{jsonQueryID},
	})
	require.NoError(t, err)
	assert.True(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Queryid:          &secondJsonQueryID,
		Requiredqueryids: []uuid.UUID{jsonQueryID, contextQueryID},
	})
	require.NoError(t, err)
	assert.False(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Queryid:          &contextQueryID,
		Requiredqueryids: []uuid.UUID{jsonQueryID, secondJsonQueryID},
	})
	require.NoError(t, err)
	assert.True(t, isdependent)
}

func TestQueriesList(t *testing.T) {
	t.Parallel()
	if testing.Short() {
		t.SkipNow()
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	net := test.GetNetwork(t, ctx)
	_ = test.CreateDB(t, ctx, cfg, net, &test.CreateDatabaseConfig{})

	queries := cfg.GetDBQueries()

	contextQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeContextFull))
	require.NoError(t, err)

	jsonQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeJsonExtractor))
	require.NoError(t, err)

	qs, err := queries.ListQueries(ctx)
	require.NoError(t, err)
	assert.Len(t, qs, 2)
	assert.ElementsMatch(t, []*repository.Fullactivequery{
		{
			ID:            jsonQueryID,
			Type:          repository.QuerytypeJsonExtractor,
			Activeversion: 0,
			Latestversion: 0,
			Config:        nil,
			Requiredids:   []uuid.UUID{},
		},
		{
			ID:            contextQueryID,
			Type:          repository.QuerytypeContextFull,
			Activeversion: 0,
			Latestversion: 0,
			Config:        nil,
			Requiredids:   []uuid.UUID{},
		},
	}, qs)

	qs, err = queries.ListQueriesById(ctx, []uuid.UUID{jsonQueryID})
	require.NoError(t, err)
	assert.Len(t, qs, 1)
	assert.ElementsMatch(t, []*repository.Fullactivequery{
		{
			ID:            jsonQueryID,
			Type:          repository.QuerytypeJsonExtractor,
			Activeversion: 0,
			Latestversion: 0,
			Config:        nil,
			Requiredids:   []uuid.UUID{},
		},
	}, qs)
}

func TestListQueryClients(t *testing.T) {
	t.Parallel()
	if testing.Short() {
		t.SkipNow()
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	net := test.GetNetwork(t, ctx)
	_ = test.CreateDB(t, ctx, cfg, net, &test.CreateDatabaseConfig{})

	queries := cfg.GetDBQueries()

	contextID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeContextFull))
	require.NoError(t, err)

	clients, err := queries.ListQueryClientIDs(ctx, &contextID)
	require.NoError(t, err)
	assert.ElementsMatch(t, []uuid.UUID{}, clients)

	clientOneID := "EXAMPLE"
	err = queries.CreateClient(ctx, &repository.CreateClientParams{
		Name: "example_client",
		ID:   clientOneID,
	})
	require.NoError(t, err)
	versionOne, err := queries.AddLatestCollectorVersion(ctx, clientOneID)
	require.NoError(t, err)
	err = queries.SetActiveCollectorVersion(ctx, &repository.SetActiveCollectorVersionParams{
		Clientid:  clientOneID,
		Versionid: versionOne,
	})
	require.NoError(t, err)
	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Clientid:     clientOneID,
		Queryid:      contextID,
		Addedversion: 1,
		Name:         "example_key",
	})
	require.NoError(t, err)

	clients, err = queries.ListQueryClientIDs(ctx, &contextID)
	require.NoError(t, err)
	assert.ElementsMatch(t, []string{clientOneID}, clients)

	clientTwoID := "EXAMPLE_DOS"
	err = queries.CreateClient(ctx, &repository.CreateClientParams{
		Name: "example_client_dos",
		ID:   clientTwoID,
	})
	require.NoError(t, err)
	versionTwo, err := queries.AddLatestCollectorVersion(ctx, clientTwoID)
	require.NoError(t, err)
	err = queries.SetActiveCollectorVersion(ctx, &repository.SetActiveCollectorVersionParams{
		Clientid:  clientTwoID,
		Versionid: versionTwo,
	})
	require.NoError(t, err)
	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Clientid:     clientTwoID,
		Queryid:      contextID,
		Addedversion: 1,
		Name:         "example_key",
	})
	require.NoError(t, err)

	clients, err = queries.ListQueryClientIDs(ctx, &contextID)
	require.NoError(t, err)
	assert.ElementsMatch(t, []string{clientOneID, clientTwoID}, clients)

	jsonID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeJsonExtractor))
	require.NoError(t, err)
	_, err = queries.AddLatestQueryVersion(ctx, jsonID)
	require.NoError(t, err)
	err = queries.AddRequiredQuery(ctx, &repository.AddRequiredQueryParams{
		Queryid:         jsonID,
		Requiredqueryid: contextID,
		Addedversion:    1,
	})
	require.NoError(t, err)
	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Clientid:     clientOneID,
		Queryid:      jsonID,
		Addedversion: 1,
		Name:         "example_key",
	})
	require.NoError(t, err)

	clients, err = queries.ListQueryClientIDs(ctx, &jsonID)
	require.NoError(t, err)
	assert.ElementsMatch(t, []string{clientOneID}, clients)
}