Support TFJob v1beta1 crd validation

cmd/root.go

@@ -28,13 +28,9 @@ var cfgFile string
 // RootCmd represents the base command when called without any subcommands
 var RootCmd = &cobra.Command{
 	Use:   "crd-validation",
-	Short: "A brief description of your application",
-	Long: `A longer description that spans multiple lines and likely contains
-examples and usage of using your application. For example:
-
-Cobra is a CLI library for Go that empowers applications.
-This application is a tool to generate the needed files
-to quickly create a Cobra application.`,
+	Short: "Generator for kubeflow crd validation schema",
+	Long: `crd-validation is a CLI library to generate the needed files to validate custom objects
+	passed by user via OpenAPI v3 schema.`,
 	// Uncomment the following line if your bare application
 	// has an action associated with it:
 	//	Run: func(cmd *cobra.Command, args []string) { },
@@ -52,13 +48,7 @@ func Execute() {
 func init() {
 	cobra.OnInitialize(initConfig)
 
-	// Here you will define your flags and configuration settings.
-	// Cobra supports Persistent Flags, which, if defined here,
-	// will be global for your application.
-
 	RootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.crd-validation.yaml)")
-	// Cobra also supports local flags, which will only run
-	// when this action is called directly.
 	RootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
 }
 

cmd/tfjob.go

@@ -15,6 +15,8 @@
 package cmd
 
 import (
+	"fmt"
+
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
 
@@ -48,11 +50,14 @@ var tfjobCmd = &cobra.Command{
 
 func generateTFJob() {
 	original := config.NewCustomResourceDefinition("tfjob")
-	var outputDir string
+	var outputDir, jobVersion string
 	if viper.Get("global") != nil {
-		outputDir = viper.Get("global").(map[string]interface{})["output"].(string)
+		outputDir = viper.GetString("global.output")
+		jobVersion = viper.GetString("tfjob.spec.version")
 	}
-	generator := crd.NewTFJobGenerator(outputDir)
+	filename := fmt.Sprintf("tfjob-crd-%v.yaml", jobVersion)
+
+	generator := crd.NewTFJobGenerator(jobVersion)
 	final := generator.Generate(original)
-	generator.Export(final)
+	generator.Export(final, outputDir, filename)
 }

pkg/crd/exporter/exporter.go

@@ -13,21 +13,16 @@ import (
 
 // Exporter exports the yaml config to file.
 type Exporter struct {
-	outputDir string
-	filename  string
-	writer    *bufio.Writer
+	writer *bufio.Writer
 }
 
-func New(outputDir, filename string) *Exporter {
-	return &Exporter{
-		outputDir: outputDir,
-		filename:  filename,
-	}
+func New() *Exporter {
+	return &Exporter{}
 }
 
 // Export exports the yaml config to file.
-func (e *Exporter) Export(final *apiextensions.CustomResourceDefinition) {
-	file, err := filepath.Abs(filepath.Join(e.outputDir, e.filename))
+func (e *Exporter) Export(final *apiextensions.CustomResourceDefinition, outputDir string, filename string) {
+	file, err := filepath.Abs(filepath.Join(outputDir, filename))
 	if err != nil {
 		log.Fatalf("Failed to open the output file %s: %v", file, err)
 	}

pkg/crd/tfjob.go

@@ -1,36 +1,44 @@
 package crd
 
 import (
-	tfv1alpha2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
-	log "github.com/sirupsen/logrus"
-	apiextensions "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
+	"fmt"
 
 	"github.com/kubeflow/crd-validation/pkg/crd/exporter"
 	"github.com/kubeflow/crd-validation/pkg/utils"
-)
-
-const (
-	// CRDName is the name for TFJob.
-	CRDName = "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2.TFJob"
-
-	generatedFile = "tfjob-crd-v1alpha2.yaml"
+	tfv1alpha2 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1alpha2"
+	tfv1beta1 "github.com/kubeflow/tf-operator/pkg/apis/tensorflow/v1beta1"
+	log "github.com/sirupsen/logrus"
+	apiextensions "k8s.io/apiextensions-apiserver/pkg/apis/apiextensions/v1beta1"
 )
 
 // TFJobGenerator is the type for TFJob CRD generator.
 type TFJobGenerator struct {
 	*exporter.Exporter
+	jobVersion string
 }
 
 // NewTFJobGenerator creates a new TFJob CRD generator.
-func NewTFJobGenerator(outputDir string) *TFJobGenerator {
+func NewTFJobGenerator(version string) *TFJobGenerator {
 	return &TFJobGenerator{
-		Exporter: exporter.New(outputDir, generatedFile),
+		Exporter:   exporter.New(),
+		jobVersion: version,
 	}
 }
 
 // Generate generates the crd.
-func (t TFJobGenerator) Generate(original *apiextensions.CustomResourceDefinition) *apiextensions.CustomResourceDefinition {
+func (t *TFJobGenerator) Generate(original *apiextensions.CustomResourceDefinition) *apiextensions.CustomResourceDefinition {
 	log.Println("Generating validation for TFJob")
-	original.Spec.Validation = utils.GetCustomResourceValidation(CRDName, tfv1alpha2.GetOpenAPIDefinitions)
+
+	// CRDName is the name for TFJob.
+	CRDName := fmt.Sprintf("github.com/kubeflow/tf-operator/pkg/apis/tensorflow/%v.TFJob", t.jobVersion)
+
+	if t.jobVersion == "v1alpha2" {
+		original.Spec.Validation = utils.GetCustomResourceValidation(CRDName, tfv1alpha2.GetOpenAPIDefinitions)
+	} else if t.jobVersion == "v1beta1" {
+		original.Spec.Validation = utils.GetCustomResourceValidation(CRDName, tfv1beta1.GetOpenAPIDefinitions)
+	} else {
+		log.Errorf("Invalid version of TFJob %s", t.jobVersion)
+	}
+
 	return original
 }