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

package repository_test

import (
	"context"
	"os"
	"path"
	"queryorchestration/internal/database"
	"queryorchestration/internal/database/repository"
	"queryorchestration/internal/serviceconfig"
	"queryorchestration/internal/test"
	"testing"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/stretchr/testify/assert"
)

func TestQueries(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping long test in short mode")
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	test.SetCfgProvider(t, cfg)
	cfg.SetBasePath(path.Join(os.Getenv("PWD"), "../../.."))
	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
		Cfg:           cfg,
		RunMigrations: true,
	})
	defer cleanup()

	queries := cfg.GetDBQueries()

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	jsonQueryConfig, err := queries.GetQueryConfig(ctx, &repository.GetQueryConfigParams{
		Queryid:      jsonQueryID,
		Addedversion: jsonQuery.Activeversion,
	})
	assert.NoError(t, err)
	assert.Equal(t, jsonConfig, jsonQueryConfig.Config)

	jsonQueryConfig, err = queries.GetQueryConfig(ctx, &repository.GetQueryConfigParams{
		Queryid:      jsonQueryID,
		Addedversion: 2,
	})
	assert.NoError(t, err)
	assert.Equal(t, []byte(`{"second": "key"}`), jsonQueryConfig.Config)

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

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

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

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

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

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

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

func TestQueryDependencyTree(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping long test in short mode")
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	test.SetCfgProvider(t, cfg)
	cfg.SetBasePath(path.Join(os.Getenv("PWD"), "../../.."))
	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
		Cfg:           cfg,
		RunMigrations: true,
	})
	defer cleanup()

	queries := cfg.GetDBQueries()

	clientID, err := queries.CreateClient(ctx, "example_client")
	assert.NoError(t, err)
	jobID, err := queries.CreateJob(ctx, clientID)
	assert.NoError(t, err)
	collID, err := queries.CreateCollector(ctx, jobID)
	assert.NoError(t, err)
	_, err = queries.AddLatestCollectorVersion(ctx, collID)
	assert.NoError(t, err)
	err = queries.SetActiveCollectorVersion(ctx, &repository.SetActiveCollectorVersionParams{
		Collectorid: collID,
		Versionid:   1,
	})
	assert.NoError(t, err)
	docID, err := queries.CreateDocument(ctx, &repository.CreateDocumentParams{
		Jobid: jobID,
		Hash:  "sample",
	})
	assert.NoError(t, err)

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

	dependents, err := queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: contextQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)

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

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: jsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)

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

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: jsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: contextQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)

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

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: jsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: contextQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jsonQueryID}, dependents)

	secondJsonQueryID, err := queries.CreateQuery(ctx, repository.Querytype(repository.QuerytypeJsonExtractor))
	assert.NoError(t, err)
	_, err = queries.AddLatestQueryVersion(ctx, secondJsonQueryID)
	assert.NoError(t, err)

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: jsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: contextQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jsonQueryID}, dependents)

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

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: secondJsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: jsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: contextQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jsonQueryID}, dependents)

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

	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: secondJsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: jsonQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{secondJsonQueryID}, dependents)
	dependents, err = queries.ListQueryDirectDependentsByDocumentID(ctx, &repository.ListQueryDirectDependentsByDocumentIDParams{
		ID:          docID,
		Requiredids: contextQueryID,
	})
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jsonQueryID}, dependents)

	isdependent, err := queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Requiredqueryid: jsonQueryID,
		ID:              []pgtype.UUID{contextQueryID},
	})
	assert.NoError(t, err)
	assert.False(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Requiredqueryid: jsonQueryID,
		ID:              []pgtype.UUID{secondJsonQueryID},
	})
	assert.NoError(t, err)
	assert.True(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Requiredqueryid: jsonQueryID,
		ID:              []pgtype.UUID{jsonQueryID},
	})
	assert.NoError(t, err)
	assert.True(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Requiredqueryid: secondJsonQueryID,
		ID:              []pgtype.UUID{jsonQueryID, contextQueryID},
	})
	assert.NoError(t, err)
	assert.False(t, isdependent)

	isdependent, err = queries.IsQueryInDependencyTree(ctx, &repository.IsQueryInDependencyTreeParams{
		Requiredqueryid: contextQueryID,
		ID:              []pgtype.UUID{jsonQueryID, secondJsonQueryID},
	})
	assert.NoError(t, err)
	assert.True(t, isdependent)
}

func TestQueriesList(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping long test in short mode")
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	test.SetCfgProvider(t, cfg)
	cfg.SetBasePath(path.Join(os.Getenv("PWD"), "../../.."))
	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
		Cfg:           cfg,
		RunMigrations: true,
	})
	defer cleanup()

	queries := cfg.GetDBQueries()

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

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

	qs, err := queries.ListQueries(ctx)
	assert.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:   []pgtype.UUID{},
		},
		{
			ID:            contextQueryID,
			Type:          repository.QuerytypeContextFull,
			Activeversion: 0,
			Latestversion: 0,
			Config:        nil,
			Requiredids:   []pgtype.UUID{},
		},
	}, qs)

	qs, err = queries.ListQueriesById(ctx, []pgtype.UUID{jsonQueryID})
	assert.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:   []pgtype.UUID{},
		},
	}, qs)
}

func TestListQueryJobs(t *testing.T) {
	if testing.Short() {
		t.Skip("Skipping long test in short mode")
	}
	ctx := context.Background()

	cfg := &serviceconfig.BaseConfig{}
	test.SetCfgProvider(t, cfg)
	cfg.SetBasePath(path.Join(os.Getenv("PWD"), "../../.."))
	_, cleanup := test.CreateDB(t, ctx, &test.CreateDatabaseConfig{
		Cfg:           cfg,
		RunMigrations: true,
	})
	defer cleanup()

	queries := cfg.GetDBQueries()

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

	jobs, err := queries.ListQueryJobIDs(ctx, contextID)
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{}, jobs)

	clientOneID, err := queries.CreateClient(ctx, "example_client")
	assert.NoError(t, err)
	jobOneID, err := queries.CreateJob(ctx, clientOneID)
	assert.NoError(t, err)
	collOneID, err := queries.CreateCollector(ctx, jobOneID)
	assert.NoError(t, err)
	_, err = queries.AddLatestCollectorVersion(ctx, collOneID)
	assert.NoError(t, err)
	err = queries.SetActiveCollectorVersion(ctx, &repository.SetActiveCollectorVersionParams{
		Collectorid: collOneID,
		Versionid:   1,
	})
	assert.NoError(t, err)
	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Collectorid:  collOneID,
		Queryid:      contextID,
		Addedversion: 1,
		Name:         "example_key",
	})
	assert.NoError(t, err)

	jobs, err = queries.ListQueryJobIDs(ctx, contextID)
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jobOneID}, jobs)

	jobTwoID, err := queries.CreateJob(ctx, clientOneID)
	assert.NoError(t, err)
	collTwoID, err := queries.CreateCollector(ctx, jobTwoID)
	assert.NoError(t, err)
	_, err = queries.AddLatestCollectorVersion(ctx, collTwoID)
	assert.NoError(t, err)
	err = queries.SetActiveCollectorVersion(ctx, &repository.SetActiveCollectorVersionParams{
		Collectorid: collTwoID,
		Versionid:   1,
	})
	assert.NoError(t, err)
	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Collectorid:  collTwoID,
		Queryid:      contextID,
		Addedversion: 1,
		Name:         "example_key",
	})
	assert.NoError(t, err)

	jobs, err = queries.ListQueryJobIDs(ctx, contextID)
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jobOneID, jobTwoID}, jobs)

	clientTwoID, err := queries.CreateClient(ctx, "example_client_two")
	assert.NoError(t, err)
	jobThreeID, err := queries.CreateJob(ctx, clientTwoID)
	assert.NoError(t, err)
	collThreeID, err := queries.CreateCollector(ctx, jobThreeID)
	assert.NoError(t, err)
	_, err = queries.AddLatestCollectorVersion(ctx, collThreeID)
	assert.NoError(t, err)
	err = queries.SetActiveCollectorVersion(ctx, &repository.SetActiveCollectorVersionParams{
		Collectorid: collThreeID,
		Versionid:   1,
	})
	assert.NoError(t, err)
	err = queries.AddCollectorQuery(ctx, &repository.AddCollectorQueryParams{
		Collectorid:  collThreeID,
		Queryid:      contextID,
		Addedversion: 1,
		Name:         "example_key",
	})
	assert.NoError(t, err)

	jobs, err = queries.ListQueryJobIDs(ctx, contextID)
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jobOneID, jobTwoID, jobThreeID}, jobs)

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

	jobs, err = queries.ListQueryJobIDs(ctx, jsonID)
	assert.NoError(t, err)
	assert.ElementsMatch(t, []pgtype.UUID{jobOneID}, jobs)
}