From b0e05072caee4dc15d9b043aa8e861edcf2614bb Mon Sep 17 00:00:00 2001
From: Robin Tang <rtang.cs@gmail.com>
Date: Thu, 11 Jul 2024 11:43:50 -0700
Subject: [PATCH] Pass context in.

---
 lib/kafkalib/connection.go | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/lib/kafkalib/connection.go b/lib/kafkalib/connection.go
index 386207d28..d02880fda 100644
--- a/lib/kafkalib/connection.go
+++ b/lib/kafkalib/connection.go
@@ -83,7 +83,7 @@ func (c Connection) Dialer(ctx context.Context) (*kafka.Dialer, error) {
 	return dialer, nil
 }
 
-func (c Connection) Transport() (*kafka.Transport, error) {
+func (c Connection) Transport(ctx context.Context) (*kafka.Transport, error) {
 	transport := &kafka.Transport{
 		DialTimeout: 10 * time.Second,
 	}
@@ -100,7 +100,7 @@ func (c Connection) Transport() (*kafka.Transport, error) {
 			transport.TLS = &tls.Config{}
 		}
 	case AwsMskIam:
-		_awsCfg, err := awsCfg.LoadDefaultConfig(context.Background())
+		_awsCfg, err := awsCfg.LoadDefaultConfig(ctx)
 		if err != nil {
 			return nil, fmt.Errorf("failed to load AWS configuration: %w", err)
 		}