refactor: make func re-usable

Author: vaishakdinesh

internal/services/dns_record/schema.go

@@ -4,8 +4,7 @@ package dns_record
 
 import (
 	"context"
-	"os"
-
+	
 	"github.com/hashicorp/terraform-plugin-framework-jsontypes/jsontypes"
 	"github.com/hashicorp/terraform-plugin-framework-timetypes/timetypes"
 	"github.com/hashicorp/terraform-plugin-framework-validators/float64validator"
@@ -23,34 +22,14 @@ import (
 
 	"github.com/cloudflare/terraform-provider-cloudflare/internal/customfield"
 	"github.com/cloudflare/terraform-provider-cloudflare/internal/customvalidator"
+	"github.com/cloudflare/terraform-provider-cloudflare/internal/utils"
 )
 
 var _ resource.ResourceWithConfigValidators = (*DNSRecordResource)(nil)
 
-// GetSchemaVersion returns the appropriate schema version based on TF_MIG_TEST environment variable.
-//
-// This function allows controlled rollout of StateUpgrader migrations:
-//   - During development/testing: Set TF_MIG_TEST=1 to enable migrations (returns postMigration version)
-//   - In production: StateUpgraders remain dormant (returns preMigration version)
-//   - For coordinated release: Remove this wrapper and set Version directly to enable all migrations at once
-//
-// Parameters:
-//   - preMigration: The version to use when migrations are disabled (typically 0)
-//   - postMigration: The version to use when migrations are enabled (typically 500)
-//
-// Example usage:
-//
-//	Version: GetSchemaVersion(0, 500)  // Returns 0 normally, 500 when TF_MIG_TEST=1
-func GetSchemaVersion(preMigration, postMigration int64) int64 {
-	if os.Getenv("TF_MIG_TEST") == "" {
-		return preMigration
-	}
-	return postMigration
-}
-
 func ResourceSchema(ctx context.Context) schema.Schema {
 	return schema.Schema{
-		Version: GetSchemaVersion(0, 500),
+		Version: utils.GetSchemaVersion(0, 500),
 		Attributes: map[string]schema.Attribute{
 			"id": schema.StringAttribute{
 				Description:   "Identifier.",

internal/services/load_balancer_pool/schema.go

@@ -4,9 +4,9 @@ package load_balancer_pool
 
 import (
 	"context"
-	"os"
 
 	"github.com/cloudflare/terraform-provider-cloudflare/internal/customfield"
+	"github.com/cloudflare/terraform-provider-cloudflare/internal/utils"
 	"github.com/hashicorp/terraform-plugin-framework-timetypes/timetypes"
 	"github.com/hashicorp/terraform-plugin-framework-validators/float64validator"
 	"github.com/hashicorp/terraform-plugin-framework-validators/listvalidator"
@@ -25,30 +25,9 @@ import (
 
 var _ resource.ResourceWithConfigValidators = (*LoadBalancerPoolResource)(nil)
 
-// GetSchemaVersion returns the appropriate schema version based on TF_MIG_TEST environment variable.
-//
-// This function allows controlled rollout of StateUpgrader migrations:
-//   - During development/testing: Set TF_MIG_TEST=1 to enable migrations (returns postMigration version)
-//   - In production: StateUpgraders remain dormant (returns preMigration version)
-//   - For coordinated release: Remove this wrapper and set Version directly to enable all migrations at once
-//
-// Parameters:
-//   - preMigration: The version to use when migrations are disabled (typically 0)
-//   - postMigration: The version to use when migrations are enabled (typically 500)
-//
-// Example usage:
-//
-//	Version: GetSchemaVersion(0, 500)  // Returns 0 normally, 500 when TF_MIG_TEST=1
-func GetSchemaVersion(preMigration, postMigration int64) int64 {
-	if os.Getenv("TF_MIG_TEST") == "" {
-		return preMigration
-	}
-	return postMigration
-}
-
 func ResourceSchema(ctx context.Context) schema.Schema {
 	return schema.Schema{
-		Version: GetSchemaVersion(0, 500),
+		Version: utils.GetSchemaVersion(0, 500),
 		Attributes: map[string]schema.Attribute{
 			"id": schema.StringAttribute{
 				Computed:      true,

internal/services/page_rule/schema.go

@@ -4,7 +4,9 @@ package page_rule
 
 import (
 	"context"
+
 	"github.com/cloudflare/terraform-provider-cloudflare/internal/customfield"
+	"github.com/cloudflare/terraform-provider-cloudflare/internal/utils"
 	"github.com/hashicorp/terraform-plugin-framework-timetypes/timetypes"
 	"github.com/hashicorp/terraform-plugin-framework-validators/int64validator"
 	"github.com/hashicorp/terraform-plugin-framework-validators/stringvalidator"
@@ -25,7 +27,7 @@ var _ resource.ResourceWithConfigValidators = (*PageRuleResource)(nil)
 
 func ResourceSchema(ctx context.Context) schema.Schema {
 	return schema.Schema{
-		Version: 500,
+		Version: utils.GetSchemaVersion(0, 500),
 		Attributes: map[string]schema.Attribute{
 			"id": schema.StringAttribute{
 				Description:   "Identifier.",

internal/utils/schema_version.go

@@ -0,0 +1,24 @@
+package utils
+
+import "os"
+
+// GetSchemaVersion returns the appropriate schema version based on TF_MIG_TEST environment variable.
+//
+// This function allows controlled rollout of StateUpgrader migrations:
+//   - During development/testing: Set TF_MIG_TEST=1 to enable migrations (returns postMigration version)
+//   - In production: StateUpgraders remain dormant (returns preMigration version)
+//   - For coordinated release: Remove this wrapper and set Version directly to enable all migrations at once
+//
+// Parameters:
+//   - preMigration: The version to use when migrations are disabled (typically 0)
+//   - postMigration: The version to use when migrations are enabled (typically 500)
+//
+// Example usage:
+//
+//	Version: GetSchemaVersion(0, 500)  // Returns 0 normally, 500 when TF_MIG_TEST=1
+func GetSchemaVersion(preMigration, postMigration int64) int64 {
+	if os.Getenv("TF_MIG_TEST") == "" {
+		return preMigration
+	}
+	return postMigration
+}