query-orchestration/cmd/queryAPI/main.go

// This API currently manages all user entities to be managed. This allows users of the service to change required entities and processes.
//
// The entities in question:
//
// - Client - A client is a single client of the Doczy project, all configurations and clients for a client will be in reference to the respective client.
//
// - Collector - A collector specifies document processing settings for a given client.
//
// - Sample - A sample is a subset of documents to be used for testing. When it is active a client will only process documents in the sample.
//
// - Export - An export is a manually started process which outputs the results for a client to a given location.
//
// Note: Query functionality has been removed from this API. See remove_query_plan.md for details.
package main

import (
	"context"
	"fmt"
	"log/slog"
	"os"

	queryapi "queryorchestration/api/queryAPI"
	"queryorchestration/internal/backgroundtask"
	"queryorchestration/internal/bot"
	"queryorchestration/internal/client"
	clientupdate "queryorchestration/internal/client/update"
	"queryorchestration/internal/cognitoauth"
	"queryorchestration/internal/collector"
	collectorset "queryorchestration/internal/collector/set"
	"queryorchestration/internal/customschema"
	"queryorchestration/internal/document"
	documentbatch "queryorchestration/internal/document/batch"
	documentdownload "queryorchestration/internal/document/download"
	documentupload "queryorchestration/internal/document/upload"
	"queryorchestration/internal/eula"
	"queryorchestration/internal/export"
	"queryorchestration/internal/fieldextraction"
	"queryorchestration/internal/folder"
	"queryorchestration/internal/label"
	"queryorchestration/internal/server/api"
	"queryorchestration/internal/serviceconfig"
	awsc "queryorchestration/internal/serviceconfig/aws"
	"queryorchestration/internal/serviceconfig/build"
	"queryorchestration/internal/serviceconfig/chatbot"
	"queryorchestration/internal/serviceconfig/objectstore"
	"queryorchestration/internal/serviceconfig/queue/clientsync"
	"queryorchestration/internal/uisettings"
	"queryorchestration/internal/usermanagement"

	"github.com/aws/aws-sdk-go-v2/service/cognitoidentityprovider"
	"github.com/getkin/kin-openapi/openapi3"
	_ "github.com/lib/pq"
)

type QueryAPIConfig struct {
	api.BaseConfig
	clientsync.ClientSyncConfig
	objectstore.ObjectStoreConfig
	// Chatbot configuration is loaded from CHATBOT_* env vars and
	// validated immediately after InitializeConfig. See plan §8 for the
	// authoritative spec; the embedding lets serviceconfig.InitializeConfig
	// walk the env tags as part of the standard parse pass.
	chatbot.ChatbotConfig
	BackgroundRunner *backgroundtask.Runner
}

func main() {
	// Print version information before any environment checks
	build.PrintVersionInfo("queryAPI")

	ctx := context.Background()

	cfg := &QueryAPIConfig{}

	cfg.RegisterHandlersFunc = func() (*openapi3.T, error) {
		return registerQueryAPIHandlers(cfg)
	}

	// Both of these operations (InitializeConfig and InitializeAuthConfig)
	// would be better off in one function but since there are services
	// that need the config but do not use auth we will keep them separate for now.

	// This must be done before the rbac.InitializeAuthProvider
	errInitializingConfig := serviceconfig.InitializeConfig(cfg)
	if errInitializingConfig != nil {
		slog.Error(errInitializingConfig.Error())
		os.Exit(1)
	}

	// Chatbot config validation runs after InitializeConfig and before
	// auth init / server.Listen() per plan §8. Range checks on the
	// numeric fields are not enforceable by env tags alone, so a
	// dedicated Validate is invoked here and bails the process on
	// invalid input the same way other config-validation failures do.
	if err := cfg.ChatbotConfig.Validate(); err != nil {
		slog.Error("chatbot config validation failed", "error", err)
		os.Exit(1)
	}

	// Authentication specific config.
	cfg.BaseURL = fmt.Sprintf("%s:%d", cfg.BaseURL, cfg.Port)
	errorInitializingAuthConfig := cfg.InitializeAuthConfig(cfg.BaseURL, cfg.GetLogger())
	fmt.Printf("base url after InitializeAuthConfig: %s\n", cfg.BaseURL)
	// join initErr with errorInitializingAuthConfig
	if errorInitializingAuthConfig != nil {
		slog.Error(errorInitializingAuthConfig.Error())
		os.Exit(1)
	}

	// Validate Permit.io configuration before starting server
	if err := cognitoauth.ValidatePermitIOConfiguration(); err != nil {
		slog.Error("Permit.io configuration validation failed", "error", err)
		os.Exit(1)
	}

	// Create Cognito SDK client for environment attribute registration and middleware use.
	// This client is created once and reused for the lifetime of the process.
	awsCfg, err := awsc.GetAWSConfig(ctx)
	if err != nil {
		slog.Error("Failed to load AWS config for Cognito client", "error", err)
		os.Exit(1)
	}
	cognitoClient := cognitoidentityprovider.NewFromConfig(awsCfg)

	// Ensure the custom:environment_id attribute is registered on the Cognito user pool.
	// This is safe to call on every startup (idempotent). Failure is non-fatal because
	// the attribute may already exist or the environment may not support this API (e.g. LocalStack).
	if err := usermanagement.EnsureEnvironmentIDAttribute(ctx, cognitoClient, cfg.GetAuthUserPoolID(), cfg.GetLogger()); err != nil {
		cfg.GetLogger().Warn("Failed to ensure Cognito environment_id attribute", "error", err)
	} else {
		cfg.GetLogger().Info("Cognito custom attribute custom:environment_id is registered")
	}

	// Initialize S3 client before creating the server (needed for background worker)
	err = cfg.SetStoreClient(ctx)
	if err != nil {
		slog.Error(err.Error())
		os.Exit(1)
	}

	server, err := api.New(ctx, cfg)
	if err != nil {
		slog.Error(err.Error())
		os.Exit(1)
	}

	server.Listen()
}

// registerQueryAPIHandlers builds the queryAPI service graph, registers
// the OpenAPI-driven echo router, and returns the swagger doc the
// server uses for request validation. Extracted from main() so the
// per-package average-complexity linter does not flag main().
func registerQueryAPIHandlers(cfg *QueryAPIConfig) (*openapi3.T, error) {
	services, err := buildQueryAPIServices(cfg)
	if err != nil {
		return nil, err
	}
	cons := queryapi.NewControllers(services, cfg)
	queryapi.RegisterHandlers(cfg.Router, cons)
	swagger, err := queryapi.GetSwagger()
	if err != nil {
		return nil, fmt.Errorf("error loading swagger: %w", err)
	}
	return swagger, nil
}

// buildQueryAPIServices constructs every service that backs the
// generated controller. Bot construction validates the chatbot config
// a second time so non-main callers (tests, alternate entry points)
// inherit the same gate as main().
func buildQueryAPIServices(cfg *QueryAPIConfig) (*queryapi.Services, error) {
	exp := export.New()
	col := collector.New(cfg)
	colupdate := collectorset.New(cfg, &collectorset.Services{
		Collector: col,
	})
	cli := client.New(cfg)
	cliUpdate := clientupdate.New(cfg, &clientupdate.Services{
		Client: cli,
	})
	doc := document.New(cfg)
	docup := documentupload.New(cfg)
	docdownload := documentdownload.New(cfg)
	docbatch := documentbatch.New(cfg)
	fieldext := fieldextraction.New(cfg)
	fld := folder.New(cfg)
	lbl := label.New(cfg)
	eul := eula.New(cfg)
	uiSvc := uisettings.New(cfg)
	customSchema := customschema.New(cfg, &customschema.SchemaValidator{}, customschema.NewSlogAuditSink(cfg.GetLogger()))
	botSvc, err := bot.New(cfg.ChatbotConfig, cfg)
	if err != nil {
		return nil, fmt.Errorf("error constructing bot service: %w", err)
	}

	return &queryapi.Services{
		Export:           exp,
		Collector:        col,
		CollectorSet:     colupdate,
		Client:           cli,
		ClientUpdate:     cliUpdate,
		Document:         doc,
		DocumentUpload:   docup,
		DocumentDownload: docdownload,
		DocumentBatch:    docbatch,
		UISettings:       uiSvc,
		FieldExtraction:  fieldext,
		CustomSchema:     customSchema,
		Folder:           fld,
		Label:            lbl,
		Eula:             eul,
		Bot:              botSvc,
	}, nil
}