Merge pull request #26 from tilezen/add-requester-pays-support

Add support for requester pays S3 requests
tilezen · Sep 2, 2023 · 9a48755 · 9a48755
2 parents 19c3db1 + 69032c4
commit 9a48755
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 8 deletions.
diff --git a/cmd/build/main.go b/cmd/build/main.go
@@ -57,6 +57,7 @@ func main() {
 	formatStr := flag.String("format", "mvt", "(For metatile generator) The format of the tile inside the metatile to extract.")
 	pathTemplateStr := flag.String("path-template", "", "(For metatile, tapalcatl2 generator) The template to use for the path part of the S3 path to the t2 archive.")
 	bucketStr := flag.String("bucket", "", "(For metatile, tapalcatl2 generator) The name of the S3 bucket to request t2 archives from.")
+	requesterPays := flag.Bool("requester-pays", false, "(For metatile, tapalcatl2 generator) Whether to make S3 requests with requester pays enabled.")
 	materializedZoomsStr := flag.String("materialized-zooms", "", "(For tapalcatl2 generator) Specifies the materialized zooms for t2 archives.")
 	flag.Parse()
 
@@ -181,7 +182,7 @@ func main() {
 		metatileSize := uint(8)
 		maxDetailZoom := maptile.Zoom(13)
 
-		jobCreator, err = tilepack.NewMetatileJobGenerator(*bucketStr, *pathTemplateStr, *layerNameStr, *formatStr, metatileSize, maxDetailZoom, zooms, bounds)
+		jobCreator, err = tilepack.NewMetatileJobGenerator(*bucketStr, *requesterPays, *pathTemplateStr, *layerNameStr, *formatStr, metatileSize, maxDetailZoom, zooms, bounds)
 	case "tapalcatl2":
 		if *bucketStr == "" {
 			log.Fatalf("Bucket name is required")
@@ -209,7 +210,7 @@ func main() {
 			materializedZooms[i] = maptile.Zoom(z)
 		}
 
-		jobCreator, err = tilepack.NewTapalcatl2JobGenerator(*bucketStr, *pathTemplateStr, *layerNameStr, materializedZooms, zooms, bounds)
+		jobCreator, err = tilepack.NewTapalcatl2JobGenerator(*bucketStr, *requesterPays, *pathTemplateStr, *layerNameStr, materializedZooms, zooms, bounds)
 	default:
 		log.Fatalf("Unknown job generator type %s", *generatorStr)
 	}

diff --git a/tilepack/metatile_job_creator.go b/tilepack/metatile_job_creator.go
@@ -28,7 +28,7 @@ func log2Uint(size uint) uint {
 	return uint(math.Log2(float64(size)))
 }
 
-func NewMetatileJobGenerator(bucket string, pathTemplate string, layerName string, format string, metatileSize uint, maxDetailZoom maptile.Zoom, zooms []maptile.Zoom, bounds orb.Bound) (JobGenerator, error) {
+func NewMetatileJobGenerator(bucket string, requesterPays bool, pathTemplate string, layerName string, format string, metatileSize uint, maxDetailZoom maptile.Zoom, zooms []maptile.Zoom, bounds orb.Bound) (JobGenerator, error) {
 	sess, err := session.NewSessionWithOptions(session.Options{
 		SharedConfigState: session.SharedConfigEnable,
 	})
@@ -47,6 +47,7 @@ func NewMetatileJobGenerator(bucket string, pathTemplate string, layerName strin
 	return &metatileJobGenerator{
 		s3Client:      downloader,
 		bucket:        bucket,
+		requesterPays: requesterPays,
 		pathTemplate:  pathTemplate,
 		layerName:     layerName,
 		metatileSize:  metatileSize,
@@ -60,6 +61,7 @@ func NewMetatileJobGenerator(bucket string, pathTemplate string, layerName strin
 type metatileJobGenerator struct {
 	s3Client      *s3manager.Downloader
 	bucket        string
+	requesterPays bool
 	pathTemplate  string
 	layerName     string
 	metatileSize  uint
@@ -78,10 +80,16 @@ func (x *metatileJobGenerator) CreateWorker() (func(id int, jobs chan *TileReque
 		for metaTileRequest := range jobs {
 			// Download the metatile archive zip to a byte buffer
 			compressedBytes := &aws.WriteAtBuffer{}
-			numBytes, err := x.s3Client.Download(compressedBytes, &s3.GetObjectInput{
+			input := &s3.GetObjectInput{
 				Bucket: aws.String(x.bucket),
 				Key:    aws.String(metaTileRequest.URL),
-			})
+			}
+
+			if x.requesterPays {
+				input.RequestPayer = aws.String("requester")
+			}
+
+			numBytes, err := x.s3Client.Download(compressedBytes, input)
 			if err != nil {
 				log.Printf("Unable to download item s3://%s/%s: %+v", x.bucket, metaTileRequest.URL, err)
 				continue

diff --git a/tilepack/t2_job_creator.go b/tilepack/t2_job_creator.go
@@ -18,7 +18,7 @@ import (
 	"github.com/paulmach/orb/maptile"
 )
 
-func NewTapalcatl2JobGenerator(bucket string, pathTemplate string, layerName string, materializedZooms []maptile.Zoom, zooms []maptile.Zoom, bounds orb.Bound) (JobGenerator, error) {
+func NewTapalcatl2JobGenerator(bucket string, requesterPays bool, pathTemplate string, layerName string, materializedZooms []maptile.Zoom, zooms []maptile.Zoom, bounds orb.Bound) (JobGenerator, error) {
 	sess, err := session.NewSessionWithOptions(session.Options{
 		SharedConfigState: session.SharedConfigEnable,
 	})
@@ -31,6 +31,7 @@ func NewTapalcatl2JobGenerator(bucket string, pathTemplate string, layerName str
 	return &tapalcatl2JobGenerator{
 		s3Client:          downloader,
 		bucket:            bucket,
+		requesterPays:     requesterPays,
 		pathTemplate:      pathTemplate,
 		layerName:         layerName,
 		materializedZooms: materializedZooms,
@@ -42,6 +43,7 @@ func NewTapalcatl2JobGenerator(bucket string, pathTemplate string, layerName str
 type tapalcatl2JobGenerator struct {
 	s3Client          *s3manager.Downloader
 	bucket            string
+	requesterPays     bool
 	pathTemplate      string
 	layerName         string
 	materializedZooms []maptile.Zoom
@@ -63,10 +65,16 @@ func (x *tapalcatl2JobGenerator) CreateWorker() (func(id int, jobs chan *TileReq
 		for request := range jobs {
 			// Download the Tapalcatl2 archive zip to a byte buffer
 			compressedBytes := &aws.WriteAtBuffer{}
-			numBytes, err := x.s3Client.Download(compressedBytes, &s3.GetObjectInput{
+			input := &s3.GetObjectInput{
 				Bucket: aws.String(x.bucket),
 				Key:    aws.String(request.URL),
-			})
+			}
+
+			if x.requesterPays {
+				input.RequestPayer = aws.String("requester")
+			}
+
+			numBytes, err := x.s3Client.Download(compressedBytes, input)
 			if err != nil {
 				log.Fatalf("Unable to download item %s: %+v", request.URL, err)
 			}