feat: load embedded provider schemas for providers found in stacks fi…

…le while early decoding

internal/features/modules/jobs/schema.go

@@ -5,34 +5,25 @@ package jobs
 
 import (
 	"context"
-	"encoding/json"
 	"errors"
 	"fmt"
-	"io"
 	"io/fs"
 	"log"
 	"time"
 
 	"github.com/hashicorp/go-multierror"
 	"github.com/hashicorp/go-version"
 	"github.com/hashicorp/hcl-lang/lang"
-	tfjson "github.com/hashicorp/terraform-json"
 	"github.com/hashicorp/terraform-ls/internal/document"
 	"github.com/hashicorp/terraform-ls/internal/features/modules/state"
 	"github.com/hashicorp/terraform-ls/internal/job"
 	"github.com/hashicorp/terraform-ls/internal/registry"
-	"github.com/hashicorp/terraform-ls/internal/schemas"
 	globalState "github.com/hashicorp/terraform-ls/internal/state"
 	op "github.com/hashicorp/terraform-ls/internal/terraform/module/operation"
 	tfaddr "github.com/hashicorp/terraform-registry-address"
 	tfregistry "github.com/hashicorp/terraform-schema/registry"
-	tfschema "github.com/hashicorp/terraform-schema/schema"
 	"github.com/zclconf/go-cty/cty"
 	ctyjson "github.com/zclconf/go-cty/cty/json"
-	"go.opentelemetry.io/otel"
-	"go.opentelemetry.io/otel/attribute"
-	"go.opentelemetry.io/otel/codes"
-	"go.opentelemetry.io/otel/trace"
 )
 
 const tracerName = "github.com/hashicorp/terraform-ls/internal/terraform/module"
@@ -73,7 +64,7 @@ func PreloadEmbeddedSchema(ctx context.Context, logger *log.Logger, fs fs.ReadDi
 	}
 
 	for _, pAddr := range missingReqs {
-		err := preloadSchemaForProviderAddr(ctx, pAddr, fs, schemaStore, logger)
+		err := globalState.PreloadSchemaForProviderAddr(ctx, pAddr, fs, schemaStore, logger)
 		if err != nil {
 			return err
 		}
@@ -82,119 +73,6 @@ func PreloadEmbeddedSchema(ctx context.Context, logger *log.Logger, fs fs.ReadDi
 	return nil
 }
 
-func preloadSchemaForProviderAddr(ctx context.Context, pAddr tfaddr.Provider, fs fs.ReadDirFS,
-	schemaStore *globalState.ProviderSchemaStore, logger *log.Logger) error {
-
-	startTime := time.Now()
-
-	if pAddr.IsLegacy() && pAddr.Type == "terraform" {
-		// The terraform provider is built into Terraform 0.11+
-		// and while it's possible, users typically don't declare
-		// entry in required_providers block for it.
-		pAddr = tfaddr.NewProvider(tfaddr.BuiltInProviderHost, tfaddr.BuiltInProviderNamespace, "terraform")
-	} else if pAddr.IsLegacy() {
-		// Since we use recent version of Terraform to generate
-		// embedded schemas, these will never contain legacy
-		// addresses.
-		//
-		// A legacy namespace may come from missing
-		// required_providers entry & implied requirement
-		// from the provider block or 0.12-style entry,
-		// such as { grafana = "1.0" }.
-		//
-		// Implying "hashicorp" namespace here mimics behaviour
-		// of all recent (0.14+) Terraform versions.
-		originalAddr := pAddr
-		pAddr.Namespace = "hashicorp"
-		logger.Printf("preloading schema for %s (implying %s)",
-			originalAddr.ForDisplay(), pAddr.ForDisplay())
-	}
-
-	ctx, rootSpan := otel.Tracer(tracerName).Start(ctx, "preloadProviderSchema",
-		trace.WithAttributes(attribute.KeyValue{
-			Key:   attribute.Key("ProviderAddress"),
-			Value: attribute.StringValue(pAddr.String()),
-		}))
-	defer rootSpan.End()
-
-	pSchemaFile, err := schemas.FindProviderSchemaFile(fs, pAddr)
-	if err != nil {
-		rootSpan.RecordError(err)
-		rootSpan.SetStatus(codes.Error, "schema file not found")
-		if errors.Is(err, schemas.SchemaNotAvailable{Addr: pAddr}) {
-			logger.Printf("preloaded schema not available for %s", pAddr)
-			return nil
-		}
-		return err
-	}
-
-	_, span := otel.Tracer(tracerName).Start(ctx, "readProviderSchemaFile",
-		trace.WithAttributes(attribute.KeyValue{
-			Key:   attribute.Key("ProviderAddress"),
-			Value: attribute.StringValue(pAddr.String()),
-		}))
-	b, err := io.ReadAll(pSchemaFile.File)
-	if err != nil {
-		span.RecordError(err)
-		span.SetStatus(codes.Error, "schema file not readable")
-		return err
-	}
-	span.SetStatus(codes.Ok, "schema file read successfully")
-	span.End()
-
-	_, span = otel.Tracer(tracerName).Start(ctx, "decodeProviderSchemaData",
-		trace.WithAttributes(attribute.KeyValue{
-			Key:   attribute.Key("ProviderAddress"),
-			Value: attribute.StringValue(pAddr.String()),
-		}))
-	jsonSchemas := tfjson.ProviderSchemas{}
-	err = json.Unmarshal(b, &jsonSchemas)
-	if err != nil {
-		span.RecordError(err)
-		span.SetStatus(codes.Error, "schema file not decodable")
-		return err
-	}
-	span.SetStatus(codes.Ok, "schema data decoded successfully")
-	span.End()
-
-	ps, ok := jsonSchemas.Schemas[pAddr.String()]
-	if !ok {
-		return fmt.Errorf("%q: no schema found in file", pAddr)
-	}
-
-	pSchema := tfschema.ProviderSchemaFromJson(ps, pAddr)
-	pSchema.SetProviderVersion(pAddr, pSchemaFile.Version)
-
-	_, span = otel.Tracer(tracerName).Start(ctx, "loadProviderSchemaDataIntoMemDb",
-		trace.WithAttributes(attribute.KeyValue{
-			Key:   attribute.Key("ProviderAddress"),
-			Value: attribute.StringValue(pAddr.String()),
-		}))
-	err = schemaStore.AddPreloadedSchema(pAddr, pSchemaFile.Version, pSchema)
-	if err != nil {
-		span.RecordError(err)
-		span.SetStatus(codes.Error, "loading schema into mem-db failed")
-		span.End()
-		existsError := &globalState.AlreadyExistsError{}
-		if errors.As(err, &existsError) {
-			// This accounts for a possible race condition
-			// where we may be preloading the same schema
-			// for different providers at the same time
-			logger.Printf("schema for %s is already loaded", pAddr)
-			return nil
-		}
-		return err
-	}
-	span.SetStatus(codes.Ok, "schema loaded successfully")
-	span.End()
-
-	elapsedTime := time.Since(startTime)
-	logger.Printf("preloaded schema for %s %s in %s", pAddr, pSchemaFile.Version, elapsedTime)
-	rootSpan.SetStatus(codes.Ok, "schema loaded successfully")
-
-	return nil
-}
-
 // GetModuleDataFromRegistry obtains data about any modules (inputs & outputs)
 // from the Registry API based on module calls which were previously parsed
 // via [LoadModuleMetadata]. The same data could be obtained via [ParseModuleManifest]

internal/features/stacks/events.go

@@ -14,6 +14,7 @@ import (
 	"github.com/hashicorp/terraform-ls/internal/job"
 	"github.com/hashicorp/terraform-ls/internal/lsp"
 	"github.com/hashicorp/terraform-ls/internal/protocol"
+	"github.com/hashicorp/terraform-ls/internal/schemas"
 	globalAst "github.com/hashicorp/terraform-ls/internal/terraform/ast"
 	"github.com/hashicorp/terraform-ls/internal/terraform/module/operation"
 )
@@ -200,9 +201,23 @@ func (f *StacksFeature) decodeStack(ctx context.Context, dir document.DirHandle,
 	}
 	ids = append(ids, metaId)
 
+	eSchemaId, err := f.stateStore.JobStore.EnqueueJob(ctx, job.Job{
+		Dir: dir,
+		Func: func(ctx context.Context) error {
+			return jobs.PreloadEmbeddedSchema(ctx, f.logger, schemas.FS,
+				f.store, f.stateStore.ProviderSchemas, path)
+		},
+		DependsOn:   job.IDs{metaId},
+		Type:        operation.OpTypePreloadEmbeddedSchema.String(),
+		IgnoreState: ignoreState,
+	})
+	if err != nil {
+		return ids, err
+	}
+	ids = append(ids, eSchemaId)
+
 	// TODO: Implement the following functions where appropriate to stacks
 	// Future: decodeDeclaredModuleCalls(ctx, dir, ignoreState)
-	// TODO: PreloadEmbeddedSchema(ctx, f.logger, schemas.FS,
 	// Future: DecodeReferenceTargets(ctx, f.Store, f.rootFeature, path)
 	// Future: DecodeReferenceOrigins(ctx, f.Store, f.rootFeature, path)
 

internal/features/stacks/jobs/schema.go

@@ -0,0 +1,62 @@
+// Copyright (c) HashiCorp, Inc.
+// SPDX-License-Identifier: MPL-2.0
+
+package jobs
+
+import (
+	"context"
+	"io/fs"
+	"log"
+
+	"github.com/hashicorp/go-version"
+	"github.com/hashicorp/terraform-ls/internal/document"
+	"github.com/hashicorp/terraform-ls/internal/features/stacks/state"
+	"github.com/hashicorp/terraform-ls/internal/job"
+	globalState "github.com/hashicorp/terraform-ls/internal/state"
+	"github.com/hashicorp/terraform-ls/internal/terraform/module/operation"
+	tfaddr "github.com/hashicorp/terraform-registry-address"
+)
+
+func PreloadEmbeddedSchema(ctx context.Context, logger *log.Logger, fs fs.ReadDirFS, stackStore *state.StackStore, schemaStore *globalState.ProviderSchemaStore, stackPath string) error {
+	record, err := stackStore.StackRecordByPath(stackPath)
+
+	if err != nil {
+		return err
+	}
+
+	// Avoid preloading schema if it is already in progress or already known
+	if record.PreloadEmbeddedSchemaState != operation.OpStateUnknown && !job.IgnoreState(ctx) {
+		return job.StateNotChangedErr{Dir: document.DirHandleFromPath(stackPath)}
+	}
+
+	err = stackStore.SetPreloadEmbeddedSchemaState(stackPath, operation.OpStateLoading)
+	if err != nil {
+		return err
+	}
+	defer stackStore.SetPreloadEmbeddedSchemaState(stackPath, operation.OpStateLoaded)
+
+	pReqs := make(map[tfaddr.Provider]version.Constraints, len(record.Meta.ProviderRequirements))
+	for _, req := range record.Meta.ProviderRequirements {
+		pReqs[req.Source] = req.VersionConstraints
+	}
+
+	missingReqs, err := schemaStore.MissingSchemas(pReqs)
+	if err != nil {
+		return err
+	}
+
+	if len(missingReqs) == 0 {
+		// avoid preloading any schemas if we already have all
+		return nil
+	}
+
+	for _, pAddr := range missingReqs {
+		err := globalState.PreloadSchemaForProviderAddr(ctx, pAddr, fs, schemaStore, logger)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+
+}

internal/features/stacks/jobs/version_test.go

@@ -22,7 +22,7 @@ func TestLoadTerraformVersion(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	ss, err := state.NewStackStore(gs.ChangeStore)
+	ss, err := state.NewStackStore(gs.ChangeStore, gs.ProviderSchemas)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -87,7 +87,7 @@ func TestLoadTerraformVersion_invalid(t *testing.T) {
 			if err != nil {
 				t.Fatal(err)
 			}
-			ss, err := state.NewStackStore(gs.ChangeStore)
+			ss, err := state.NewStackStore(gs.ChangeStore, gs.ProviderSchemas)
 			if err != nil {
 				t.Fatal(err)
 			}

internal/features/stacks/stacks_feature.go

@@ -28,7 +28,7 @@ type StacksFeature struct {
 }
 
 func NewStacksFeature(bus *eventbus.EventBus, stateStore *globalState.StateStore, fs jobs.ReadOnlyFS) (*StacksFeature, error) {
-	store, err := state.NewStackStore(stateStore.ChangeStore)
+	store, err := state.NewStackStore(stateStore.ChangeStore, stateStore.ProviderSchemas)
 	if err != nil {
 		return nil, err
 	}

internal/features/stacks/state/schema.go

@@ -30,7 +30,7 @@ var dbSchema = &memdb.DBSchema{
 	},
 }
 
-func NewStackStore(changeStore *globalState.ChangeStore) (*StackStore, error) {
+func NewStackStore(changeStore *globalState.ChangeStore, providerSchemasStore *globalState.ProviderSchemaStore) (*StackStore, error) {
 	db, err := memdb.NewMemDB(dbSchema)
 	if err != nil {
 		return nil, err
@@ -39,9 +39,10 @@ func NewStackStore(changeStore *globalState.ChangeStore) (*StackStore, error) {
 	discardLogger := log.New(io.Discard, "", 0)
 
 	return &StackStore{
-		db:          db,
-		tableName:   stackTableName,
-		logger:      discardLogger,
-		changeStore: changeStore,
+		db:                   db,
+		tableName:            stackTableName,
+		logger:               discardLogger,
+		changeStore:          changeStore,
+		providerSchemasStore: providerSchemasStore,
 	}, nil
 }

internal/features/stacks/state/stack_record.go

@@ -16,6 +16,10 @@ import (
 type StackRecord struct {
 	path string
 
+	// PreloadEmbeddedSchemaState tracks if we tried loading all provider
+	// schemas from our embedded schema data
+	PreloadEmbeddedSchemaState operation.OpState
+
 	Meta      StackMetadata
 	MetaErr   error
 	MetaState operation.OpState
@@ -42,10 +46,17 @@ func (m *StackRecord) Copy() *StackRecord {
 	}
 
 	newRecord := &StackRecord{
-		path:             m.path,
+		path: m.path,
+
+		PreloadEmbeddedSchemaState: m.PreloadEmbeddedSchemaState,
+
 		Meta:             m.Meta.Copy(),
 		ParsingErr:       m.ParsingErr,
 		DiagnosticsState: m.DiagnosticsState.Copy(),
+
+		RequiredTerraformVersion:      m.RequiredTerraformVersion,
+		RequiredTerraformVersionErr:   m.RequiredTerraformVersionErr,
+		RequiredTerraformVersionState: m.RequiredTerraformVersionState,
 	}
 
 	if m.ParsedFiles != nil {

internal/features/stacks/state/stack_store.go

@@ -21,7 +21,8 @@ type StackStore struct {
 	tableName string
 	logger    *log.Logger
 
-	changeStore *globalState.ChangeStore
+	changeStore          *globalState.ChangeStore
+	providerSchemasStore *globalState.ProviderSchemaStore
 }
 
 func (s *StackStore) SetLogger(logger *log.Logger) {
@@ -314,6 +315,25 @@ func (s *StackStore) UpdateMetadata(path string, meta *tfstack.Meta, mErr error)
 	return nil
 }
 
+func (s *StackStore) SetPreloadEmbeddedSchemaState(path string, state operation.OpState) error {
+	txn := s.db.Txn(true)
+	defer txn.Abort()
+
+	record, err := stackCopyByPath(txn, path)
+	if err != nil {
+		return err
+	}
+
+	record.PreloadEmbeddedSchemaState = state
+	err = txn.Insert(s.tableName, record)
+	if err != nil {
+		return err
+	}
+
+	txn.Commit()
+	return nil
+}
+
 func (s *StackStore) add(txn *memdb.Txn, stackPath string) error {
 	// TODO: Introduce Exists method to Txn?
 	obj, err := txn.First(s.tableName, "id", stackPath)