Repartition a glue table with the AWS Go SDK

glue_repartition.go

package main

import (
	"fmt"
	"log"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/glue"
)


func repartition(database, tableName, path string, newPartitions []string) {

	sess := session.Must(session.NewSession(&aws.Config{
		Region: aws.String("eu-west-1"),
	}))

	svc := glue.New(sess)

	partitions := getAllPartitions(svc, database, tableName)

	newPartitions = compare(newPartitions, partitions)

	if len(newPartitions) != 0 {
		fmt.Println("Missing Partitions: ", newPartitions)
		createPartitions(svc, newPartitions, database, tableName, path)
	}
	log.Println("Done")
}

func getAllPartitions(svc *glue.Glue, db, table string) []string {

	params := &glue.GetPartitionsInput{
		DatabaseName: aws.String(db),
		TableName:    aws.String(table),
	}

	var partitions []string
	err := svc.GetPartitionsPages(params,
		func(page *glue.GetPartitionsOutput, lastPage bool) bool {
			for _, p := range page.Partitions {
				partitions = append(partitions, aws.StringValueSlice(p.Values)...)
			}
			if lastPage {
				return true
			}
			return false
		})

	if err != nil {
    panic(err)
	}
	return partitions
}

func createPartitions(svc *glue.Glue, partitions []string, db, table, s3Path string) {
	t, _ := svc.GetTable(
		&glue.GetTableInput{
			DatabaseName: aws.String(db),
			Name:         aws.String(table),
		},
	)
	for _, v := range partitions {
		t.Table.StorageDescriptor.Location = aws.String(s3Path + v)
		fmt.Println("Creating partition: ", v)
		_, err := svc.CreatePartition(
			&glue.CreatePartitionInput{
				DatabaseName: aws.String(db),
				TableName:    aws.String(table),
				PartitionInput: &glue.PartitionInput{
					StorageDescriptor: t.Table.StorageDescriptor,
					Values:            []*string{&v},
				},
			},
		)

		if err != nil {
			fmt.Println("Error: ", err)
		}
	}
}

func compare(a, b []string) []string {
	for i := len(a) - 1; i >= 0; i-- {
		for _, v := range b {
			if a[i] == v {
				a = append(a[:i], a[i+1:]...)
				break
			}
		}
	}
	return a
}