Merge pull request #13 from DIMO-Network/custom-formater

Add optional formating for custom generator.
DIMO-Network · May 23, 2024 · 2175803 · 2175803
2 parents 8085430 + 9bd05cf
commit 2175803
Show file tree

Hide file tree

Showing 7 changed files with 34 additions and 12 deletions.
diff --git a/cmd/codegen/codegen.go b/cmd/codegen/codegen.go
@@ -25,9 +25,10 @@ func main() {
 	generators := flag.String("generators", "all", "Comma separated list of generators to run. Options: all, model, convert, custom.")
 	// Convert flags
 	withTest := flag.Bool("convert.with-test", false, "Generate test functions for conversion functions. Default is true.")
-	// GQL flags
-	gqlOutFile := flag.String("custom.output-file", "", "Path of the generate gql file that is appened to the outputDir.")
-	gqlTemplateFile := flag.String("custom.template-file", "", "Path to the template file. Which is executed with codegen.TemplateData data.")
+	// Custom flags
+	customOutFile := flag.String("custom.output-file", "", "Path of the generate gql file that is appened to the outputDir.")
+	customTemplateFile := flag.String("custom.template-file", "", "Path to the template file. Which is executed with codegen.TemplateData data.")
+	customFormat := flag.Bool("custom.format", false, "Format the generated file with goimports.")
 	// Migration flags
 	migrationFileName := flag.String("migration.file-name", "", "Name of the migration file. Default is the model name.")
 
@@ -64,8 +65,9 @@ func main() {
 		PackageName: *packageName,
 		OutputDir:   *outputDir,
 		Custom: custom.Config{
-			OutputFile:   *gqlOutFile,
-			TemplateFile: *gqlTemplateFile,
+			OutputFile:   *customOutFile,
+			TemplateFile: *customTemplateFile,
+			Format:       *customFormat,
 		},
 		Convert: convert.Config{
 			WithTest: *withTest,

diff --git a/internal/generator/custom/custom.go b/internal/generator/custom/custom.go
@@ -2,13 +2,15 @@
 package custom
 
 import (
+	"bytes"
 	"fmt"
 	"os"
 	"path"
 	"path/filepath"
 	"strings"
 	"text/template"
 
+	"github.com/DIMO-Network/model-garage/pkg/codegen"
 	"github.com/DIMO-Network/model-garage/pkg/schema"
 )
 
@@ -18,9 +20,10 @@ var customFileFormat = "%s.txt"
 type Config struct {
 	// OutputFile is the name of the model to generate the custom file.
 	OutputFile string
-
 	// TemplateFile is the path to the template file.
 	TemplateFile string
+	// Format controls whether the generated file is formatted with goimports.
+	Format bool
 }
 
 // Generate creates a new Custom file.
@@ -49,10 +52,19 @@ func Generate(tmplData *schema.TemplateData, outputDir string, cfg Config) error
 		}
 	}()
 
-	err = customFileTmpl.Execute(customFileOutputFile, &tmplData)
+	var outBuf bytes.Buffer
+	err = customFileTmpl.Execute(&outBuf, &tmplData)
 	if err != nil {
 		return fmt.Errorf("error executing Custom template: %w", err)
 	}
+	if cfg.Format {
+		err = codegen.FormatAndWriteToFile(outBuf.Bytes(), filePath)
+		if err != nil {
+			return fmt.Errorf("error writing file: %w", err)
+		}
+		return nil
+	}
+	_, _ = customFileOutputFile.Write(outBuf.Bytes())
 
 	return nil
 }

diff --git a/pkg/clickhouseinfra/clickhouseinfra.go b/pkg/clickhouseinfra/clickhouseinfra.go
@@ -19,8 +19,9 @@ import (
 )
 
 const (
-	defaultUser      = "default"
-	defaultDB        = "dimo"
+	defaultUser = "default"
+	defaultDB   = "dimo"
+	// SecureNativePort is the secure port for the ClickHouse container.
 	SecureNativePort = nat.Port("9440/tcp")
 )
 
@@ -95,7 +96,8 @@ func (c *Container) GetClickHouseAsConn() (clickhouse.Conn, error) {
 			Database: c.DbName,
 		},
 		TLS: &tls.Config{
-			RootCAs: c.RootCAs,
+			MinVersion: tls.VersionTLS12,
+			RootCAs:    c.RootCAs,
 		},
 	})
 	if err != nil {
@@ -119,7 +121,8 @@ func (c *Container) GetClickhouseAsDB(ctx context.Context) (*sql.DB, error) {
 			Database: c.DbName,
 		},
 		TLS: &tls.Config{
-			RootCAs: c.RootCAs,
+			MinVersion: tls.VersionTLS12,
+			RootCAs:    c.RootCAs,
 		},
 	})
 	const retries = 3

diff --git a/pkg/migrations/migrations.go b/pkg/migrations/migrations.go
@@ -1,3 +1,4 @@
+// Package migrations provides the functionality to run goose migrations on a clickhouse database.
 package migrations
 
 import (

diff --git a/pkg/schema/signal.go b/pkg/schema/signal.go
@@ -1,3 +1,4 @@
+// Package schema contains the types and functions for generating the schema from the spec and definition files.
 package schema
 
 import (

diff --git a/pkg/vss/convert/payloadv1.go b/pkg/vss/convert/payloadv1.go
@@ -32,6 +32,9 @@ func SignalsFromV1Payload(ctx context.Context, tokenGetter TokenIDGetter, jsonDa
 	}
 
 	source, err := SourceFromV1Data(jsonData)
+	if err != nil {
+		return nil, fmt.Errorf("error getting source: %w", err)
+	}
 	baseSignal := vss.Signal{
 		TokenID:   tokenID,
 		Timestamp: ts,

diff --git a/pkg/vss/convert/payloadv2_test.go b/pkg/vss/convert/payloadv2_test.go
@@ -129,7 +129,7 @@ var fullV2InputJSON = `{
 }`
 
 var expectedV2Signals = []vss.Signal{
-	{TokenID: tokenID, Timestamp: time.Date(2024, time.April, 18, 17, 20, 26, 633000000, time.UTC), Name: "powertrainCombustionEngineECT", ValueNumber: 107, ValueString: "", Source: "dimo/integration/123"}, //nolint // false positive
+	{TokenID: tokenID, Timestamp: time.Date(2024, time.April, 18, 17, 20, 26, 633000000, time.UTC), Name: "powertrainCombustionEngineECT", ValueNumber: 107, ValueString: "", Source: "dimo/integration/123"},
 	{TokenID: tokenID, Timestamp: time.Date(2024, time.April, 18, 17, 20, 27, 173000000, time.UTC), Name: "powertrainCombustionEngineMAF", ValueNumber: 475.79, ValueString: "", Source: "dimo/integration/123"},
 	{TokenID: tokenID, Timestamp: time.Date(2024, time.April, 18, 17, 20, 29, 314000000, time.UTC), Name: "oBDEngineLoad", ValueNumber: 12.54912, ValueString: "", Source: "dimo/integration/123"},
 	{TokenID: tokenID, Timestamp: time.Date(2024, time.April, 18, 17, 20, 29, 844000000, time.UTC), Name: "powertrainCombustionEngineTPS", ValueNumber: 23.529600000000002, ValueString: "", Source: "dimo/integration/123"},