feat: sync credentials every 10 mins for aws in duckdb

BemiHQ · Jan 17, 2025 · 9a0807a · 9a0807a
1 parent 3377e26
commit 9a0807a
Showing 1 changed file with 46 additions and 24 deletions.
diff --git a/src/duckdb.go b/src/duckdb.go
@@ -7,6 +7,7 @@ import (
 	"os"
 	"regexp"
 	"strings"
+	"time"
 
 	_ "github.com/marcboeker/go-duckdb"
 )
@@ -20,8 +21,9 @@ var DEFAULT_BOOT_QUERIES = []string{
 }
 
 type Duckdb struct {
-	db     *sql.DB
-	config *Config
+	refreshQuit chan struct{}
+	db          *sql.DB
+	config      *Config
 }
 
 func NewDuckdb(config *Config) *Duckdb {
@@ -30,8 +32,9 @@ func NewDuckdb(config *Config) *Duckdb {
 	PanicIfError(err)
 
 	duckdb := &Duckdb{
-		db:     db,
-		config: config,
+		db:          db,
+		config:      config,
+		refreshQuit: make(chan struct{}),
 	}
 
 	bootQueries := readDuckdbInitFile(config)
@@ -45,26 +48,20 @@ func NewDuckdb(config *Config) *Duckdb {
 
 	switch config.StorageType {
 	case STORAGE_TYPE_S3:
-		switch config.Aws.CredentialsType {
-		case AWS_CREDENTIALS_TYPE_STATIC:
-			query := "CREATE SECRET aws_s3_secret (TYPE S3, KEY_ID '$accessKeyId', SECRET '$secretAccessKey', REGION '$region', ENDPOINT '$endpoint', SCOPE '$s3Bucket')"
-			_, err = duckdb.ExecContext(ctx, query, map[string]string{
-				"accessKeyId":     config.Aws.AccessKeyId,
-				"secretAccessKey": config.Aws.SecretAccessKey,
-				"region":          config.Aws.Region,
-				"endpoint":        config.Aws.S3Endpoint,
-				"s3Bucket":        "s3://" + config.Aws.S3Bucket,
-			})
-			PanicIfError(err)
-		case AWS_CREDENTIALS_TYPE_DEFAULT:
-			query := "CREATE SECRET aws_s3_secret (TYPE S3, PROVIDER CREDENTIAL_CHAIN, REGION '$region', ENDPOINT '$endpoint', SCOPE '$s3Bucket')"
-			_, err = duckdb.ExecContext(ctx, query, map[string]string{
-				"region":   config.Aws.Region,
-				"endpoint": config.Aws.S3Endpoint,
-				"s3Bucket": "s3://" + config.Aws.S3Bucket,
-			})
-			PanicIfError(err)
-		}
+		duckdb.setAwsCredentials(ctx)
+		ticker := time.NewTicker(10 * time.Minute)
+		time.Tick(10 * time.Minute)
+		go func() {
+			for {
+				select {
+				case <-ticker.C:
+					duckdb.setAwsCredentials(ctx)
+				case <-duckdb.refreshQuit:
+					ticker.Stop()
+					return
+				}
+			}
+		}()
 
 		if config.LogLevel == LOG_LEVEL_TRACE {
 			_, err = duckdb.ExecContext(ctx, "SET enable_http_logging=true", nil)
@@ -75,6 +72,30 @@ func NewDuckdb(config *Config) *Duckdb {
 	return duckdb
 }
 
+func (duckdb *Duckdb) setAwsCredentials(ctx context.Context) {
+	config := duckdb.config
+	switch config.Aws.CredentialsType {
+	case AWS_CREDENTIALS_TYPE_STATIC:
+		query := "CREATE OR REPLACE SECRET aws_s3_secret (TYPE S3, KEY_ID '$accessKeyId', SECRET '$secretAccessKey', REGION '$region', ENDPOINT '$endpoint', SCOPE '$s3Bucket')"
+		_, err := duckdb.ExecContext(ctx, query, map[string]string{
+			"accessKeyId":     config.Aws.AccessKeyId,
+			"secretAccessKey": config.Aws.SecretAccessKey,
+			"region":          config.Aws.Region,
+			"endpoint":        config.Aws.S3Endpoint,
+			"s3Bucket":        "s3://" + config.Aws.S3Bucket,
+		})
+		PanicIfError(err)
+	case AWS_CREDENTIALS_TYPE_DEFAULT:
+		query := "CREATE OR REPLACE SECRET aws_s3_secret (TYPE S3, PROVIDER CREDENTIAL_CHAIN, REGION '$region', ENDPOINT '$endpoint', SCOPE '$s3Bucket')"
+		_, err := duckdb.ExecContext(ctx, query, map[string]string{
+			"region":   config.Aws.Region,
+			"endpoint": config.Aws.S3Endpoint,
+			"s3Bucket": "s3://" + config.Aws.S3Bucket,
+		})
+		PanicIfError(err)
+	}
+}
+
 func (duckdb *Duckdb) ExecContext(ctx context.Context, query string, args map[string]string) (sql.Result, error) {
 	LogDebug(duckdb.config, "Querying DuckDB:", query, args)
 	return duckdb.db.ExecContext(ctx, replaceNamedStringArgs(query, args))
@@ -91,6 +112,7 @@ func (duckdb *Duckdb) PrepareContext(ctx context.Context, query string) (*sql.St
 }
 
 func (duckdb *Duckdb) Close() {
+	close(duckdb.refreshQuit)
 	duckdb.db.Close()
 }