Files
query-orchestration/internal/serviceconfig/queue/config.go
T
Michael McGuinness 752fb2e2c0 Merged in feature/chc_1 (pull request #120)
Text Extraction Clean Up + Parallelization

* merged_Call

* directchildren

* bitofimprovements

* go

* parallel

* fixfullsuite

* pondforconcurrency

* comments

* muchdone

* bitsofclean

* stabilisedtests

* snappy

* threadpooltests

* childelements

* testspassed
2025-04-25 17:02:50 +00:00

94 lines
2.3 KiB
Go

package queue
import (
"context"
"fmt"
"log/slog"
"time"
awsc "queryorchestration/internal/serviceconfig/aws"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/config"
"github.com/aws/aws-sdk-go-v2/service/sqs"
"github.com/aws/aws-sdk-go-v2/service/sqs/types"
)
type QueueConfig struct {
AWSEndpointUrlSQS string `env:"AWS_ENDPOINT_URL_SQS"`
QueueClient SQSClient
}
type ConfigProvider interface {
awsc.ConfigProvider
SetQueueClient(context.Context) error
SetQueueClientFromSQSCfg(SQSClient)
GetQueueClient() SQSClient
PingQueueByURL(context.Context, string) error
SendToQueue(context.Context, *SendParams) error
ReceiveFromQueue(context.Context, *ReceiveParams) (*sqs.ReceiveMessageOutput, error)
DeleteFromQueue(context.Context, *DeleteParams) error
SetSQSEndpoint(string)
GetSQSEndpoint() string
}
func (c *QueueConfig) GetQueueClient() SQSClient {
return c.QueueClient
}
func (c *QueueConfig) GetSQSEndpoint() string {
return c.AWSEndpointUrlSQS
}
func (c *QueueConfig) PingQueueByURL(ctx context.Context, url string) error {
timeout := time.After(30 * time.Second)
tick := time.NewTicker(2 * time.Second)
defer tick.Stop()
input := &sqs.GetQueueAttributesInput{
QueueUrl: aws.String(url),
AttributeNames: []types.QueueAttributeName{types.QueueAttributeNameApproximateNumberOfMessages},
}
for {
select {
case <-timeout:
return fmt.Errorf("unable to ping queue by url: %s", url)
case <-tick.C:
r, err := c.QueueClient.GetQueueAttributes(ctx, input)
if err == nil {
slog.Info("Successful queue ping", "url", url, "approx_msgs", r.Attributes[string(types.QueueAttributeNameApproximateNumberOfMessages)])
return nil
}
slog.Debug("Attempted queue ping", "url", url, "address", c.GetSQSEndpoint())
case <-ctx.Done():
return ctx.Err()
}
}
}
func (s *QueueConfig) SetQueueClientFromSQSCfg(client SQSClient) {
s.QueueClient = client
}
func (c *QueueConfig) SetQueueClient(ctx context.Context) error {
cfg, err := awsc.GetAWSConfigWithOpts(ctx, func(lo *config.LoadOptions) error {
if c.AWSEndpointUrlSQS != "" {
lo.BaseEndpoint = c.AWSEndpointUrlSQS
}
return nil
})
if err != nil {
return err
}
c.QueueClient = sqs.NewFromConfig(cfg)
return nil
}
func (c *QueueConfig) SetSQSEndpoint(e string) {
c.AWSEndpointUrlSQS = e
}