fix: Late setup for format and stage (cloudquery#14528)

disq · web-flow · commit dcf1a36aa08f · 2023-10-13T16:56:14.000Z
Shouldn't attempt to `CREATE STAGE` and `CREATE FORMAT` before migrate starts so we can have `cloudquery migrate` with fewer permissions required.
diff --git a/plugins/destination/snowflake/client/client.go b/plugins/destination/snowflake/client/client.go
@@ -5,6 +5,7 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
+	"sync"
 
 	"github.com/cloudquery/plugin-sdk/v4/plugin"
 	"github.com/cloudquery/plugin-sdk/v4/writers/batchwriter"
@@ -19,12 +20,15 @@ type Client struct {
 	logger zerolog.Logger
 	spec   Spec
 	writer *batchwriter.BatchWriter
+
+	setupWriteOnce *sync.Once
 }
 
-func New(ctx context.Context, logger zerolog.Logger, spec []byte, _ plugin.NewClientOptions) (plugin.Client, error) {
+func New(_ context.Context, logger zerolog.Logger, spec []byte, _ plugin.NewClientOptions) (plugin.Client, error) {
 	var err error
 	c := &Client{
-		logger: logger.With().Str("module", "sf-dest").Logger(),
+		logger:         logger.With().Str("module", "sf-dest").Logger(),
+		setupWriteOnce: &sync.Once{},
 	}
 	if err := json.Unmarshal(spec, &c.spec); err != nil {
 		return nil, fmt.Errorf("failed to unmarshal snowflake spec: %w", err)
@@ -43,12 +47,6 @@ func New(ctx context.Context, logger zerolog.Logger, spec []byte, _ plugin.NewCl
 		return nil, err
 	}
 	c.db = db
-	if _, err := c.db.ExecContext(ctx, createOrReplaceFileFormat); err != nil {
-		return nil, fmt.Errorf("failed to create file format %s: %w", createOrReplaceFileFormat, err)
-	}
-	if _, err := c.db.ExecContext(ctx, createOrReplaceStage); err != nil {
-		return nil, fmt.Errorf("failed to create stage %s: %w", createOrReplaceStage, err)
-	}
 	return c, nil
 }
 
diff --git a/plugins/destination/snowflake/client/write.go b/plugins/destination/snowflake/client/write.go
@@ -20,6 +20,10 @@ const (
 )
 
 func (c *Client) Write(ctx context.Context, msgs <-chan message.WriteMessage) error {
+	if err := c.setupWrite(ctx); err != nil {
+		return err
+	}
+
 	if err := c.writer.Write(ctx, msgs); err != nil {
 		return err
 	}
@@ -68,6 +72,20 @@ func (c *Client) WriteTableBatch(ctx context.Context, name string, msgs message.
 	return err
 }
 
+func (c *Client) setupWrite(ctx context.Context) error {
+	var setupErr error
+	c.setupWriteOnce.Do(func() {
+		if _, err := c.db.ExecContext(ctx, createOrReplaceFileFormat); err != nil {
+			setupErr = fmt.Errorf("failed to create file format %s: %w", createOrReplaceFileFormat, err)
+			return
+		}
+		if _, err := c.db.ExecContext(ctx, createOrReplaceStage); err != nil {
+			setupErr = fmt.Errorf("failed to create stage %s: %w", createOrReplaceStage, err)
+		}
+	})
+	return setupErr
+}
+
 // escapePath properly escapes the `\` character in window's file paths.
 func escapePath(p string) string {
 	return strings.ReplaceAll(p, "\\", "\\\\")
