schema generate + schema create

parseablehq · Dec 16, 2024 · dc9380f · dc9380f
1 parent e370625
commit dc9380f
Show file tree

Hide file tree

Showing 2 changed files with 104 additions and 18 deletions.
diff --git a/cmd/generate.go b/cmd/generate.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net/http"
 	"os"
+	"pb/pkg/common"
 	internalHTTP "pb/pkg/http"
 
 	"github.com/spf13/cobra"
@@ -17,24 +18,24 @@ const (
 )
 
 var GenerateSchemaCmd = &cobra.Command{
-	Use:     "schema",
+	Use:     "generate",
 	Short:   "Generate Schema for JSON",
-	Example: "pb generate schema --file=test.json",
+	Example: "pb schema generate --file=test.json",
 	RunE: func(cmd *cobra.Command, _ []string) error {
 		// Get the file path from the `--file` flag
 		filePath, err := cmd.Flags().GetString("file")
 		if err != nil {
-			return fmt.Errorf("failed to read file flag: %w", err)
+			return fmt.Errorf(common.Red+"failed to read file flag: %w"+common.Reset, err)
 		}
 
 		if filePath == "" {
-			return fmt.Errorf("file flag is required")
+			return fmt.Errorf(common.Red + "file flag is required" + common.Reset)
 		}
 
 		// Read the file content
 		fileContent, err := os.ReadFile(filePath)
 		if err != nil {
-			return fmt.Errorf("failed to read file %s: %w", filePath, err)
+			return fmt.Errorf(common.Red+"failed to read file %s: %w"+common.Reset, filePath, err)
 		}
 
 		// Initialize HTTP client
@@ -43,7 +44,7 @@ var GenerateSchemaCmd = &cobra.Command{
 		// Create the HTTP request
 		req, err := client.NewRequest(http.MethodPost, generateStaticSchemaPath, bytes.NewBuffer(fileContent))
 		if err != nil {
-			return fmt.Errorf("failed to create new request: %w", err)
+			return fmt.Errorf(common.Red+"failed to create new request: %w"+common.Reset, err)
 		}
 
 		// Set Content-Type header
@@ -52,34 +53,109 @@ var GenerateSchemaCmd = &cobra.Command{
 		// Execute the request
 		resp, err := client.Client.Do(req)
 		if err != nil {
-			return fmt.Errorf("request execution failed: %w", err)
+			return fmt.Errorf(common.Red+"request execution failed: %w"+common.Reset, err)
 		}
 		defer resp.Body.Close()
 
 		// Check for non-200 status codes
 		if resp.StatusCode != http.StatusOK {
 			body, _ := io.ReadAll(resp.Body)
-			fmt.Printf("Error response: %s\n", string(body))
-			return fmt.Errorf("non-200 status code received: %s", resp.Status)
+			fmt.Printf(common.Red+"Error response: %s\n"+common.Reset, string(body))
+			return fmt.Errorf(common.Red+"non-200 status code received: %s"+common.Reset, resp.Status)
 		}
 
 		// Parse and print the response
 		respBody, err := io.ReadAll(resp.Body)
 		if err != nil {
-			return fmt.Errorf("failed to read response body: %w", err)
+			return fmt.Errorf(common.Red+"failed to read response body: %w"+common.Reset, err)
 		}
 
 		var prettyJSON bytes.Buffer
 		if err := json.Indent(&prettyJSON, respBody, "", "  "); err != nil {
-			return fmt.Errorf("failed to format response as JSON: %w", err)
+			return fmt.Errorf(common.Red+"failed to format response as JSON: %w"+common.Reset, err)
 		}
 
-		fmt.Println(prettyJSON.String())
+		fmt.Println(common.Green + prettyJSON.String() + common.Reset)
+		return nil
+	},
+}
+
+var CreateSchemaCmd = &cobra.Command{
+	Use:     "create",
+	Short:   "Create Schema for a Parseable stream",
+	Example: "pb schema create --stream=my_stream --file=schema.json",
+	RunE: func(cmd *cobra.Command, _ []string) error {
+		// Get the stream name from the `--stream` flag
+		streamName, err := cmd.Flags().GetString("stream")
+		if err != nil {
+			return fmt.Errorf(common.Red+"failed to read stream flag: %w"+common.Reset, err)
+		}
+
+		if streamName == "" {
+			return fmt.Errorf(common.Red + "stream flag is required" + common.Reset)
+		}
+
+		// Get the file path from the `--file` flag
+		filePath, err := cmd.Flags().GetString("file")
+		if err != nil {
+			return fmt.Errorf(common.Red+"failed to read config flag: %w"+common.Reset, err)
+		}
+
+		if filePath == "" {
+			return fmt.Errorf(common.Red + "file path flag is required" + common.Reset)
+		}
+
+		// Read the JSON schema file
+		schemaContent, err := os.ReadFile(filePath)
+		if err != nil {
+			return fmt.Errorf(common.Red+"failed to read schema file %s: %w"+common.Reset, filePath, err)
+		}
+
+		// Initialize HTTP client
+		client := internalHTTP.DefaultClient(&DefaultProfile)
+
+		// Construct the API path
+		apiPath := fmt.Sprintf("/logstream/%s", streamName)
+
+		// Create the HTTP PUT request
+		req, err := client.NewRequest(http.MethodPut, apiPath, bytes.NewBuffer(schemaContent))
+		if err != nil {
+			return fmt.Errorf(common.Red+"failed to create new request: %w"+common.Reset, err)
+		}
+
+		// Set custom headers
+		req.Header.Set("Content-Type", "application/json")
+		req.Header.Set("X-P-Static-Schema-Flag", "true")
+
+		// Execute the request
+		resp, err := client.Client.Do(req)
+		if err != nil {
+			return fmt.Errorf(common.Red+"request execution failed: %w"+common.Reset, err)
+		}
+		defer resp.Body.Close()
+
+		// Check for non-200 status codes
+		if resp.StatusCode != http.StatusOK {
+			body, _ := io.ReadAll(resp.Body)
+			fmt.Printf(common.Red+"Error response: %s\n"+common.Reset, string(body))
+			return fmt.Errorf(common.Red+"non-200 status code received: %s"+common.Reset, resp.Status)
+		}
+
+		// Parse and print the response
+		respBody, err := io.ReadAll(resp.Body)
+		if err != nil {
+			return fmt.Errorf(common.Red+"failed to read response body: %w"+common.Reset, err)
+		}
+
+		fmt.Println(common.Green + string(respBody) + common.Reset)
 		return nil
 	},
 }
 
 func init() {
 	// Add the `--file` flag to the command
 	GenerateSchemaCmd.Flags().StringP("file", "f", "", "Path to the JSON file to generate schema")
+	CreateSchemaCmd.Flags().StringP("stream", "s", "", "Name of the stream to associate with the schema")
+	CreateSchemaCmd.Flags().StringP("file", "f", "", "Path to the JSON file to create schema")
+
 }
diff --git a/main.go b/main.go
@@ -92,10 +92,19 @@ var profile = &cobra.Command{
 	},
 }
 
-var generate = &cobra.Command{
-	Use:               "generate",
-	Short:             "Generate schema for json data files",
-	Long:              "\nuse generate schema to generate schema for json files.",
+var schema = &cobra.Command{
+	Use:   "schema",
+	Short: "Generate or create schemas for JSON data or Parseable streams",
+	Long: `The "schema" command allows you to either:
+  - Generate a schema automatically from a JSON file for analysis or integration.
+  - Create a custom schema for Parseable streams (PB streams) to structure and process your data.
+
+Examples:
+  - To generate a schema from a JSON file:
+      pb schema generate --file=data.json
+  - To create a schema for a PB stream:
+      pb schema create --stream-name=my_stream --config=data.json
+`,
 	PersistentPreRunE: combinedPreRun,
 	PersistentPostRun: func(cmd *cobra.Command, args []string) {
 		if os.Getenv("PB_ANALYTICS") == "disable" {
@@ -217,7 +226,8 @@ func main() {
 	query.AddCommand(pb.QueryCmd)
 	query.AddCommand(pb.SavedQueryList)
 
-	generate.AddCommand(pb.GenerateSchemaCmd)
+	schema.AddCommand(pb.GenerateSchemaCmd)
+	schema.AddCommand(pb.CreateSchemaCmd)
 
 	install.AddCommand(pb.InstallOssCmd)
 
@@ -230,7 +240,7 @@ func main() {
 
 	cli.AddCommand(pb.AutocompleteCmd)
 	cli.AddCommand(install)
-	cli.AddCommand(generate)
+	cli.AddCommand(schema)
 
 	// Set as command
 	pb.VersionCmd.Run = func(_ *cobra.Command, _ []string) {