diff --git a/main.go b/main.go index 9c2991d..7a11ed0 100644 --- a/main.go +++ b/main.go @@ -13,6 +13,7 @@ var ( config = kingpin.Flag("config", "Configuration file with S3 and Redshift credentials").Required().File() accessKey = kingpin.Flag("access-key", "AWS access key. Defaults to $AWS_ACCESS_KEY_ID").OverrideDefaultFromEnvar("AWS_ACCESS_KEY_ID").Required().String() secretKey = kingpin.Flag("secret-access-key", "AWS secret access key. Defaults to $AWS_SECRET_").OverrideDefaultFromEnvar("AWS_SECRET_ACCESS_KEY").Required().String() + sessionToken = kingpin.Flag("session-token", "AWS session token.").String() project = kingpin.Flag("project", "Google Project ID").OverrideDefaultFromEnvar("GCLOUD_PROJECT").Required().String() overwrite = kingpin.Flag("overwrite", "Overwrite BigQuery table").Bool() usePartitionedTables = kingpin.Flag("partition", "Create time partitioned BigQuery tables.").Bool() @@ -67,6 +68,7 @@ func main() { } awsConfig.S3.AccessKey = *accessKey awsConfig.S3.SecretKey = *secretKey + awsConfig.S3.SessionToken = *sessionToken config := &Configuration{ AWS: awsConfig, diff --git a/redshift/config.go b/redshift/config.go index 7e5d5c8..a8d7714 100644 --- a/redshift/config.go +++ b/redshift/config.go @@ -88,13 +88,17 @@ type RedshiftConnectionDetails struct { } type S3Configuration struct { - Bucket string `yaml:"bucket"` - AccessKey string `yaml:"access_key"` - SecretKey string `yaml:"secret_key"` + Bucket string `yaml:"bucket"` + AccessKey string `yaml:"access_key"` + SecretKey string `yaml:"secret_key"` + SessionToken string `yaml:"session_token"` } func (c S3Configuration) ToRedshiftCredentialsClause() string { - return fmt.Sprintf("aws_access_key_id=%s;aws_secret_access_key=%s", c.AccessKey, c.SecretKey) + if c.SessionToken == "" { + return fmt.Sprintf("aws_access_key_id=%s;aws_secret_access_key=%s", c.AccessKey, c.SecretKey) + } + return fmt.Sprintf("aws_access_key_id=%s;aws_secret_access_key=%s;token=%s", c.AccessKey, c.SecretKey, c.SessionToken) } func (c *RedshiftConnectionDetails) URLString() string {