Add preview level functions for services and methods to OpenAPI code

openapi/code/entity.go

@@ -23,6 +23,16 @@ func (f *Field) IsOptionalObject() bool {
 	return f.Entity != nil && !f.Required && (f.Entity.IsObject() || f.Entity.IsExternal())
 }
 
+// IsPrivatePreview flags object being in private preview.
+func (f *Field) IsPrivatePreview() bool {
+	return f.Schema != nil && isPrivatePreview(&f.Schema.Node)
+}
+
+// IsPublicPreview flags object being in public preview.
+func (f *Field) IsPublicPreview() bool {
+	return f.Schema != nil && isPublicPreview(&f.Schema.Node)
+}
+
 type EnumEntry struct {
 	Named
 	Entity *Entity
@@ -59,6 +69,9 @@ type Entity struct {
 	// if entity has required fields, this is the order of them
 	RequiredOrder []string
 	fields        map[string]Field
+
+	// Schema references the OpenAPI schema this entity was created from.
+	Schema *openapi.Schema
 }
 
 // FullName includes package name and untransformed name of the entity
@@ -242,3 +255,13 @@ func (e *Entity) IsAllRequiredFieldsPrimitive() bool {
 	}
 	return true
 }
+
+// IsPrivatePreview flags object being in private preview.
+func (e *Entity) IsPrivatePreview() bool {
+	return e.Schema != nil && isPrivatePreview(&e.Schema.Node)
+}
+
+// IsPublicPreview flags object being in public preview.
+func (e *Entity) IsPublicPreview() bool {
+	return e.Schema != nil && isPublicPreview(&e.Schema.Node)
+}

openapi/code/load.go

@@ -38,7 +38,7 @@ func NewFromFile(name string) (*Batch, error) {
 			}
 			batch.packages[tag.Package] = pkg
 		}
-		err := pkg.Load(spec, &tag)
+		err := pkg.Load(spec, tag)
 		if err != nil {
 			return nil, fmt.Errorf("fail to load %s: %w", tag.Name, err)
 		}

openapi/code/method.go

@@ -326,3 +326,13 @@ func (m *Method) CanHaveResponseBody() bool {
 func (m *Method) TitleVerb() string {
 	return strings.Title(strings.ToLower(m.Verb))
 }
+
+// IsPrivatePreview flags object being in private preview.
+func (m *Method) IsPrivatePreview() bool {
+	return isPrivatePreview(&m.operation.Node)
+}
+
+// IsPublicPreview flags object being in public preview.
+func (m *Method) IsPublicPreview() bool {
+	return isPublicPreview(&m.operation.Node)
+}

openapi/code/package.go

@@ -142,7 +142,8 @@ func (pkg *Package) schemaToEntity(s *openapi.Schema, path []string, hasName boo
 		Named: Named{
 			Description: s.Description,
 		},
-		enum: map[string]EnumEntry{},
+		Schema: s,
+		enum:   map[string]EnumEntry{},
 	}
 	// pull embedded types up, if they can be defined at package level
 	if s.IsDefinable() && !hasName {
@@ -324,7 +325,7 @@ func (pkg *Package) HasWaits() bool {
 }
 
 // Load takes OpenAPI specification and loads a service model
-func (pkg *Package) Load(spec *openapi.Specification, tag *openapi.Tag) error {
+func (pkg *Package) Load(spec *openapi.Specification, tag openapi.Tag) error {
 	for k, v := range spec.Components.Schemas {
 		split := strings.Split(k, ".")
 		if split[0] != pkg.Name {
@@ -348,6 +349,7 @@ func (pkg *Package) Load(spec *openapi.Specification, tag *openapi.Tag) error {
 						Name:        tag.Service,
 						Description: tag.Description,
 					},
+					tag: &tag,
 				}
 				pkg.services[tag.Service] = svc
 			}

openapi/code/preview.go

@@ -0,0 +1,15 @@
+package code
+
+import (
+	"strings"
+
+	"github.com/databricks/databricks-sdk-go/openapi"
+)
+
+func isPrivatePreview(n *openapi.Node) bool {
+	return strings.ToLower(n.Preview) == "private"
+}
+
+func isPublicPreview(n *openapi.Node) bool {
+	return strings.ToLower(n.Preview) == "public"
+}

openapi/code/service.go

@@ -17,6 +17,7 @@ type Service struct {
 	Package             *Package
 	methods             map[string]*Method
 	ByPathParamsMethods []*Shortcut
+	tag                 *openapi.Tag
 }
 
 // FullName holds package name and service name
@@ -298,3 +299,13 @@ func (svc *Service) Waits() (waits []*Wait) {
 	})
 	return waits
 }
+
+// IsPrivatePreview flags object being in private preview.
+func (svc *Service) IsPrivatePreview() bool {
+	return isPrivatePreview(&svc.tag.Node)
+}
+
+// IsPublicPreview flags object being in public preview.
+func (svc *Service) IsPublicPreview() bool {
+	return isPublicPreview(&svc.tag.Node)
+}

openapi/model.go

@@ -9,9 +9,9 @@ import (
 )
 
 type Node struct {
-	Description  string `json:"description,omitempty"`
-	Availability string `json:"x-databricks-availability,omitempty"`
-	Ref          string `json:"$ref,omitempty"`
+	Description string `json:"description,omitempty"`
+	Preview     string `json:"x-databricks-preview,omitempty"`
+	Ref         string `json:"$ref,omitempty"`
 }
 
 // IsRef flags object being a reference to a component