1 files changed, 165 insertions, 0 deletions
diff --git a/vendor/github.com/authzed/zed/internal/commands/completion.go b/vendor/github.com/authzed/zed/internal/commands/completion.go
new file mode 100644
index 0000000..dd24a74
--- /dev/null
+++ b/vendor/github.com/authzed/zed/internal/commands/completion.go
@@ -0,0 +1,165 @@
+package commands
+
+import (
+	"errors"
+	"strings"
+
+	"github.com/spf13/cobra"
+
+	v1 "github.com/authzed/authzed-go/proto/authzed/api/v1"
+	"github.com/authzed/spicedb/pkg/schemadsl/compiler"
+
+	"github.com/authzed/zed/internal/client"
+)
+
+type CompletionArgumentType int
+
+const (
+	ResourceType CompletionArgumentType = iota
+	ResourceID
+	Permission
+	SubjectType
+	SubjectID
+	SubjectTypeWithOptionalRelation
+)
+
+func FileExtensionCompletions(extension ...string) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return func(_ *cobra.Command, _ []string, _ string) ([]string, cobra.ShellCompDirective) {
+		return extension, cobra.ShellCompDirectiveFilterFileExt
+	}
+}
+
+func GetArgs(fields ...CompletionArgumentType) func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+	return func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
+		// Read the current schema, if any.
+		schema, err := readSchema(cmd)
+		if err != nil {
+			return nil, cobra.ShellCompDirectiveError
+		}
+
+		// Find the specified resource type, if any.
+		var resourceType string
+	loop:
+		for index, arg := range args {
+			field := fields[index]
+			switch field {
+			case ResourceType:
+				resourceType = arg
+				break loop
+
+			case ResourceID:
+				pieces := strings.Split(arg, ":")
+				if len(pieces) >= 1 {
+					resourceType = pieces[0]
+					break loop
+				}
+			}
+		}
+
+		// Handle : on resource and subject IDs.
+		if strings.HasSuffix(toComplete, ":") && (fields[len(args)] == ResourceID || fields[len(args)] == SubjectID) {
+			comps := []string{}
+			comps = cobra.AppendActiveHelp(comps, "Please enter an object ID")
+			return comps, cobra.ShellCompDirectiveNoFileComp
+		}
+
+		// Handle # on subject types. If the toComplete contains a valid subject,
+		// then we should return the relation names. Note that we cannot do this
+		// on the # character because shell autocompletion won't send it to us.
+		if len(args) == len(fields)-1 && toComplete != "" && fields[len(args)] == SubjectTypeWithOptionalRelation {
+			for _, objDef := range schema.ObjectDefinitions {
+				subjectType := toComplete
+				if objDef.Name == subjectType {
+					relationNames := make([]string, 0)
+					relationNames = append(relationNames, subjectType)
+					for _, relation := range objDef.Relation {
+						relationNames = append(relationNames, subjectType+"#"+relation.Name)
+					}
+					return relationNames, cobra.ShellCompDirectiveNoFileComp
+				}
+			}
+		}
+
+		if len(args) >= len(fields) {
+			// If we have all the arguments, return no completions.
+			return nil, cobra.ShellCompDirectiveNoFileComp
+		}
+
+		// Return the completions.
+		currentFieldType := fields[len(args)]
+		switch currentFieldType {
+		case ResourceType:
+			fallthrough
+
+		case SubjectType:
+			fallthrough
+
+		case SubjectID:
+			fallthrough
+
+		case SubjectTypeWithOptionalRelation:
+			fallthrough
+
+		case ResourceID:
+			resourceTypeNames := make([]string, 0, len(schema.ObjectDefinitions))
+			for _, objDef := range schema.ObjectDefinitions {
+				resourceTypeNames = append(resourceTypeNames, objDef.Name)
+			}
+
+			flags := cobra.ShellCompDirectiveNoFileComp
+			if currentFieldType == ResourceID || currentFieldType == SubjectID || currentFieldType == SubjectTypeWithOptionalRelation {
+				flags |= cobra.ShellCompDirectiveNoSpace
+			}
+
+			return resourceTypeNames, flags
+
+		case Permission:
+			if resourceType == "" {
+				return nil, cobra.ShellCompDirectiveNoFileComp
+			}
+
+			relationNames := make([]string, 0)
+			for _, objDef := range schema.ObjectDefinitions {
+				if objDef.Name == resourceType {
+					for _, relation := range objDef.Relation {
+						relationNames = append(relationNames, relation.Name)
+					}
+				}
+			}
+			return relationNames, cobra.ShellCompDirectiveNoFileComp
+		}
+
+		return nil, cobra.ShellCompDirectiveDefault
+	}
+}
+
+func readSchema(cmd *cobra.Command) (*compiler.CompiledSchema, error) {
+	// TODO: we should find a way to cache this
+	client, err := client.NewClient(cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	request := &v1.ReadSchemaRequest{}
+
+	resp, err := client.ReadSchema(cmd.Context(), request)
+	if err != nil {
+		return nil, err
+	}
+
+	schemaText := resp.SchemaText
+	if len(schemaText) == 0 {
+		return nil, errors.New("no schema defined")
+	}
+
+	compiledSchema, err := compiler.Compile(
+		compiler.InputSchema{Source: "schema", SchemaString: schemaText},
+		compiler.AllowUnprefixedObjectType(),
+		compiler.SkipValidation(),
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	return compiledSchema, nil
+}