Test different tracking proposals and implement chosen one

pkg/sdk/client.go

@@ -3,6 +3,7 @@ package sdk
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
 	"fmt"
 	"log"
 	"os"
@@ -132,7 +133,7 @@ func NewClient(cfg *gosnowflake.Config) (*Client, error) {
 			logger := instrumentedsql.LoggerFunc(func(ctx context.Context, s string, kv ...interface{}) {
 				switch s {
 				case "sql-conn-query", "sql-conn-exec":
-					log.Printf("[DEBUG] %s: %v (%s)\n", s, kv, ctx.Value(snowflakeAccountLocatorContextKey))
+					log.Printf("[DEBUG] %s: %v (%s)\n", s, kv, ctx.Value(SnowflakeAccountLocatorContextKey))
 				default:
 					return
 				}
@@ -264,20 +265,27 @@ func (c *Client) Close() error {
 	return nil
 }
 
-type snowflakeAccountLocatorContext string
+type ContextKey string
 
 const (
-	snowflakeAccountLocatorContextKey snowflakeAccountLocatorContext = "snowflake_account_locator"
+	SnowflakeAccountLocatorContextKey ContextKey = "snowflake_account_locator"
+	MetadataContextKey                ContextKey = "metadata"
+	DashboardTrackingPrefix                      = "dashboard_tracking_"
 )
 
+func ContextWithMetadata(ctx context.Context, metadata map[string]string) context.Context {
+	return context.WithValue(ctx, MetadataContextKey, metadata)
+}
+
 // Exec executes a query that does not return rows.
 func (c *Client) exec(ctx context.Context, sql string) (sql.Result, error) {
 	if c.dryRun {
 		c.traceLogs = append(c.traceLogs, sql)
 		log.Printf("[DEBUG] sql-conn-exec-dry: %v\n", sql)
 		return nil, nil
 	}
-	ctx = context.WithValue(ctx, snowflakeAccountLocatorContextKey, c.accountLocator)
+	ctx = context.WithValue(ctx, SnowflakeAccountLocatorContextKey, c.accountLocator)
+	sql = appendQueryMetadata(ctx, sql)
 	result, err := c.db.ExecContext(ctx, sql)
 	return result, decodeDriverError(err)
 }
@@ -289,7 +297,8 @@ func (c *Client) query(ctx context.Context, dest interface{}, sql string) error
 		log.Printf("[DEBUG] sql-conn-query-dry: %v\n", sql)
 		return nil
 	}
-	ctx = context.WithValue(ctx, snowflakeAccountLocatorContextKey, c.accountLocator)
+	ctx = context.WithValue(ctx, SnowflakeAccountLocatorContextKey, c.accountLocator)
+	sql = appendQueryMetadata(ctx, sql)
 	return decodeDriverError(c.db.SelectContext(ctx, dest, sql))
 }
 
@@ -300,6 +309,20 @@ func (c *Client) queryOne(ctx context.Context, dest interface{}, sql string) err
 		log.Printf("[DEBUG] sql-conn-query-one-dry: %v\n", sql)
 		return nil
 	}
-	ctx = context.WithValue(ctx, snowflakeAccountLocatorContextKey, c.accountLocator)
+	ctx = context.WithValue(ctx, SnowflakeAccountLocatorContextKey, c.accountLocator)
+	sql = appendQueryMetadata(ctx, sql)
 	return decodeDriverError(c.db.GetContext(ctx, dest, sql))
 }
+
+func appendQueryMetadata(ctx context.Context, query string) string {
+	if ctx.Value(MetadataContextKey) != nil {
+		metadataMap := ctx.Value(MetadataContextKey)
+		bytes, err := json.Marshal(metadataMap)
+		if err != nil {
+			log.Printf("[ERROR] failed to marshal the metadata: %v\n", err)
+		} else {
+			return fmt.Sprintf("%s --%s%s", query, DashboardTrackingPrefix, string(bytes))
+		}
+	}
+	return query
+}

pkg/sdk/client_integration_test.go

@@ -3,7 +3,10 @@ package sdk
 import (
 	"context"
 	"database/sql"
+	"encoding/json"
+	"fmt"
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/acceptance/testprofiles"
@@ -132,3 +135,72 @@ func TestClient_NewClientDriverLoggingLevel(t *testing.T) {
 		assert.Equal(t, "trace", gosnowflake.GetLogger().GetLogLevel())
 	})
 }
+
+func TestClient_AdditionalMetadata(t *testing.T) {
+	client := defaultTestClient(t)
+
+	// needed for using information_schema
+	databaseId := randomAccountObjectIdentifier()
+	require.NoError(t, client.Databases.Create(context.Background(), databaseId, &CreateDatabaseOptions{}))
+	t.Cleanup(func() {
+		require.NoError(t, client.Databases.Drop(context.Background(), databaseId, &DropDatabaseOptions{}))
+	})
+
+	metadata := map[string]string{
+		"version": "v1.0.0",
+		"method":  "create",
+	}
+
+	assertQueryMetadata := func(t *testing.T, queryId string) {
+		t.Helper()
+		result, err := client.QueryUnsafe(context.Background(), fmt.Sprintf("SELECT QUERY_ID, QUERY_TEXT FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY(RESULT_LIMIT => 2)) WHERE QUERY_ID = '%s'", queryId))
+		require.NoError(t, err)
+		require.Len(t, result, 1)
+		require.Equal(t, queryId, *result[0]["QUERY_ID"])
+		var parsedMetadata map[string]string
+		queryText := (*result[0]["QUERY_TEXT"]).(string)
+		queryMetadata := strings.Split(queryText, fmt.Sprintf("--%s", DashboardTrackingPrefix))[1]
+		err = json.Unmarshal([]byte(queryMetadata), &parsedMetadata)
+		require.NoError(t, err)
+		require.Equal(t, metadata, parsedMetadata)
+	}
+
+	t.Run("query one", func(t *testing.T) {
+		queryIdChan := make(chan string, 1)
+		ctx := context.Background()
+		ctx = ContextWithMetadata(ctx, metadata)
+		ctx = gosnowflake.WithQueryIDChan(ctx, queryIdChan)
+		row := struct {
+			One int `db:"ONE"`
+		}{}
+		err := client.queryOne(ctx, &row, "SELECT 1 AS ONE")
+		require.NoError(t, err)
+
+		assertQueryMetadata(t, <-queryIdChan)
+	})
+
+	t.Run("query", func(t *testing.T) {
+		queryIdChan := make(chan string, 1)
+		ctx := context.Background()
+		ctx = ContextWithMetadata(ctx, metadata)
+		ctx = gosnowflake.WithQueryIDChan(ctx, queryIdChan)
+		var rows []struct {
+			One int `db:"ONE"`
+		}
+		err := client.query(ctx, &rows, "SELECT 1 AS ONE")
+		require.NoError(t, err)
+
+		assertQueryMetadata(t, <-queryIdChan)
+	})
+
+	t.Run("exec", func(t *testing.T) {
+		queryIdChan := make(chan string, 1)
+		ctx := context.Background()
+		ctx = ContextWithMetadata(ctx, metadata)
+		ctx = gosnowflake.WithQueryIDChan(ctx, queryIdChan)
+		_, err := client.exec(ctx, "SELECT 1")
+		require.NoError(t, err)
+
+		assertQueryMetadata(t, <-queryIdChan)
+	})
+}

pkg/sdk/integration_test_imports.go

@@ -15,22 +15,22 @@ import (
 // ExecForTests is an exact copy of exec (that is unexported), that some integration tests/helpers were using
 // TODO: remove after we have all usages covered by SDK (for now it means implementing stages, tables, and tags)
 func (c *Client) ExecForTests(ctx context.Context, sql string) (sql.Result, error) {
-	ctx = context.WithValue(ctx, snowflakeAccountLocatorContextKey, c.accountLocator)
+	ctx = context.WithValue(ctx, SnowflakeAccountLocatorContextKey, c.accountLocator)
 	result, err := c.db.ExecContext(ctx, sql)
 	return result, decodeDriverError(err)
 }
 
 // QueryOneForTests is an exact copy of queryOne (that is unexported), that some integration tests/helpers were using
 // TODO: remove after introducing all resources using this
 func (c *Client) QueryOneForTests(ctx context.Context, dest interface{}, sql string) error {
-	ctx = context.WithValue(ctx, snowflakeAccountLocatorContextKey, c.accountLocator)
+	ctx = context.WithValue(ctx, SnowflakeAccountLocatorContextKey, c.accountLocator)
 	return decodeDriverError(c.db.GetContext(ctx, dest, sql))
 }
 
 // QueryForTests is an exact copy of query (that is unexported), that some integration tests/helpers were using
 // TODO: remove after introducing all resources using this
 func (c *Client) QueryForTests(ctx context.Context, dest interface{}, sql string) error {
-	ctx = context.WithValue(ctx, snowflakeAccountLocatorContextKey, c.accountLocator)
+	ctx = context.WithValue(ctx, SnowflakeAccountLocatorContextKey, c.accountLocator)
 	return decodeDriverError(c.db.SelectContext(ctx, dest, sql))
 }
 

pkg/sdk/testint/basic_object_tracking_integration_test.go

@@ -10,6 +10,8 @@ import (
 	"testing"
 )
 
+// Research for basic object tracking done as part of SNOW-1737787
+
 // https://docs.snowflake.com/en/sql-reference/parameters#query-tag
 func TestInt_ContextQueryTags(t *testing.T) {
 	client := testClient(t)
@@ -64,7 +66,7 @@ func TestInt_QueryComment(t *testing.T) {
 
 	queryIdChan := make(chan string, 1)
 	metadata := `{"comment": "some comment"}`
-	_, err = client.QueryUnsafe(gosnowflake.WithQueryIDChan(ctx, queryIdChan), fmt.Sprintf(`SELECT 1; --%s`, metadata)) // TODO: Check in Snowhouse (check with Sebastian if this approach will be possible to query in snowhouse)
+	_, err = client.QueryUnsafe(gosnowflake.WithQueryIDChan(ctx, queryIdChan), fmt.Sprintf(`SELECT 1; --%s`, metadata))
 	require.NoError(t, err)
 	queryId := <-queryIdChan
 
@@ -75,12 +77,6 @@ func TestInt_QueryComment(t *testing.T) {
 	require.Equal(t, metadata, strings.Split((*result[0]["QUERY_TEXT"]).(string), "--")[1])
 }
 
-func TestInt_QueryComment_ClientSupport(t *testing.T) {
-
-}
-
-// - check comment characters (if they're valid characters etc.) - could be solved by receiving object and turning into e.g. JSON will be done at a lower level
-
 func TestInt_AppName(t *testing.T) {
 	// https://community.snowflake.com/s/article/How-to-see-application-name-added-in-the-connection-string-in-Snowsight
 	t.Skip("there no way to check client application name by querying Snowflake's")
@@ -95,4 +91,4 @@ func TestInt_AppName(t *testing.T) {
 	require.NoError(t, err)
 }
 
-// TODO: trying to use connection parameters
+// TODO(SNOW-1805150): Document potential usage of connection string