chore: Add support for usage tracking to data sources

pkg/datasources/accounts.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"log"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/snowflakeroles"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
@@ -110,19 +113,18 @@ var accountsSchema = map[string]*schema.Schema{
 // Accounts Snowflake Accounts resource.
 func Accounts() *schema.Resource {
 	return &schema.Resource{
-		Read:   ReadAccounts,
-		Schema: accountsSchema,
+		ReadContext: TrackingReadWrapper(datasources.Accounts, ReadAccounts),
+		Schema:      accountsSchema,
 	}
 }
 
 // ReadAccounts lists accounts.
-func ReadAccounts(d *schema.ResourceData, meta interface{}) error {
+func ReadAccounts(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	client := meta.(*provider.Context).Client
-	ctx := context.Background()
 
 	ok, err := client.ContextFunctions.IsRoleInSession(ctx, snowflakeroles.Orgadmin)
 	if err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if !ok {
 		log.Printf("[DEBUG] ORGADMIN role is not in current session, cannot read accounts")
@@ -136,7 +138,7 @@ func ReadAccounts(d *schema.ResourceData, meta interface{}) error {
 	}
 	accounts, err := client.Accounts.Show(ctx, opts)
 	if err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	d.SetId("accounts")
 	accountsFlatten := []map[string]interface{}{}
@@ -161,7 +163,7 @@ func ReadAccounts(d *schema.ResourceData, meta interface{}) error {
 		accountsFlatten = append(accountsFlatten, m)
 	}
 	if err := d.Set("accounts", accountsFlatten); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	return nil
 }

pkg/datasources/alerts.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"log"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
 
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
@@ -76,15 +79,14 @@ var alertsSchema = map[string]*schema.Schema{
 // Alerts Snowflake Roles resource.
 func Alerts() *schema.Resource {
 	return &schema.Resource{
-		Read:   ReadAlerts,
-		Schema: alertsSchema,
+		ReadContext: TrackingReadWrapper(datasources.Alerts, ReadAlerts),
+		Schema:      alertsSchema,
 	}
 }
 
 // ReadAlerts Reads the database metadata information.
-func ReadAlerts(d *schema.ResourceData, meta interface{}) error {
+func ReadAlerts(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	client := meta.(*provider.Context).Client
-	ctx := context.Background()
 
 	d.SetId("alerts_read")
 
@@ -114,7 +116,7 @@ func ReadAlerts(d *schema.ResourceData, meta interface{}) error {
 	if err != nil {
 		log.Printf("[DEBUG] failed to list alerts in schema (%s)", d.Id())
 		d.SetId("")
-		return err
+		return diag.FromErr(err)
 	}
 
 	alerts := make([]map[string]any, 0, len(listAlerts))
@@ -127,7 +129,7 @@ func ReadAlerts(d *schema.ResourceData, meta interface{}) error {
 	}
 
 	if err := d.Set("alerts", alerts); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	return nil
 }

pkg/datasources/common.go

@@ -1,8 +1,13 @@
 package datasources
 
 import (
+	"context"
+
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/tracking"
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/resources"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
@@ -147,3 +152,10 @@ func handleExtendedIn(d *schema.ResourceData, setField **sdk.ExtendedIn) error {
 	}
 	return nil
 }
+
+func TrackingReadWrapper(datasourceName datasources.Datasource, readImplementation schema.ReadContextFunc) schema.ReadContextFunc {
+	return func(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
+		ctx = tracking.NewContext(ctx, tracking.NewVersionedDatasourceMetadata(datasourceName))
+		return readImplementation(ctx, d, meta)
+	}
+}

pkg/datasources/connections.go

@@ -3,6 +3,8 @@ package datasources
 import (
 	"context"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/resources"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/schemas"
@@ -34,7 +36,7 @@ var connectionsSchema = map[string]*schema.Schema{
 
 func Connections() *schema.Resource {
 	return &schema.Resource{
-		ReadContext: ReadConnections,
+		ReadContext: TrackingReadWrapper(datasources.Connections, ReadConnections),
 		Schema:      connectionsSchema,
 		Description: "Datasource used to get details of filtered connections. Filtering is aligned with the current possibilities for [SHOW CONNECTIONS](https://docs.snowflake.com/en/sql-reference/sql/show-connections) query. The results of SHOW is encapsulated in one output collection `connections`.",
 	}

pkg/datasources/cortex_search_services.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"log"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
 
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/sdk"
@@ -108,13 +111,13 @@ var cortexSearchServicesSchema = map[string]*schema.Schema{
 // CortexSearchServices Snowflake Cortex search services resource.
 func CortexSearchServices() *schema.Resource {
 	return &schema.Resource{
-		Read:   ReadCortexSearchServices,
-		Schema: cortexSearchServicesSchema,
+		ReadContext: TrackingReadWrapper(datasources.CortexSearchServices, ReadCortexSearchServices),
+		Schema:      cortexSearchServicesSchema,
 	}
 }
 
 // ReadCortexSearchServices Reads the cortex search services metadata information.
-func ReadCortexSearchServices(d *schema.ResourceData, meta interface{}) error {
+func ReadCortexSearchServices(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	client := meta.(*provider.Context).Client
 	request := sdk.NewShowCortexSearchServiceRequest()
 
@@ -167,7 +170,7 @@ func ReadCortexSearchServices(d *schema.ResourceData, meta interface{}) error {
 	if err != nil {
 		log.Printf("[DEBUG] snowflake_cortex_search_services.go: %v", err)
 		d.SetId("")
-		return err
+		return diag.FromErr(err)
 	}
 	d.SetId("cortex_search_services")
 	records := make([]map[string]any, 0, len(dts))
@@ -181,7 +184,7 @@ func ReadCortexSearchServices(d *schema.ResourceData, meta interface{}) error {
 		records = append(records, record)
 	}
 	if err := d.Set("cortex_search_services", records); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	return nil
 }

pkg/datasources/current_account.go

@@ -5,6 +5,9 @@ import (
 	"fmt"
 	"log"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
 
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@@ -33,15 +36,14 @@ var currentAccountSchema = map[string]*schema.Schema{
 // CurrentAccount the Snowflake current account resource.
 func CurrentAccount() *schema.Resource {
 	return &schema.Resource{
-		Read:   ReadCurrentAccount,
-		Schema: currentAccountSchema,
+		ReadContext: TrackingReadWrapper(datasources.CurrentAccount, ReadCurrentAccount),
+		Schema:      currentAccountSchema,
 	}
 }
 
 // ReadCurrentAccount read the current snowflake account information.
-func ReadCurrentAccount(d *schema.ResourceData, meta interface{}) error {
+func ReadCurrentAccount(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	client := meta.(*provider.Context).Client
-	ctx := context.Background()
 
 	current, err := client.ContextFunctions.CurrentSessionDetails(ctx)
 	if err != nil {
@@ -53,11 +55,11 @@ func ReadCurrentAccount(d *schema.ResourceData, meta interface{}) error {
 	d.SetId(fmt.Sprintf("%s.%s", current.Account, current.Region))
 	accountErr := d.Set("account", current.Account)
 	if accountErr != nil {
-		return accountErr
+		return diag.FromErr(accountErr)
 	}
 	regionErr := d.Set("region", current.Region)
 	if regionErr != nil {
-		return regionErr
+		return diag.FromErr(regionErr)
 	}
 	url, err := current.AccountURL()
 	if err != nil {
@@ -67,7 +69,7 @@ func ReadCurrentAccount(d *schema.ResourceData, meta interface{}) error {
 
 	urlErr := d.Set("url", url)
 	if urlErr != nil {
-		return urlErr
+		return diag.FromErr(urlErr)
 	}
 	return nil
 }

pkg/datasources/current_role.go

@@ -4,6 +4,9 @@ import (
 	"context"
 	"log"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/helpers"
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
@@ -19,14 +22,13 @@ var currentRoleSchema = map[string]*schema.Schema{
 
 func CurrentRole() *schema.Resource {
 	return &schema.Resource{
-		Read:   ReadCurrentRole,
-		Schema: currentRoleSchema,
+		ReadContext: TrackingReadWrapper(datasources.CurrentRole, ReadCurrentRole),
+		Schema:      currentRoleSchema,
 	}
 }
 
-func ReadCurrentRole(d *schema.ResourceData, meta interface{}) error {
+func ReadCurrentRole(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	client := meta.(*provider.Context).Client
-	ctx := context.Background()
 
 	role, err := client.ContextFunctions.CurrentRole(ctx)
 	if err != nil {
@@ -38,7 +40,7 @@ func ReadCurrentRole(d *schema.ResourceData, meta interface{}) error {
 	d.SetId(helpers.EncodeSnowflakeID(role))
 	err = d.Set("name", role.Name())
 	if err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	return nil
 }

pkg/datasources/database.go

@@ -3,6 +3,9 @@ package datasources
 import (
 	"context"
 
+	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/provider/datasources"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
+
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/helpers"
 
 	"github.com/Snowflake-Labs/terraform-provider-snowflake/pkg/internal/provider"
@@ -54,52 +57,51 @@ var databaseSchema = map[string]*schema.Schema{
 // Database the Snowflake Database resource.
 func Database() *schema.Resource {
 	return &schema.Resource{
-		Read:   ReadDatabase,
-		Schema: databaseSchema,
+		ReadContext: TrackingReadWrapper(datasources.Database, ReadDatabase),
+		Schema:      databaseSchema,
 	}
 }
 
 // ReadDatabase read the database meta-data information.
-func ReadDatabase(d *schema.ResourceData, meta interface{}) error {
+func ReadDatabase(ctx context.Context, d *schema.ResourceData, meta any) diag.Diagnostics {
 	client := meta.(*provider.Context).Client
-	ctx := context.Background()
 	name := d.Get("name").(string)
 	id := sdk.NewAccountObjectIdentifier(name)
 	database, err := client.Databases.ShowByID(ctx, id)
 	if err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	d.SetId(helpers.EncodeResourceIdentifier(database.ID()))
 	if err := d.Set("name", database.Name); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if err := d.Set("comment", database.Comment); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if err := d.Set("owner", database.Owner); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if err := d.Set("is_default", database.IsDefault); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if err := d.Set("is_current", database.IsCurrent); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	var origin string
 	if database.Origin != nil {
 		origin = database.Origin.FullyQualifiedName()
 	}
 	if err := d.Set("origin", origin); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if err := d.Set("retention_time", database.RetentionTime); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if err := d.Set("created_on", database.CreatedOn.String()); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	if err := d.Set("options", database.Options); err != nil {
-		return err
+		return diag.FromErr(err)
 	}
 	return nil
 }