chore: Basic object tracking (#3205)

## Changes
- All proposals for basic object tracking tested
- Added functions (and tested them) that allow us to use Golang's
context to chosen usage tracking

## Next pr
-
#3205 (comment)

pkg/acceptance/helpers/information_schema_client.go

@@ -0,0 +1,44 @@
+package helpers
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+	"github.com/stretchr/testify/require"
+)
+
+type InformationSchemaClient struct {
+	context *TestClientContext
+	ids     *IdsGenerator
+}
+
+func NewInformationSchemaClient(context *TestClientContext, idsGenerator *IdsGenerator) *InformationSchemaClient {
+	return &InformationSchemaClient{
+		context: context,
+		ids:     idsGenerator,
+	}
+}
+
+func (c *InformationSchemaClient) client() *sdk.Client {
+	return c.context.client
+}
+
+func (c *InformationSchemaClient) GetQueryTextByQueryId(t *testing.T, queryId string) string {
+	t.Helper()
+	result, err := c.client().QueryUnsafe(context.Background(), fmt.Sprintf("SELECT QUERY_TEXT FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY(RESULT_LIMIT => 20)) WHERE QUERY_ID = '%s'", queryId))
+	require.NoError(t, err)
+	require.Len(t, result, 1)
+	require.NotNil(t, result[0]["QUERY_TEXT"])
+	return (*result[0]["QUERY_TEXT"]).(string)
+}
+
+func (c *InformationSchemaClient) GetQueryTagByQueryId(t *testing.T, queryId string) string {
+	t.Helper()
+	result, err := c.client().QueryUnsafe(context.Background(), fmt.Sprintf("SELECT QUERY_TAG FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY(RESULT_LIMIT => 20)) WHERE QUERY_ID = '%s'", queryId))
+	require.NoError(t, err)
+	require.Len(t, result, 1)
+	require.NotNil(t, result[0]["QUERY_TAG"])
+	return (*result[0]["QUERY_TAG"]).(string)
+}

pkg/acceptance/helpers/test_client.go

@@ -36,6 +36,7 @@ type TestClient struct {
 	FileFormat                   *FileFormatClient
 	Function                     *FunctionClient
 	Grant                        *GrantClient
+	InformationSchema            *InformationSchemaClient
 	MaskingPolicy                *MaskingPolicyClient
 	MaterializedView             *MaterializedViewClient
 	NetworkPolicy                *NetworkPolicyClient
@@ -108,6 +109,7 @@ func NewTestClient(c *sdk.Client, database string, schema string, warehouse stri
 		FileFormat:                   NewFileFormatClient(context, idsGenerator),
 		Function:                     NewFunctionClient(context, idsGenerator),
 		Grant:                        NewGrantClient(context, idsGenerator),
+		InformationSchema:            NewInformationSchemaClient(context, idsGenerator),
 		MaskingPolicy:                NewMaskingPolicyClient(context, idsGenerator),
 		MaterializedView:             NewMaterializedViewClient(context, idsGenerator),
 		NetworkPolicy:                NewNetworkPolicyClient(context, idsGenerator),

pkg/acceptance/helpers/user_client.go

@@ -68,6 +68,14 @@ func (c *UserClient) Alter(t *testing.T, id sdk.AccountObjectIdentifier, opts *s
 	require.NoError(t, err)
 }
 
+func (c *UserClient) AlterCurrentUser(t *testing.T, opts *sdk.AlterUserOptions) {
+	t.Helper()
+	id, err := c.context.client.ContextFunctions.CurrentUser(context.Background())
+	require.NoError(t, err)
+	err = c.client().Alter(context.Background(), id, opts)
+	require.NoError(t, err)
+}
+
 func (c *UserClient) DropUserFunc(t *testing.T, id sdk.AccountObjectIdentifier) func() {
 	t.Helper()
 	ctx := context.Background()

pkg/internal/tracking/context.go

@@ -0,0 +1,73 @@
+package tracking
+
+import (
+	"context"
+	"errors"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
+)
+
+const (
+	ProviderVersion string = "v0.99.0" // TODO(SNOW-1814934): Currently hardcoded, make it computed
+	MetadataPrefix  string = "terraform_provider_usage_tracking"
+)
+
+type key struct{}
+
+var metadataContextKey key
+
+type Operation string
+
+const (
+	CreateOperation     Operation = "create"
+	ReadOperation       Operation = "read"
+	UpdateOperation     Operation = "update"
+	DeleteOperation     Operation = "delete"
+	ImportOperation     Operation = "import"
+	CustomDiffOperation Operation = "custom_diff"
+)
+
+type Metadata struct {
+	Version   string    `json:"version,omitempty"`
+	Resource  string    `json:"resource,omitempty"`
+	Operation Operation `json:"operation,omitempty"`
+}
+
+func (m Metadata) validate() error {
+	errs := make([]error, 0)
+	if m.Version == "" {
+		errs = append(errs, errors.New("version for metadata should not be empty"))
+	}
+	if m.Resource == "" {
+		errs = append(errs, errors.New("resource name for metadata should not be empty"))
+	}
+	if m.Operation == "" {
+		errs = append(errs, errors.New("operation for metadata should not be empty"))
+	}
+	return errors.Join(errs...)
+}
+
+func NewMetadata(version string, resource resources.Resource, operation Operation) Metadata {
+	return Metadata{
+		Version:   version,
+		Resource:  resource.String(),
+		Operation: operation,
+	}
+}
+
+func NewVersionedMetadata(resource resources.Resource, operation Operation) Metadata {
+	return Metadata{
+		Version:   ProviderVersion,
+		Resource:  resource.String(),
+		Operation: operation,
+	}
+}
+
+func NewContext(ctx context.Context, metadata Metadata) context.Context {
+	return context.WithValue(ctx, metadataContextKey, metadata)
+}
+
+func FromContext(ctx context.Context) (Metadata, bool) {
+	metadata, ok := ctx.Value(metadataContextKey).(Metadata)
+	return metadata, ok
+}

pkg/internal/tracking/context_test.go

@@ -0,0 +1,45 @@
+package tracking
+
+import (
+	"context"
+	"testing"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_Context(t *testing.T) {
+	metadata := NewMetadata("123", resources.Account, CreateOperation)
+	newMetadata := NewMetadata("321", resources.Database, UpdateOperation)
+	ctx := context.Background()
+
+	// no metadata in context
+	value := ctx.Value(metadataContextKey)
+	require.Nil(t, value)
+
+	retrievedMetadata, ok := FromContext(ctx)
+	require.False(t, ok)
+	require.Empty(t, retrievedMetadata)
+
+	// add metadata by hand
+	ctx = context.WithValue(ctx, metadataContextKey, metadata)
+
+	value = ctx.Value(metadataContextKey)
+	require.NotNil(t, value)
+	require.Equal(t, metadata, value)
+
+	retrievedMetadata, ok = FromContext(ctx)
+	require.True(t, ok)
+	require.Equal(t, metadata, retrievedMetadata)
+
+	// add metadata with NewContext function (overrides previous value)
+	ctx = NewContext(ctx, newMetadata)
+
+	value = ctx.Value(metadataContextKey)
+	require.NotNil(t, value)
+	require.Equal(t, newMetadata, value)
+
+	retrievedMetadata, ok = FromContext(ctx)
+	require.True(t, ok)
+	require.Equal(t, newMetadata, retrievedMetadata)
+}

pkg/internal/tracking/query.go

@@ -0,0 +1,31 @@
+package tracking
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+)
+
+func AppendMetadata(sql string, metadata Metadata) (string, error) {
+	bytes, err := json.Marshal(metadata)
+	if err != nil {
+		return "", fmt.Errorf("failed to marshal the metadata: %w", err)
+	} else {
+		return fmt.Sprintf("%s --%s %s", sql, MetadataPrefix, string(bytes)), nil
+	}
+}
+
+func ParseMetadata(sql string) (Metadata, error) {
+	parts := strings.Split(sql, fmt.Sprintf("--%s", MetadataPrefix))
+	if len(parts) != 2 {
+		return Metadata{}, fmt.Errorf("failed to parse metadata from sql, incorrect number of parts, expected: 2, got: %d", len(parts))
+	}
+	var metadata Metadata
+	if err := json.Unmarshal([]byte(strings.TrimSpace(parts[1])), &metadata); err != nil {
+		return Metadata{}, fmt.Errorf("failed to unmarshal metadata from sql: %s, err = %w", sql, err)
+	}
+	if err := metadata.validate(); err != nil {
+		return Metadata{}, err
+	}
+	return metadata, nil
+}

pkg/internal/tracking/query_test.go

@@ -0,0 +1,65 @@
+package tracking
+
+import (
+	"encoding/json"
+	"fmt"
+	"testing"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
+	"github.com/stretchr/testify/require"
+)
+
+func TestAppendMetadata(t *testing.T) {
+	metadata := NewMetadata("123", resources.Account, CreateOperation)
+	sql := "SELECT 1"
+
+	bytes, err := json.Marshal(metadata)
+	require.NoError(t, err)
+
+	expectedSql := fmt.Sprintf("%s --%s %s", sql, MetadataPrefix, string(bytes))
+
+	newSql, err := AppendMetadata(sql, metadata)
+	require.NoError(t, err)
+	require.Equal(t, expectedSql, newSql)
+}
+
+func TestParseMetadata(t *testing.T) {
+	metadata := NewMetadata("123", resources.Account, CreateOperation)
+	bytes, err := json.Marshal(metadata)
+	require.NoError(t, err)
+	sql := fmt.Sprintf("SELECT 1 --%s %s", MetadataPrefix, string(bytes))
+
+	parsedMetadata, err := ParseMetadata(sql)
+	require.NoError(t, err)
+	require.Equal(t, metadata, parsedMetadata)
+}
+
+func TestParseInvalidMetadataKeys(t *testing.T) {
+	sql := fmt.Sprintf(`SELECT 1 --%s {"key": "value"}`, MetadataPrefix)
+
+	parsedMetadata, err := ParseMetadata(sql)
+	require.ErrorContains(t, err, "version for metadata should not be empty")
+	require.ErrorContains(t, err, "resource name for metadata should not be empty")
+	require.ErrorContains(t, err, "operation for metadata should not be empty")
+	require.Equal(t, Metadata{}, parsedMetadata)
+}
+
+func TestParseInvalidMetadataJson(t *testing.T) {
+	sql := fmt.Sprintf(`SELECT 1 --%s "key": "value"`, MetadataPrefix)
+
+	parsedMetadata, err := ParseMetadata(sql)
+	require.ErrorContains(t, err, "failed to unmarshal metadata from sql")
+	require.Equal(t, Metadata{}, parsedMetadata)
+}
+
+func TestParseMetadataFromInvalidSqlCommentPrefix(t *testing.T) {
+	metadata := NewMetadata("123", resources.Account, CreateOperation)
+	sql := "SELECT 1"
+
+	bytes, err := json.Marshal(metadata)
+	require.NoError(t, err)
+
+	parsedMetadata, err := ParseMetadata(fmt.Sprintf("%s --invalid_prefix %s", sql, string(bytes)))
+	require.ErrorContains(t, err, "failed to parse metadata from sql")
+	require.Equal(t, Metadata{}, parsedMetadata)
+}

pkg/resources/common.go

@@ -5,6 +5,10 @@ import (
 	"regexp"
 	"strings"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/tracking"
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
 
 	"github.com/hashicorp/go-cty/cty"
@@ -101,3 +105,45 @@ func ImportName[T sdk.AccountObjectIdentifier | sdk.DatabaseObjectIdentifier | s
 
 	return []*schema.ResourceData{d}, nil
 }
+
+func TrackingImportWrapper(resourceName resources.Resource, importImplementation schema.StateContextFunc) schema.StateContextFunc {
+	return func(ctx context.Context, d *schema.ResourceData, meta any) ([]*schema.ResourceData, error) {
+		ctx = tracking.NewContext(ctx, tracking.NewVersionedMetadata(resourceName, tracking.ImportOperation))
+		return importImplementation(ctx, d, meta)
+	}
+}
+
+func TrackingCreateWrapper(resourceName resources.Resource, createImplementation schema.CreateContextFunc) schema.CreateContextFunc {
+	return func(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
+		ctx = tracking.NewContext(ctx, tracking.NewVersionedMetadata(resourceName, tracking.CreateOperation))
+		return createImplementation(ctx, d, meta)
+	}
+}
+
+func TrackingReadWrapper(resourceName resources.Resource, readImplementation schema.ReadContextFunc) schema.ReadContextFunc {
+	return func(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
+		ctx = tracking.NewContext(ctx, tracking.NewVersionedMetadata(resourceName, tracking.ReadOperation))
+		return readImplementation(ctx, d, meta)
+	}
+}
+
+func TrackingUpdateWrapper(resourceName resources.Resource, updateImplementation schema.UpdateContextFunc) schema.UpdateContextFunc {
+	return func(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
+		ctx = tracking.NewContext(ctx, tracking.NewVersionedMetadata(resourceName, tracking.UpdateOperation))
+		return updateImplementation(ctx, d, meta)
+	}
+}
+
+func TrackingDeleteWrapper(resourceName resources.Resource, deleteImplementation schema.DeleteContextFunc) schema.DeleteContextFunc {
+	return func(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
+		ctx = tracking.NewContext(ctx, tracking.NewVersionedMetadata(resourceName, tracking.DeleteOperation))
+		return deleteImplementation(ctx, d, meta)
+	}
+}
+
+func TrackingCustomDiffWrapper(resourceName resources.Resource, customdiffImplementation schema.CustomizeDiffFunc) schema.CustomizeDiffFunc {
+	return func(ctx context.Context, diff *schema.ResourceDiff, meta any) error {
+		ctx = tracking.NewContext(ctx, tracking.NewVersionedMetadata(resourceName, tracking.CustomDiffOperation))
+		return customdiffImplementation(ctx, diff, meta)
+	}
+}

pkg/resources/schema.go

@@ -8,6 +8,8 @@ import (
 	"slices"
 	"strings"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/resources"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/helpers"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/collections"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
@@ -89,23 +91,24 @@ var schemaSchema = map[string]*schema.Schema{
 // Schema returns a pointer to the resource representing a schema.
 func Schema() *schema.Resource {
 	return &schema.Resource{
-		CreateContext: CreateContextSchema,
-		ReadContext:   ReadContextSchema(true),
-		UpdateContext: UpdateContextSchema,
-		DeleteContext: DeleteContextSchema,
+		CreateContext: TrackingCreateWrapper(resources.Schema, CreateContextSchema),
+		ReadContext:   TrackingReadWrapper(resources.Schema, ReadContextSchema(true)),
+		UpdateContext: TrackingUpdateWrapper(resources.Schema, UpdateContextSchema),
+		DeleteContext: TrackingDeleteWrapper(resources.Schema, DeleteContextSchema),
 		Description:   "Resource used to manage schema objects. For more information, check [schema documentation](https://docs.snowflake.com/en/sql-reference/sql/create-schema).",
 
-		CustomizeDiff: customdiff.All(
+		CustomizeDiff: TrackingCustomDiffWrapper(resources.Schema, customdiff.All(
 			ComputedIfAnyAttributeChanged(schemaSchema, ShowOutputAttributeName, "name", "comment", "with_managed_access", "is_transient"),
 			ComputedIfAnyAttributeChanged(schemaSchema, DescribeOutputAttributeName, "name"),
 			ComputedIfAnyAttributeChanged(schemaSchema, FullyQualifiedNameAttributeName, "name"),
 			ComputedIfAnyAttributeChanged(schemaParametersSchema, ParametersAttributeName, collections.Map(sdk.AsStringList(sdk.AllSchemaParameters), strings.ToLower)...),
+			// TODO(SNOW-1804424 - next pr): handle custom context in parameters customdiff
 			schemaParametersCustomDiff,
-		),
+		)),
 
 		Schema: collections.MergeMaps(schemaSchema, schemaParametersSchema),
 		Importer: &schema.ResourceImporter{
-			StateContext: ImportSchema,
+			StateContext: TrackingImportWrapper(resources.Schema, ImportSchema),
 		},
 
 		SchemaVersion: 2,