redriver.go 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340
  1. package redriver
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "strconv"
  8. "github.com/aws/aws-sdk-go-v2/aws"
  9. "github.com/aws/aws-sdk-go-v2/service/sqs"
  10. "github.com/aws/aws-sdk-go-v2/service/sqs/types"
  11. "github.com/fgm/izidic"
  12. "code.osinet.fr/fgm/sqs_demo/services"
  13. )
  14. const (
  15. // MessageSystemAttributeNameDeadLetterQueueSourceArn is an undocumented
  16. // types.Message attribute, used by the SQS console to support redrive.
  17. MessageSystemAttributeNameDeadLetterQueueSourceArn types.MessageSystemAttributeName = "DeadLetterQueueSourceArn"
  18. )
  19. type ItemsKeys struct {
  20. MessageID string
  21. ReceiptHandle string
  22. }
  23. // Redriver is a redrive-oriented facade in front of the sqs.Client API.
  24. type Redriver interface {
  25. ListQueues(ctx context.Context, prefix string) (QueueUrls []string, err error)
  26. GetDLQ(ctx context.Context, qName string) (QueueUrls string, err error)
  27. GetQueueInfo(ctx context.Context, qName string) (*QueueInfo, error)
  28. GetQueueItems(ctx context.Context, qName string) ([]Message, error)
  29. DeleteItems(ctx context.Context, qName string, itemsIDs []ItemsKeys) error
  30. Purge(ctx context.Context, qName string) error
  31. RedriveItems(ctx context.Context, qName string, itemIDs []ItemsKeys) error
  32. }
  33. type redriver struct {
  34. Wait int32
  35. io.Writer
  36. *sqs.Client
  37. }
  38. func (r *redriver) ListQueues(ctx context.Context, prefix string) (QueueUrls []string, err error) {
  39. lqi := &sqs.ListQueuesInput{
  40. MaxResults: aws.Int32(1000),
  41. NextToken: nil,
  42. QueueNamePrefix: aws.String(prefix),
  43. }
  44. lqo, err := r.Client.ListQueues(ctx, lqi)
  45. if err != nil {
  46. return nil, fmt.Errorf("listing queues: %w", err)
  47. }
  48. return lqo.QueueUrls, nil
  49. }
  50. func (*redriver) parseQueueInfoRedriveAllowPolicy(qName string, qao sqs.GetQueueAttributesOutput) (qrap *QueueInfoAttributesRedriveAllowPolicy, err error) {
  51. srap, ok := qao.Attributes[string(types.QueueAttributeNameRedriveAllowPolicy)]
  52. if !ok || srap == "" {
  53. return nil, nil
  54. }
  55. qrap = &QueueInfoAttributesRedriveAllowPolicy{}
  56. if err = json.Unmarshal([]byte(srap), qrap); err != nil {
  57. return nil, fmt.Errorf(
  58. "failed parsing redrive policy for queue %q %w", qName, err)
  59. }
  60. return qrap, nil
  61. }
  62. func (*redriver) parseQueueInfoRedrivePolicy(qName string, qao sqs.GetQueueAttributesOutput) (qrp *QueueInfoAttributesRedrivePolicy, err error) {
  63. srp, ok := qao.Attributes[string(types.QueueAttributeNameRedrivePolicy)]
  64. if !ok || srp == "" {
  65. return nil, nil
  66. }
  67. qrp = &QueueInfoAttributesRedrivePolicy{}
  68. if err = json.Unmarshal([]byte(srp), qrp); err != nil {
  69. return qrp, fmt.Errorf(
  70. "failed parsing redrive policy for queue %q %w", qName, err)
  71. }
  72. if _, err = URLFromARNString(qrp.DeadLetterTargetARN); err != nil {
  73. return qrp, fmt.Errorf(
  74. "failed converting queue %q ARN to URL: %w", qName, err)
  75. }
  76. return qrp, nil
  77. }
  78. func (r *redriver) GetDLQ(ctx context.Context, qName string) (QueueUrl string, err error) {
  79. qui := &sqs.GetQueueUrlInput{QueueName: &qName}
  80. qu, err := r.GetQueueUrl(ctx, qui)
  81. if err != nil {
  82. return "", fmt.Errorf("failed getting URL for queue %q: %w", qName, err)
  83. }
  84. qai := &sqs.GetQueueAttributesInput{
  85. QueueUrl: qu.QueueUrl,
  86. AttributeNames: []types.QueueAttributeName{types.QueueAttributeNameRedrivePolicy},
  87. }
  88. qao, err := r.Client.GetQueueAttributes(ctx, qai)
  89. if err != nil {
  90. return "", fmt.Errorf("failed getting DLQ for queue %q: %w", qName, err)
  91. }
  92. if qao == nil {
  93. return "", fmt.Errorf("redrive policy for queue %q is empty", qName)
  94. }
  95. qrp, err := r.parseQueueInfoRedrivePolicy(qName, *qao)
  96. if err != nil {
  97. return "", fmt.Errorf("failed parsing redrive policy for queue %q: %w", qName, err)
  98. }
  99. if qrp == nil {
  100. return "", nil // Queue has no DLQ: this is not an error
  101. }
  102. qURL, _ := URLFromARNString(qrp.DeadLetterTargetARN) // Already checked in parseQueueInfoRedrivePolicy
  103. return qURL, nil
  104. }
  105. func (r *redriver) GetQueueInfo(ctx context.Context, qName string) (*QueueInfo, error) {
  106. qui := &sqs.GetQueueUrlInput{QueueName: &qName}
  107. qu, err := r.GetQueueUrl(ctx, qui)
  108. if err != nil {
  109. return nil, fmt.Errorf("failed getting URL for queue %q: %w", qName, err)
  110. }
  111. if qu.QueueUrl == nil {
  112. return nil, fmt.Errorf("URL for queue %q is empty", qName)
  113. }
  114. qai := &sqs.GetQueueAttributesInput{
  115. QueueUrl: qu.QueueUrl,
  116. AttributeNames: []types.QueueAttributeName{types.QueueAttributeNameAll},
  117. }
  118. qao, err := r.Client.GetQueueAttributes(ctx, qai)
  119. if err != nil {
  120. return nil, fmt.Errorf("failed getting all attributes for queue %q: %w", qName, err)
  121. }
  122. if qao == nil {
  123. return nil, fmt.Errorf("no attributes returned for queue %q", qName)
  124. }
  125. qi, err := r.parseQueueInfoAttributes(qName, qu, qao)
  126. if err != nil {
  127. return nil, err
  128. }
  129. return qi, err
  130. }
  131. func (r *redriver) parseQueueInfoAttributes(qName string, qu *sqs.GetQueueUrlOutput, qao *sqs.GetQueueAttributesOutput) (*QueueInfo, error) {
  132. var (
  133. qi *QueueInfo
  134. errCount int
  135. anom, anomd, anomnv int64
  136. created, changed int64 // timestamps
  137. delay, max, retention, vto, wait int64
  138. )
  139. u := *qu.QueueUrl
  140. a, err := ARNFromURL(u)
  141. if err != nil {
  142. return nil, fmt.Errorf("ARN for queue %q cannot be parsed from URL %q: %w", qName, u, err)
  143. }
  144. for _, field := range []struct {
  145. name types.QueueAttributeName
  146. value *int64
  147. sDef string
  148. def int64
  149. }{
  150. {types.QueueAttributeNameApproximateNumberOfMessages, &anom, "-1", -1},
  151. {types.QueueAttributeNameApproximateNumberOfMessagesDelayed, &anomd, "-1", -1},
  152. {types.QueueAttributeNameApproximateNumberOfMessagesNotVisible, &anomnv, "-1", -1},
  153. {types.QueueAttributeNameCreatedTimestamp, &created, "0", 0},
  154. {types.QueueAttributeNameDelaySeconds, &delay, "-1", -1},
  155. {types.QueueAttributeNameLastModifiedTimestamp, &changed, "0", 0},
  156. {types.QueueAttributeNameMaximumMessageSize, &max, "262144", 1 << (8 + 10)}, // 256ko
  157. {types.QueueAttributeNameMessageRetentionPeriod, &max, "1209600", 14 * 24 * 60 * 60}, // 2 weeks
  158. {types.QueueAttributeNameReceiveMessageWaitTimeSeconds, &wait, "0", 0}, // short polling
  159. {types.QueueAttributeNameVisibilityTimeout, &vto, "0", 0}, // short polling
  160. } {
  161. s, ok := qao.Attributes[string(field.name)]
  162. if !ok {
  163. errCount++
  164. s = field.sDef
  165. }
  166. n, err := strconv.Atoi(s)
  167. if err != nil {
  168. errCount++
  169. }
  170. *field.value = int64(n)
  171. }
  172. qrp, err := r.parseQueueInfoRedrivePolicy(qName, *qao)
  173. if err != nil {
  174. return nil, fmt.Errorf("failed parsing redrive policy for queue %q: %w", qName, err)
  175. }
  176. qrap, err := r.parseQueueInfoRedriveAllowPolicy(qName, *qao)
  177. if err != nil {
  178. return nil, fmt.Errorf("failed parsing redrive allow policy for queue %q: %w", qName, err)
  179. }
  180. qi = &QueueInfo{
  181. Name: qName,
  182. Url: u,
  183. Attributes: &QueueInfoAttributes{
  184. ApproximateNumberOfMessages: anom,
  185. ApproximateNumberOfMessagesDelayed: anomd,
  186. ApproximateNumberOfMessagesNotVisible: anomnv,
  187. CreatedTimestamp: created,
  188. DelaySeconds: delay,
  189. LastModifiedTimestamp: changed,
  190. MaximumMessageSize: max,
  191. MessageRetentionPeriod: retention,
  192. QueueARN: a.String(),
  193. ReceiveMessageWaitTimeSeconds: wait,
  194. VisibilityTimeout: vto,
  195. RedrivePolicy: qrp,
  196. RedriveAllowPolicy: qrap,
  197. },
  198. }
  199. return qi, nil
  200. }
  201. func (*redriver) parseReceivedAttributes(msg types.Message) (*SystemMessageAttributes, error) {
  202. var (
  203. afrt, arc, sent int64
  204. errCount int
  205. err error
  206. )
  207. for _, field := range []struct {
  208. name types.MessageSystemAttributeName
  209. value *int64
  210. sDef string
  211. def int64
  212. }{
  213. {types.MessageSystemAttributeNameApproximateFirstReceiveTimestamp, &afrt, "0", 0},
  214. {types.MessageSystemAttributeNameApproximateReceiveCount, &arc, "0", 0},
  215. {types.MessageSystemAttributeNameSentTimestamp, &sent, "0", 0},
  216. } {
  217. s, ok := msg.Attributes[string(field.name)]
  218. if !ok {
  219. errCount++
  220. s = field.sDef
  221. }
  222. n, err := strconv.Atoi(s)
  223. if err != nil {
  224. errCount++
  225. }
  226. *field.value = int64(n)
  227. }
  228. ma := SystemMessageAttributes{
  229. ApproximateFirstReceiveTimestamp: afrt,
  230. ApproximateReceiveCount: arc,
  231. DeadLetterQueueSourceARN: msg.Attributes[string(MessageSystemAttributeNameDeadLetterQueueSourceArn)],
  232. SenderId: msg.Attributes[string(types.MessageSystemAttributeNameSenderId)],
  233. SentTimestamp: sent,
  234. }
  235. if errCount > 0 {
  236. err = fmt.Errorf("encounted %d errors parsing system message attributes", errCount)
  237. return nil, err
  238. }
  239. return &ma, nil
  240. }
  241. func (r *redriver) GetQueueItems(ctx context.Context, qName string) ([]Message, error) {
  242. qui := &sqs.GetQueueUrlInput{QueueName: &qName}
  243. qu, err := r.GetQueueUrl(ctx, qui)
  244. if err != nil {
  245. return nil, fmt.Errorf("failed getting URL for queue %q: %w", qName, err)
  246. }
  247. if qu.QueueUrl == nil {
  248. return nil, fmt.Errorf("URL for queue %q is empty", qName)
  249. }
  250. rmi := sqs.ReceiveMessageInput{
  251. QueueUrl: qu.QueueUrl,
  252. AttributeNames: []types.QueueAttributeName{types.QueueAttributeNameAll},
  253. MaxNumberOfMessages: 10,
  254. MessageAttributeNames: []string{".*"},
  255. VisibilityTimeout: 0,
  256. WaitTimeSeconds: r.Wait,
  257. }
  258. rmo, err := r.Client.ReceiveMessage(ctx, &rmi)
  259. ms := make([]Message, 0, len(rmo.Messages))
  260. for _, m := range rmo.Messages {
  261. ma, err := r.parseReceivedAttributes(m)
  262. if err != nil {
  263. return nil, err
  264. }
  265. for _, field := range []struct {
  266. name string
  267. value any
  268. }{
  269. {"body", m.Body},
  270. {"md5 of body", m.MD5OfBody},
  271. {"md5 of message attributes", m.MD5OfMessageAttributes},
  272. {"message ID", m.MessageId},
  273. {"receipt handle", m.ReceiptHandle},
  274. } {
  275. if field.value == nil {
  276. return nil, fmt.Errorf("missing field %s on message", field.name)
  277. }
  278. }
  279. j, err := JSONableFromMessageAttributeValues(m.MessageAttributes)
  280. if err != nil {
  281. return nil, fmt.Errorf("failed decoding message boby: %w", err)
  282. }
  283. m := Message{
  284. Attributes: ma,
  285. Body: *m.Body,
  286. Md5OfBody: *m.MD5OfBody,
  287. Md5OfMessageAttributes: *m.MD5OfMessageAttributes,
  288. MessageAttributes: j,
  289. MessageId: *m.MessageId,
  290. ReceiptHandle: *m.ReceiptHandle,
  291. }
  292. ms = append(ms, m)
  293. }
  294. return ms, err
  295. }
  296. func (r *redriver) DeleteItems(ctx context.Context, qName string, itemsIDs []ItemsKeys) error {
  297. // TODO implement me
  298. panic("implement me")
  299. }
  300. func (r *redriver) Purge(ctx context.Context, qName string) error {
  301. // TODO implement me
  302. panic("implement me")
  303. }
  304. func (r *redriver) RedriveItems(ctx context.Context, qName string, itemIDs []ItemsKeys) error {
  305. // TODO implement me
  306. panic("implement me")
  307. }
  308. func RedriverService(dic *izidic.Container) (any, error) {
  309. cli := dic.MustService(services.SvcClient).(*sqs.Client)
  310. w := dic.MustParam(services.PWriter).(io.Writer)
  311. wait := int32(dic.MustParam(services.PWait).(int))
  312. return &redriver{
  313. Wait: wait,
  314. Writer: w,
  315. Client: cli,
  316. }, nil
  317. }