refactor: Add DecryptToken method

keboola · Dec 17, 2024 · e307174 · e307174
1 parent fa355ca
commit e307174
Show file tree

Hide file tree

Showing 5 changed files with 23 additions and 33 deletions.
diff --git a/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/file.go b/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/file.go
@@ -171,14 +171,9 @@ func (b *Bridge) importFile(ctx context.Context, file plugin.File, stats statist
 	metadata := cloudencrypt.Metadata{"sink": file.SinkKey.String()}
 
 	// Decrypt token
-	var token keboola.Token
-	if existingToken.EncryptedToken != nil {
-		token, err = b.tokenEncryptor.Decrypt(ctx, existingToken.EncryptedToken, metadata)
-		if err != nil {
-			return err
-		}
-	} else {
-		token = *existingToken.Token
+	token, err := existingToken.DecryptToken(ctx, b.tokenEncryptor, metadata)
+	if err != nil {
+		return err
 	}
 
 	// Authorized API

diff --git a/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/job.go b/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/job.go
@@ -84,14 +84,9 @@ func (b *Bridge) CleanJob(ctx context.Context, job model.Job) (err error, delete
 	metadata := cloudencrypt.Metadata{"sink": job.SinkKey.String()}
 
 	// Decrypt token
-	var token keboola.Token
-	if existingToken.EncryptedToken != nil {
-		token, err = b.tokenEncryptor.Decrypt(ctx, existingToken.EncryptedToken, metadata)
-		if err != nil {
-			return err, false
-		}
-	} else {
-		token = *existingToken.Token
+	token, err := existingToken.DecryptToken(ctx, b.tokenEncryptor, metadata)
+	if err != nil {
+		return err, false
 	}
 
 	// Get job details from storage API

diff --git a/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/slice.go b/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/slice.go
@@ -40,14 +40,9 @@ func (b *Bridge) uploadSlice(ctx context.Context, volume *diskreader.Volume, sli
 	metadata := cloudencrypt.Metadata{"sink": slice.SinkKey.String()}
 
 	// Decrypt token
-	var token keboola.Token
-	if existingToken.EncryptedToken != nil {
-		token, err = b.tokenEncryptor.Decrypt(ctx, existingToken.EncryptedToken, metadata)
-		if err != nil {
-			return err
-		}
-	} else {
-		token = *existingToken.Token
+	token, err := existingToken.DecryptToken(ctx, b.tokenEncryptor, metadata)
+	if err != nil {
+		return err
 	}
 
 	// Error when sending the event is not a fatal error

diff --git a/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/token.go b/internal/pkg/service/stream/sink/type/tablesink/keboola/bridge/token.go
@@ -85,15 +85,9 @@ func (b *Bridge) tokenForSink(ctx context.Context, now time.Time, sink definitio
 		}
 
 		// Decrypt token
-		var token keboola.Token
-		if existingToken.EncryptedToken != nil {
-			token, err = b.tokenEncryptor.Decrypt(ctx, existingToken.EncryptedToken, metadata)
-			if err != nil {
-				return keboola.Token{}, err
-			}
-		} else {
-			// Backwards compatibility, should be dropped later
-			token = *existingToken.Token
+		token, err := existingToken.DecryptToken(ctx, b.tokenEncryptor, metadata)
+		if err != nil {
+			return keboola.Token{}, err
 		}
 
 		// Operation is not called from the API and there is a token in the database, so we are using the token.

diff --git a/internal/pkg/service/stream/sink/type/tablesink/keboola/token.go b/internal/pkg/service/stream/sink/type/tablesink/keboola/token.go
@@ -1,7 +1,10 @@
 package keboola
 
 import (
+	"context"
+
 	"github.com/keboola/go-client/pkg/keboola"
+	"github.com/keboola/go-cloud-encrypt/pkg/cloudencrypt"
 
 	"github.com/keboola/keboola-as-code/internal/pkg/service/stream/definition/key"
 )
@@ -21,3 +24,11 @@ func (token Token) ID() string {
 	}
 	return token.Token.ID
 }
+
+func (token Token) DecryptToken(ctx context.Context, encryptor *cloudencrypt.GenericEncryptor[keboola.Token], metadata cloudencrypt.Metadata) (keboola.Token, error) {
+	if token.EncryptedToken != nil {
+		return encryptor.Decrypt(ctx, token.EncryptedToken, metadata)
+	}
+
+	return *token.Token, nil
+}