From d11fb1ffbcb87bd460af783bcdceaab98b422c04 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergen=20Yal=C3=A7=C4=B1n?= <yalcinsergen97@gmail.com>
Date: Thu, 23 Jan 2025 15:51:06 +0300
Subject: [PATCH] Parametrize the registry name of the provider
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Sergen Yalçın <yalcinsergen97@gmail.com>
(cherry picked from commit 5dbd74a606bca4f725330094fd5a792d32df80dc)
---
 pkg/terraform/files.go | 15 +++++++++++----
 pkg/terraform/store.go |  3 +++
 2 files changed, 14 insertions(+), 4 deletions(-)

diff --git a/pkg/terraform/files.go b/pkg/terraform/files.go
index 26a8fe1c..74129ab0 100644
--- a/pkg/terraform/files.go
+++ b/pkg/terraform/files.go
@@ -34,6 +34,8 @@ const (
 	errUnmarshalTFState  = "cannot unmarshal tfstate file"
 	errFmtNonString      = "cannot work with a non-string id: %s"
 	errReadMainTF        = "cannot read main.tf.json file"
+
+	defaultRegistry = `provider["registry.terraform.io/%s"]`
 )
 
 // FileProducerOption allows you to configure FileProducer
@@ -190,7 +192,7 @@ func (fp *FileProducer) WriteMainTF() (ProviderHandle, error) {
 
 // EnsureTFState writes the Terraform state that should exist in the filesystem
 // to start any Terraform operation.
-func (fp *FileProducer) EnsureTFState(_ context.Context, tfID string) error {
+func (fp *FileProducer) EnsureTFState(_ context.Context, tfID string) error { //nolint:gocyclo // easier to follow as a unit
 	// TODO(muvaf): Reduce the cyclomatic complexity by separating the attributes
 	// generation into its own function/interface.
 	empty, err := fp.isStateEmpty()
@@ -229,14 +231,19 @@ func (fp *FileProducer) EnsureTFState(_ context.Context, tfID string) error {
 	s := json.NewStateV4()
 	s.TerraformVersion = fp.Setup.Version
 	s.Lineage = string(fp.Resource.GetUID())
+
+	registry := fp.Setup.Requirement.Registry
+	if registry == "" {
+		registry = defaultRegistry
+	}
+
 	s.Resources = []json.ResourceStateV4{
 		{
 			Mode: "managed",
 			Type: fp.Resource.GetTerraformResourceType(),
 			Name: fp.Resource.GetName(),
-			// TODO(muvaf): we should get the full URL from Dockerfile since
-			// providers don't have to be hosted in registry.terraform.io
-			ProviderConfig: fmt.Sprintf(`provider["registry.terraform.io/%s"]`, fp.Setup.Requirement.Source),
+			// Support for private/non-default registries
+			ProviderConfig: fmt.Sprintf(registry, fp.Setup.Requirement.Source),
 			Instances: []json.InstanceObjectStateV4{
 				{
 					SchemaVersion: uint64(fp.Resource.GetTerraformSchemaVersion()),
diff --git a/pkg/terraform/store.go b/pkg/terraform/store.go
index 09ffae94..1ad290ef 100644
--- a/pkg/terraform/store.go
+++ b/pkg/terraform/store.go
@@ -47,6 +47,9 @@ type ProviderRequirement struct {
 
 	// Version of the provider. An example value is "4.0"
 	Version string
+
+	// Registry of the provider. An example value is `provider["registry.terraform.io/%s"]`
+	Registry string
 }
 
 // ProviderConfiguration holds the setup configuration body