chore: Basic object tracking

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/context_functions.go

@@ -20,6 +20,8 @@ type ContextFunctions interface {
 	CurrentSession(ctx context.Context) (string, error)
 	CurrentUser(ctx context.Context) (AccountObjectIdentifier, error)
 	CurrentSessionDetails(ctx context.Context) (*CurrentSessionDetails, error)
+
+	// TODO(SNOW-1805152): Remove this and utilize gosnowflake.WithQueryIDChan instead whenever query id is needed
 	LastQueryId(ctx context.Context) (string, error)
 
 	// Session Object functions.

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

@@ -0,0 +1,94 @@
+package testint
+
+import (
+	"context"
+	"fmt"
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+	"github.com/snowflakedb/gosnowflake"
+	"github.com/stretchr/testify/require"
+	"strings"
+	"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)
+	ctx := context.Background()
+
+	sessionId, err := client.ContextFunctions.CurrentSession(ctx)
+	require.NoError(t, err)
+
+	queryTag := "some query tag"
+	require.NoError(t, client.Parameters.SetSessionParameterOnAccount(ctx, sdk.SessionParameterQueryTag, queryTag))
+	t.Cleanup(func() {
+		_, err = client.QueryUnsafe(ctx, "ALTER SESSION UNSET QUERY_TAG")
+		require.NoError(t, err)
+	})
+
+	queryId := executeQueryAndReturnQueryId(t, context.Background(), client)
+
+	result, err := client.QueryUnsafe(ctx, fmt.Sprintf("SELECT QUERY_ID, QUERY_TAG FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY_BY_SESSION(SESSION_ID => %s, RESULT_LIMIT => 2)) WHERE QUERY_ID = '%s'", sessionId, queryId))
+	require.NoError(t, err)
+	require.Len(t, result, 1)
+	require.Equal(t, queryId, *result[0]["QUERY_ID"])
+	require.Equal(t, queryTag, *result[0]["QUERY_TAG"])
+
+	newQueryTag := "some other query tag"
+	ctxWithQueryTag := gosnowflake.WithQueryTag(context.Background(), newQueryTag)
+	newQueryId := executeQueryAndReturnQueryId(t, ctxWithQueryTag, client)
+
+	result, err = client.QueryUnsafe(ctx, fmt.Sprintf("SELECT QUERY_ID, QUERY_TAG FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY_BY_SESSION(SESSION_ID => %s, RESULT_LIMIT => 2)) WHERE QUERY_ID = '%s'", sessionId, newQueryId))
+	require.NoError(t, err)
+	require.Len(t, result, 1)
+	require.Equal(t, newQueryId, *result[0]["QUERY_ID"])
+	require.Equal(t, newQueryTag, *result[0]["QUERY_TAG"])
+}
+
+func executeQueryAndReturnQueryId(t *testing.T, ctx context.Context, client *sdk.Client) string {
+	queryIdChan := make(chan string, 1)
+	ctx = gosnowflake.WithQueryIDChan(ctx, queryIdChan)
+
+	_, err := client.QueryUnsafe(ctx, "SELECT 1")
+	require.NoError(t, err)
+
+	return <-queryIdChan
+}
+
+// https://select.dev/posts/snowflake-query-tags#using-query-comments-instead-of-query-tags
+func TestInt_QueryComment(t *testing.T) {
+	client := testClient(t)
+	ctx := context.Background()
+
+	sessionId, err := client.ContextFunctions.CurrentSession(ctx)
+	require.NoError(t, err)
+
+	queryIdChan := make(chan string, 1)
+	metadata := `{"comment": "some comment"}`
+	_, err = client.QueryUnsafe(gosnowflake.WithQueryIDChan(ctx, queryIdChan), fmt.Sprintf(`SELECT 1; --%s`, metadata))
+	require.NoError(t, err)
+	queryId := <-queryIdChan
+
+	result, err := client.QueryUnsafe(ctx, fmt.Sprintf("SELECT QUERY_ID, QUERY_TEXT FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY_BY_SESSION(SESSION_ID => %s, RESULT_LIMIT => 2)) WHERE QUERY_ID = '%s'", sessionId, queryId))
+	require.NoError(t, err)
+	require.Len(t, result, 1)
+	require.Equal(t, queryId, *result[0]["QUERY_ID"])
+	require.Equal(t, metadata, strings.Split((*result[0]["QUERY_TEXT"]).(string), "--")[1])
+}
+
+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")
+
+	version := "v0.99.0"
+	config := sdk.DefaultConfig()
+	config.Application = fmt.Sprintf("terraform-provider-snowflake:%s", version)
+	client, err := sdk.NewClient(config)
+	require.NoError(t, err)
+
+	_, err = client.QueryUnsafe(context.Background(), "SELECT 1")
+	require.NoError(t, err)
+}
+
+// TODO(SNOW-1805150): Document potential usage of connection string