github.com/tilt-dev/tilt@v0.36.0/internal/tiltfile/tiltextension/plugin.go

// Package extension implements Tilt extensions.
// This is not the internal Starkit abstraction, but the user-visible feature.
// In a Tiltfile, you can write `load("ext://foo", "bar")` to load the function bar
// from the extension foo.
package tiltextension

import (
	"context"
	"fmt"
	"strings"

	"go.starlark.net/starlark"

	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/tilt-dev/tilt/internal/controllers/apiset"
	"github.com/tilt-dev/tilt/internal/controllers/core/extension"
	"github.com/tilt-dev/tilt/internal/controllers/core/extensionrepo"
	"github.com/tilt-dev/tilt/internal/tiltfile/starkit"
	tiltfilev1alpha1 "github.com/tilt-dev/tilt/internal/tiltfile/v1alpha1"
	"github.com/tilt-dev/tilt/pkg/apis"
	"github.com/tilt-dev/tilt/pkg/apis/core/v1alpha1"
	"github.com/tilt-dev/tilt/pkg/logger"
)

const (
	extensionPrefix = "ext://"
	defaultRepoName = "default"
)

type Plugin struct {
	repoReconciler ExtRepoReconciler
	extReconciler  ExtReconciler
}

func NewPlugin(repoReconciler *extensionrepo.Reconciler, extReconciler *extension.Reconciler) *Plugin {
	return &Plugin{
		repoReconciler: repoReconciler,
		extReconciler:  extReconciler,
	}
}

func NewFakePlugin(repoReconciler ExtRepoReconciler, extReconciler ExtReconciler) *Plugin {
	return &Plugin{
		repoReconciler: repoReconciler,
		extReconciler:  extReconciler,
	}
}

type State struct {
	ExtsLoaded map[string]bool
}

func (e Plugin) NewState() interface{} {
	return State{
		ExtsLoaded: make(map[string]bool),
	}
}

func (e *Plugin) OnStart(env *starkit.Environment) error {
	env.AddLoadInterceptor(e)
	return nil
}

func (e *Plugin) recordExtensionLoaded(ctx context.Context, t *starlark.Thread, moduleName string) {
	err := starkit.SetState(t, func(existing State) (State, error) {
		existing.ExtsLoaded[moduleName] = true
		return existing, nil
	})
	if err != nil {
		logger.Get(ctx).Debugf("error updating state on Tilt extensions loader: %v", err)
	}
}

func (e *Plugin) LocalPath(t *starlark.Thread, arg string) (localPath string, err error) {
	if !strings.HasPrefix(arg, extensionPrefix) {
		return "", nil
	}

	ctx, err := starkit.ContextFromThread(t)
	if err != nil {
		return "", err
	}

	moduleName := strings.TrimPrefix(arg, extensionPrefix)
	defer func() {
		if err == nil {
			// NOTE(maia): Maybe in future we want to track if there was an error or not?
			// For now, only record on successful load.
			e.recordExtensionLoaded(ctx, t, moduleName)
		}
	}()

	starkitModel, err := starkit.ModelFromThread(t)
	if err != nil {
		return "", err
	}

	objSet, err := tiltfilev1alpha1.GetState(starkitModel)
	if err != nil {
		return "", err
	}

	ext := e.ensureExtension(t, objSet, moduleName)
	repo := e.ensureRepo(t, objSet, ext.Spec.RepoName)
	repoStatus := e.repoReconciler.ForceApply(ctx, repo)
	if repoStatus.Error != "" {
		return "", fmt.Errorf("loading extension repo %s: %s", repo.Name, repoStatus.Error)
	}
	if repoStatus.Path == "" {
		return "", fmt.Errorf("extension repo not resolved: %s", repo.Name)
	}

	repoResolved := repo.DeepCopy()
	repoResolved.Status = repoStatus
	extStatus := e.extReconciler.ForceApply(ext, repoResolved)
	if extStatus.Error != "" {
		return "", fmt.Errorf("loading extension %s: %s", ext.Name, extStatus.Error)
	}
	if extStatus.Path == "" {
		return "", fmt.Errorf("extension not resolved: %s", ext.Name)
	}

	return extStatus.Path, nil
}

// Check to see if an extension has already been registered.
//
// If it has, returns the existing object (which should only have a spec).
//
// Otherwise, infers an extension object that points to the default repo.
func (e *Plugin) ensureExtension(t *starlark.Thread, objSet apiset.ObjectSet, moduleName string) *v1alpha1.Extension {
	extName := apis.SanitizeName(moduleName)

	extSet := objSet.GetOrCreateTypedSet(&v1alpha1.Extension{})
	if existing, exists := extSet[extName]; exists {
		ext := existing.(*v1alpha1.Extension)
		metav1.SetMetaDataAnnotation(&ext.ObjectMeta, v1alpha1.AnnotationManagedBy, "tiltfile.loader")
		return ext
	}

	repoSet := objSet.GetOrCreateTypedSet(&v1alpha1.ExtensionRepo{})

	return e.registerExtension(t, extSet, repoSet, extName, moduleName)
}

// In cases where an extension is not already registered, this function will search for an extension
// repo that can satisfy the requested extension, with a fallback to an extension in the default
// repository.
func (e *Plugin) registerExtension(t *starlark.Thread, extSet, repoSet apiset.TypedObjectSet, extName, moduleName string) *v1alpha1.Extension {
	loadHost, extPath, tryRegister := strings.Cut(moduleName, "/")

	// Safety fallback (in case this is called without already previously calling ensureExtension)
	existing := extSet[extName]
	if existing != nil {
		return existing.(*v1alpha1.Extension)
	}

	// If the supplied module name does not contain a / then there's no point looking for matching
	// extension repositories. We can just return an extension named extName in the default
	// repository
	if !tryRegister {
		return e.registerDefaultExtension(t, extSet, extName, moduleName)
	}

	// Otherwise, look for a repository that can satisfy this lookup
	for _, v := range repoSet {
		repo := v.(*v1alpha1.ExtensionRepo)
		if repo.Spec.LoadHost == "" || repo.Spec.LoadHost != loadHost {
			continue
		}

		// This repo load_host matches the first component of the module name
		// So we can register this as an extension on that repo
		ext := &v1alpha1.Extension{
			ObjectMeta: metav1.ObjectMeta{
				Name: extName,
				Annotations: map[string]string{
					v1alpha1.AnnotationManagedBy: "tiltfile.loader",
				},
			},
			Spec: v1alpha1.ExtensionSpec{
				RepoName: repo.Name,
				RepoPath: extPath,
			},
		}

		extSet[extName] = ext
		return ext
	}

	return e.registerDefaultExtension(t, extSet, extName, moduleName)
}

// Registers an extension named moduleName in the default extension repository. Used as a fallback.
func (e *Plugin) registerDefaultExtension(t *starlark.Thread, extSet apiset.TypedObjectSet, extName, moduleName string) *v1alpha1.Extension {
	// Safety fallback
	existing := extSet[extName]
	if existing != nil {
		return existing.(*v1alpha1.Extension)
	}

	defaultExt := &v1alpha1.Extension{
		ObjectMeta: metav1.ObjectMeta{
			Name: extName,
			Annotations: map[string]string{
				v1alpha1.AnnotationManagedBy: "tiltfile.loader",
			},
		},
		Spec: v1alpha1.ExtensionSpec{
			RepoName: defaultRepoName,
			RepoPath: moduleName,
		},
	}

	extSet[extName] = defaultExt
	return defaultExt
}

// Check to see if an extension repo has already been registered.
//
// If it has, returns the existing object (which should only have a spec).
//
// Otherwise, register the default repo.
func (e *Plugin) ensureRepo(t *starlark.Thread, objSet apiset.ObjectSet, repoName string) *v1alpha1.ExtensionRepo {
	defaultRepo := &v1alpha1.ExtensionRepo{
		ObjectMeta: metav1.ObjectMeta{
			Name: repoName,
		},
		Spec: v1alpha1.ExtensionRepoSpec{
			URL: "https://github.com/tilt-dev/tilt-extensions",
		},
	}

	typedSet := objSet.GetOrCreateTypedSet(defaultRepo)
	existing, exists := typedSet[repoName]
	if exists {
		return existing.(*v1alpha1.ExtensionRepo)
	}

	typedSet[repoName] = defaultRepo
	return defaultRepo
}

var (
	_ starkit.LoadInterceptor = (*Plugin)(nil)
	_ starkit.StatefulPlugin  = (*Plugin)(nil)
)

func MustState(model starkit.Model) State {
	state, err := GetState(model)
	if err != nil {
		panic(err)
	}
	return state
}

func GetState(m starkit.Model) (State, error) {
	var state State
	err := m.Load(&state)
	return state, err
}