Merge branch 'main' of bitbucket.org:aarete/query-orchestration into feature/cognito

cmd/queryAPI/main.go

@@ -23,6 +23,7 @@ import (
 	"queryorchestration/internal/serviceconfig"
 
 	"queryorchestration/internal/client"
+	clientupdate "queryorchestration/internal/client/update"
 	"queryorchestration/internal/collector"
 	collectorset "queryorchestration/internal/collector/set"
 	"queryorchestration/internal/document"
@@ -61,6 +62,9 @@ func main() {
 			Collector: col,
 		})
 		cli := client.New(cfg)
+		cliUpdate := clientupdate.New(cfg, &clientupdate.Services{
+			Client: cli,
+		})
 		doc := document.New(cfg)
 		quetest := querytest.New(cfg, &querytest.Services{
 			Collector: col,
@@ -79,6 +83,7 @@ func main() {
 			QueryUpdate:  qupdate,
 			QueryTest:    quetest,
 			Client:       cli,
+			ClientUpdate: cliUpdate,
 			Document:     doc,
 		}
 

cmd/textractGenerator_test/textdetection.go

@@ -16,6 +16,8 @@ import (
 	"queryorchestration/internal/serviceconfig/objectstore"
 	"queryorchestration/internal/serviceconfig/queue/querysync"
 	"queryorchestration/internal/serviceconfig/textract"
+	"sync"
+	"time"
 
 	awstextract "github.com/aws/aws-sdk-go-v2/service/textract"
 	"github.com/aws/aws-sdk-go-v2/service/textract/types"
@@ -87,44 +89,60 @@ func main() {
 		os.Exit(1)
 	}
 
-	slog.Info("detecting text")
+	slog.Info("detecting text", "page_count", count)
 
+	startTime := time.Now()
 	allResults := map[string][]awstextract.AnalyzeDocumentOutput{}
 	allResults["base"] = make([]awstextract.AnalyzeDocumentOutput, count)
-	allResults["table"] = []awstextract.AnalyzeDocumentOutput{}
+	var wg sync.WaitGroup
 	for i := range count {
-		result, err := textSvc.GetBasePage(ctx, pdf, i)
-		if err != nil {
-			slog.Error("error detecting text", "error", err)
-			os.Exit(1)
-		}
+		wg.Add(1)
+		go func() {
+			slog.Info("getting page", "index", i)
 
+			result, err := textSvc.GetBasePage(ctx, pdf, i)
+			if err != nil {
+				slog.Error("error detecting text", "error", err)
+				os.Exit(1)
+			}
+
+			allResults["base"][i] = result
+			wg.Done()
+		}()
+	}
+	wg.Wait()
+
+	allResults["table"] = []awstextract.AnalyzeDocumentOutput{}
+	for i, result := range allResults["base"] {
+		hasTable := false
 		for _, block := range result.Blocks {
-			feature := textSvc.GetBlockFeature(block)
-			if feature == documenttext.FeatureTable {
-				pageBytes, err := pdf.GetPageAsPNG(ctx, i)
-				if err != nil {
-					slog.Error("error getting page", "error", err)
-					os.Exit(1)
-				}
-
-				res, err := cfg.GetTextractClient().AnalyzeDocument(ctx, &awstextract.AnalyzeDocumentInput{
-					Document: &types.Document{
-						Bytes: pageBytes,
-					},
-					FeatureTypes: []types.FeatureType{
-						types.FeatureTypeTables,
-					},
-				})
-				if err != nil {
-					slog.Error("error detecting table", "error", err)
-					os.Exit(1)
-				}
-				allResults["table"] = append(allResults["table"], *res)
+			if block.BlockType == types.BlockTypeLayoutTable {
+				hasTable = true
 			}
 		}
 
-		allResults["base"][i] = result
+		if hasTable {
+			slog.Info("getting table", "index", i)
+			pageBytes, err := pdf.GetPageAsPNG(ctx, i)
+			if err != nil {
+				slog.Error("error getting page", "error", err)
+				os.Exit(1)
+			}
+
+			res, err := cfg.GetTextractClient().AnalyzeDocument(ctx, &awstextract.AnalyzeDocumentInput{
+				Document: &types.Document{
+					Bytes: pageBytes,
+				},
+				FeatureTypes: []types.FeatureType{
+					types.FeatureTypeTables,
+				},
+			})
+			if err != nil {
+				slog.Error("error detecting table", "error", err)
+				os.Exit(1)
+			}
+			allResults["table"] = append(allResults["table"], *res)
+		}
 	}
 
 	generatedBase := "generated"
@@ -159,5 +177,5 @@ func main() {
 		os.Exit(1)
 	}
 
-	fmt.Println("Successfully generated " + outputFilename)
+	slog.Info("extracted text", "output", outputFilename, "time", time.Since(startTime))
 }