Back to protogen package approach

pablojimpas · Jul 10, 2024 · 7670819 · 7670819
1 parent 252c3f0
commit 7670819
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 96 deletions.
diff --git a/cmd/protoc-gen-sqlc/main.go b/cmd/protoc-gen-sqlc/main.go
@@ -5,40 +5,24 @@
 package main
 
 import (
-	"fmt"
-	"log/slog"
-	"os"
+	"flag"
 
 	"github.com/pablojimpas/protoc-gen-sqlc/internal/converter"
-
-	"google.golang.org/protobuf/proto"
+	"github.com/pablojimpas/protoc-gen-sqlc/internal/sqlc/template"
+	"google.golang.org/protobuf/compiler/protogen"
 	"google.golang.org/protobuf/types/pluginpb"
 )
 
 func main() {
-	resp, err := converter.ConvertFrom(os.Stdin)
-	if err != nil {
-		message := fmt.Sprintf("Failed to read input: %v", err)
-		slog.Error(message)
-		renderResponse(&pluginpb.CodeGeneratorResponse{
-			Error: &message,
-		})
-		os.Exit(1)
-	}
-
-	renderResponse(resp)
-}
-
-func renderResponse(resp *pluginpb.CodeGeneratorResponse) {
-	data, err := proto.Marshal(resp)
-	if err != nil {
-		slog.Error("failed to marshal response", slog.Any("error", err))
-		return
-	}
-
-	_, err = os.Stdout.Write(data)
-	if err != nil {
-		slog.Error("failed to write response", slog.Any("error", err))
-		return
-	}
+	protogen.Options{
+		ParamFunc: flag.CommandLine.Set,
+	}.Run(
+		func(p *protogen.Plugin) error {
+			p.SupportedFeatures = uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
+			opts := template.Options{}
+			converter.GenerateSchema(p, opts)
+			converter.GenerateQueries(p, opts)
+			return nil
+		},
+	)
 }
diff --git a/internal/converter/converter.go b/internal/converter/converter.go
@@ -6,75 +6,16 @@ package converter
 
 import (
 	"fmt"
-	"io"
 	"log/slog"
 
 	"github.com/pablojimpas/protoc-gen-sqlc/internal/core"
 	"github.com/pablojimpas/protoc-gen-sqlc/internal/sqlc/template"
 	"google.golang.org/protobuf/compiler/protogen"
-	"google.golang.org/protobuf/proto"
-	"google.golang.org/protobuf/reflect/protodesc"
-	"google.golang.org/protobuf/types/descriptorpb"
-	"google.golang.org/protobuf/types/pluginpb"
 )
 
-func ConvertFrom(rd io.Reader) (*pluginpb.CodeGeneratorResponse, error) {
-	slog.SetLogLoggerLevel(slog.LevelDebug)
-
-	input, err := io.ReadAll(rd)
-	if err != nil {
-		return nil, fmt.Errorf("failed to read request: %w", err)
-	}
-
-	req := &pluginpb.CodeGeneratorRequest{}
-	err = proto.Unmarshal(input, req)
-	if err != nil {
-		return nil, fmt.Errorf("can't unmarshal input: %w", err)
-	}
-
-	return Convert(req)
-}
-
-func Convert(req *pluginpb.CodeGeneratorRequest) (*pluginpb.CodeGeneratorResponse, error) {
-	files := []*pluginpb.CodeGeneratorResponse_File{}
-	genFiles := make(map[string]struct{}, len(req.FileToGenerate))
-	for _, file := range req.FileToGenerate {
-		genFiles[file] = struct{}{}
-	}
-
-	// We need this to resolve dependencies when making protodesc versions of the files
-	resolver, err := protodesc.NewFiles(&descriptorpb.FileDescriptorSet{
-		File: req.GetProtoFile(),
-	})
-	if err != nil {
-		return nil, err
-	}
-
-	for _, fileDesc := range req.GetProtoFile() {
-		if _, ok := genFiles[fileDesc.GetName()]; !ok {
-			slog.Debug("skip generating file because it wasn't requested", slog.String("name", fileDesc.GetName()))
-			continue
-		}
-
-		slog.Debug("generating file", slog.String("name", fileDesc.GetName()))
-
-		_, err := protodesc.NewFile(fileDesc, resolver)
-		if err != nil {
-			slog.Error("error loading file", slog.Any("error", err))
-			return nil, err
-		}
-	}
-
-	features := uint64(pluginpb.CodeGeneratorResponse_FEATURE_PROTO3_OPTIONAL)
-	return &pluginpb.CodeGeneratorResponse{
-		File:              files,
-		SupportedFeatures: &features,
-	}, nil
-}
-
-func generateSchema(p *protogen.Plugin, opts template.Options) {
-	slog.Debug("generating schema.pb.sql")
-	gf := p.NewGeneratedFile("schema.pb.sql", "")
+func GenerateSchema(p *protogen.Plugin, opts template.Options) {
+	slog.Debug("generating schema.sql")
+	gf := p.NewGeneratedFile("schema.sql", "")
 
 	err := template.ApplySchema(gf, template.SchemaParams{Schema: schema, Options: opts, HeaderParams: template.HeaderParams{}})
 	if err != nil {
@@ -83,7 +24,7 @@ func generateSchema(p *protogen.Plugin, opts template.Options) {
 	}
 }
 
-func generateQueries(p *protogen.Plugin, opts template.Options) {
+func GenerateQueries(p *protogen.Plugin, opts template.Options) {
 	for _, name := range p.Request.FileToGenerate {
 		f := p.FilesByPath[name]
 
@@ -93,13 +34,13 @@ func generateQueries(p *protogen.Plugin, opts template.Options) {
 		}
 
 		slog.Debug("processing queries for file", slog.String("name", name))
-		slog.Debug("generating queries in", slog.String("name", fmt.Sprintf("%s.pb.sql", f.GeneratedFilenamePrefix)))
+		slog.Debug("generating queries in", slog.String("name", fmt.Sprintf("%s.sql", f.GeneratedFilenamePrefix)))
 
-		gf := p.NewGeneratedFile(fmt.Sprintf("%s.pb.sql", f.GeneratedFilenamePrefix), f.GoImportPath)
+		gf := p.NewGeneratedFile(fmt.Sprintf("%s.sql", f.GeneratedFilenamePrefix), f.GoImportPath)
 
 		err := template.ApplyCrud(gf, template.CrudParams{
 			GoName:       f.Messages[0].GoIdent.GoName,
-			PrimaryKey:   "id",
+			PrimaryKey:   schema.Tables[0].Constraints[0].Columns[0],
 			Table:        schema.Tables[0],
 			Options:      opts,
 			HeaderParams: template.HeaderParams{Sources: []string{*f.Proto.Name}},