refactor: minimum viable changes necessary for PluginSDK2 to work

azurerm/internal/acceptance/testcase.go

@@ -4,9 +4,10 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
+
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
-	"github.com/terraform-providers/terraform-provider-azuread/azuread"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/acceptance/helpers"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/acceptance/testclient"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/acceptance/types"
@@ -82,16 +83,12 @@ func RunTestsInSequence(t *testing.T, tests map[string]map[string]func(t *testin
 }
 
 func (td TestData) runAcceptanceTest(t *testing.T, testCase resource.TestCase) {
-	testCase.ProviderFactories = map[string]terraform.ResourceProviderFactory{
-		"azuread": func() (terraform.ResourceProvider, error) {
-			aad := azuread.Provider()
-			return aad, nil
-		},
-		"azurerm": func() (terraform.ResourceProvider, error) {
+	testCase.ProviderFactories = map[string]func() (*schema.Provider, error){
+		"azurerm": func() (*schema.Provider, error) {
 			azurerm := provider.TestAzureProvider()
 			return azurerm, nil
 		},
-		"azurerm-alt": func() (terraform.ResourceProvider, error) {
+		"azurerm-alt": func() (*schema.Provider, error) {
 			azurerm := provider.TestAzureProvider()
 			return azurerm, nil
 		},
@@ -101,12 +98,12 @@ func (td TestData) runAcceptanceTest(t *testing.T, testCase resource.TestCase) {
 }
 
 func (td TestData) runAcceptanceSequentialTest(t *testing.T, testCase resource.TestCase) {
-	testCase.ProviderFactories = map[string]terraform.ResourceProviderFactory{
-		"azuread": func() (terraform.ResourceProvider, error) {
-			aad := azuread.Provider()
-			return aad, nil
+	testCase.ProviderFactories = map[string]func() (*schema.Provider, error){
+		"azurerm": func() (*schema.Provider, error) {
+			azurerm := provider.TestAzureProvider()
+			return azurerm, nil
 		},
-		"azurerm": func() (terraform.ResourceProvider, error) {
+		"azurerm-alt": func() (*schema.Provider, error) {
 			azurerm := provider.TestAzureProvider()
 			return azurerm, nil
 		},

azurerm/internal/provider/provider.go

@@ -1,6 +1,7 @@
 package provider
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"os"
@@ -9,23 +10,22 @@ import (
 	"github.com/hashicorp/go-azure-helpers/authentication"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/validation"
-	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/clients"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/features"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/resourceproviders"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/sdk"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
 )
 
-func AzureProvider() terraform.ResourceProvider {
+func AzureProvider() *schema.Provider {
 	return azureProvider(false)
 }
 
-func TestAzureProvider() terraform.ResourceProvider {
+func TestAzureProvider() *schema.Provider {
 	return azureProvider(true)
 }
 
-func azureProvider(supportLegacyTestSuite bool) terraform.ResourceProvider {
+func azureProvider(supportLegacyTestSuite bool) *schema.Provider {
 	// avoids this showing up in test output
 	debugLog := func(f string, v ...interface{}) {
 		if os.Getenv("TF_LOG") == "" {
@@ -319,12 +319,12 @@ func providerConfigure(p *schema.Provider) schema.ConfigureFunc {
 			// platform level tracing
 			CustomCorrelationRequestID: os.Getenv("ARM_CORRELATION_REQUEST_ID"),
 		}
-		client, err := clients.Build(p.StopContext(), clientBuilder)
+		client, err := clients.Build(context.TODO(), clientBuilder)
 		if err != nil {
 			return nil, err
 		}
 
-		client.StopContext = p.StopContext()
+		client.StopContext = context.TODO()
 
 		if !skipProviderRegistration {
 			// List all the available providers and their registration state to avoid unnecessary

azurerm/internal/provider/provider_test.go

@@ -4,18 +4,16 @@ import (
 	"fmt"
 	"testing"
 	"time"
-
-	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 )
 
 func TestProvider(t *testing.T) {
-	if err := TestAzureProvider().(*schema.Provider).InternalValidate(); err != nil {
+	if err := TestAzureProvider().InternalValidate(); err != nil {
 		t.Fatalf("err: %s", err)
 	}
 }
 
 func TestDataSourcesSupportCustomTimeouts(t *testing.T) {
-	provider := TestAzureProvider().(*schema.Provider)
+	provider := TestAzureProvider()
 	for dataSourceName, dataSource := range provider.DataSourcesMap {
 		t.Run(fmt.Sprintf("DataSource/%s", dataSourceName), func(t *testing.T) {
 			t.Logf("[DEBUG] Testing Data Source %q..", dataSourceName)
@@ -51,7 +49,7 @@ func TestDataSourcesSupportCustomTimeouts(t *testing.T) {
 }
 
 func TestResourcesSupportCustomTimeouts(t *testing.T) {
-	provider := TestAzureProvider().(*schema.Provider)
+	provider := TestAzureProvider()
 	for resourceName, resource := range provider.ResourcesMap {
 		t.Run(fmt.Sprintf("Resource/%s", resourceName), func(t *testing.T) {
 			t.Logf("[DEBUG] Testing Resource %q..", resourceName)

azurerm/internal/services/loadbalancer/loadbalancer.go

@@ -4,7 +4,7 @@ import (
 	"context"
 
 	"github.com/Azure/azure-sdk-for-go/services/network/mgmt/2020-11-01/network"
-	"github.com/hashicorp/terraform-plugin-sdk/helper/schema"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/services/loadbalancer/parse"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/tf/pluginsdk"
 )

azurerm/internal/services/recoveryservices/backup_policy_file_share_resource.go

@@ -238,7 +238,7 @@ func resourceBackupProtectionPolicyFileShare() *pluginsdk.Resource {
 
 		// if daily, we need daily retention
 		// if weekly daily cannot be set, and we need weekly
-		CustomizeDiff: func(diff *pluginsdk.ResourceDiff, v interface{}) error {
+		CustomizeDiff: func(ctx context.Context, diff *pluginsdk.ResourceDiff, v interface{}) error {
 			_, hasDaily := diff.GetOk("retention_daily")
 			_, hasWeekly := diff.GetOk("retention_weekly")
 

azurerm/internal/tf/pluginsdk/customize_diff.go

@@ -20,10 +20,9 @@ type (
 //
 // If multiple functions returns errors, the result is a multierror.
 func CustomDiffWithAll(funcs ...CustomizeDiffFunc) schema.CustomizeDiffFunc {
-	return func(d *schema.ResourceDiff, meta interface{}) error {
+	return func(ctx context.Context, d *schema.ResourceDiff, meta interface{}) error {
 		var err error
 		for _, f := range funcs {
-			ctx := context.TODO()
 			thisErr := f(ctx, d, meta)
 			if thisErr != nil {
 				err = multierror.Append(err, thisErr)
@@ -39,9 +38,9 @@ func CustomDiffWithAll(funcs ...CustomizeDiffFunc) schema.CustomizeDiffFunc {
 //
 // If all functions succeed, the combined function also succeeds.
 func CustomDiffInSequence(funcs ...CustomizeDiffFunc) schema.CustomizeDiffFunc {
-	return func(d *schema.ResourceDiff, meta interface{}) error {
+	return func(ctx context.Context, d *schema.ResourceDiff, meta interface{}) error {
 		for _, f := range funcs {
-			err := f(context.TODO(), d, meta)
+			err := f(ctx, d, meta)
 			if err != nil {
 				return err
 			}

azurerm/internal/tf/pluginsdk/shim.go

@@ -11,16 +11,16 @@ import (
 // which allows us to switch out the version of the Plugin SDK being used
 // without breaking open PR's
 func CustomizeDiffShim(diffFunc CustomizeDiffFunc) schema.CustomizeDiffFunc {
-	return func(diff *schema.ResourceDiff, i interface{}) error {
-		return diffFunc(context.TODO(), diff, i)
+	return func(ctx context.Context, diff *schema.ResourceDiff, i interface{}) error {
+		return diffFunc(ctx, diff, i)
 	}
 }
 
 // ValueChangeConditionShim is a shim around the Terraform Plugin SDK
 // which allows us to switch out the version of the Plugin SDK being used
 // without breaking open PR's
 func ValueChangeConditionShim(shimFunc ValueChangeConditionFunc) customdiff.ValueChangeConditionFunc {
-	return func(old, new, meta interface{}) bool {
-		return shimFunc(context.TODO(), old, new, meta)
+	return func(ctx context.Context, old, new, meta interface{}) bool {
+		return shimFunc(ctx, old, new, meta)
 	}
 }

azurerm/internal/tf/pluginsdk/state_upgrades.go

@@ -46,12 +46,10 @@ func StateUpgrades(upgrades map[int]StateUpgrade) []StateUpgrader {
 		resource := Resource{
 			Schema: upgrade.Schema(),
 		}
-		// TODO: with Plugin SDK 1.x we'll need to add a wrapper here to inject ctx
-
 		out = append(out, StateUpgrader{
 			Type: resource.CoreConfigSchema().ImpliedType(),
-			Upgrade: func(rawState map[string]interface{}, meta interface{}) (map[string]interface{}, error) {
-				return upgrade.UpgradeFunc()(context.TODO(), rawState, meta)
+			Upgrade: func(ctx context.Context, rawState map[string]interface{}, meta interface{}) (map[string]interface{}, error) {
+				return upgrade.UpgradeFunc()(ctx, rawState, meta)
 			},
 			Version: version,
 		})

azurerm/provider.go

@@ -1,10 +1,10 @@
 package azurerm
 
 import (
-	"github.com/hashicorp/terraform-plugin-sdk/v2/terraform"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/provider"
 )
 
-func Provider() terraform.ResourceProvider {
+func Provider() *schema.Provider {
 	return provider.AzureProvider()
 }

main.go

@@ -4,9 +4,8 @@ import (
 	"context"
 	"flag"
 	"log"
-	"os"
 
-	"github.com/hashicorp/terraform-plugin-sdk/plugin"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/plugin"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm"
 )
 
@@ -20,9 +19,6 @@ func main() {
 	flag.Parse()
 
 	if debugMode {
-		// This is needed so Terraform doesn't default to expecting protocol 4.
-		// TODO: remove below line once the provider migrates to plugin SDK v2.
-		os.Setenv("PLUGIN_PROTOCOL_VERSIONS", "5")
 		err := plugin.Debug(context.Background(), "registry.terraform.io/hashicorp/azurerm",
 			&plugin.ServeOpts{
 				ProviderFunc: azurerm.Provider,