From a07fbcec8fa56197f042c05170bba3bf66882bcf Mon Sep 17 00:00:00 2001
From: Alex Jeannopoulos <alexj@backpocket.com>
Date: Wed, 8 Jul 2020 11:20:37 -0400
Subject: [PATCH] added gogo.proto check using GOPATH not hardcoded

---
 go.mod  |  1 +
 main.go | 36 +++++++++++++++++++-----------------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/go.mod b/go.mod
index 7ae1304..de86cb7 100644
--- a/go.mod
+++ b/go.mod
@@ -24,6 +24,7 @@ require (
 	github.com/karrick/godirwalk v1.15.6 // indirect
 	github.com/kr/pretty v0.2.0 // indirect
 	github.com/lib/pq v1.3.0
+	github.com/logrusorgru/aurora v2.0.3+incompatible // indirect
 	github.com/mattn/go-sqlite3 v2.0.3+incompatible
 	github.com/mwitkow/go-proto-validators v0.3.0 // indirect
 	github.com/ompluscator/dynamic-struct v1.2.0
diff --git a/main.go b/main.go
index f73bfda..20e64c5 100644
--- a/main.go
+++ b/main.go
@@ -5,11 +5,11 @@ import (
 	"database/sql"
 	"encoding/json"
 	"fmt"
+	"go/build"
 	"io"
 	"io/ioutil"
 	"os"
 	"os/exec"
-	"os/user"
 	"path/filepath"
 	"strings"
 
@@ -20,6 +20,7 @@ import (
 	"github.com/jimsmart/schema"
 	_ "github.com/jinzhu/gorm/dialects/mysql"
 	_ "github.com/lib/pq"
+	"github.com/logrusorgru/aurora"
 	_ "github.com/mattn/go-sqlite3"
 
 	"github.com/smallnest/gen/dbmeta"
@@ -46,6 +47,7 @@ var (
 	module          = goopt.String([]string{"--module"}, "example.com/example", "module path")
 	overwrite       = goopt.Flag([]string{"--overwrite"}, []string{"--no-overwrite"}, "Overwrite existing files (default)", "disable overwriting files")
 	windows         = goopt.Flag([]string{"--windows"}, []string{}, "use windows line endings in generated files", "")
+	noColorOutput   = goopt.Flag([]string{"--no-color"}, []string{}, "disable color output", "")
 
 	contextFileName  = goopt.String([]string{"--context"}, "", "context file (json) to populate context with")
 	mappingFileName  = goopt.String([]string{"--mapping"}, "", "mapping file (json) to map sql types to golang/protobuf etc")
@@ -89,6 +91,7 @@ var (
 
 	baseTemplates *packr.Box
 	tableInfos    map[string]*dbmeta.ModelInfo
+	au            aurora.Aurora
 )
 
 func init() {
@@ -147,6 +150,7 @@ func main() {
 	//for i, arg := range os.Args {
 	//	fmt.Printf("[%2d] %s\n", i, arg)
 	//}
+	au = aurora.NewAurora(!*noColorOutput)
 
 	baseTemplates = packr.New("gen", "./template")
 
@@ -722,27 +726,25 @@ func createProtocCmdLine(protoBufDir, protoBufOutDir, protoBufFile string) ([]st
 		}
 	}
 
-	usr, err := user.Current()
-	if err == nil {
-		dir := usr.HomeDir
-		srcPath := filepath.Join(dir, "go/src")
-
-		//srcDirExists := dbmeta.Exists(srcPath)
+	gopath := os.Getenv("GOPATH")
+	if gopath == "" {
+		gopath = build.Default.GOPATH
+	}
 
-		gogoPath := filepath.Join(dir, "go/src/github.com/gogo/protobuf/gogoproto/gogo.proto")
-		gogoImportExists := dbmeta.Exists(gogoPath)
+	srcPath := filepath.Join(gopath, "src")
+	gogoPath := filepath.Join(gopath, "src/github.com/gogo/protobuf/gogoproto/gogo.proto")
+	gogoImportExists := dbmeta.Exists(gogoPath)
 
-		//fmt.Printf("path    : %s   srcDirExists: %t\n", srcPath, srcDirExists)
-		//fmt.Printf("gogoPath: %s   gogoImportExists: %t\n", gogoPath, gogoImportExists)
+	//fmt.Printf("path    : %s   srcDirExists: %t\n", srcPath, srcDirExists)
+	//fmt.Printf("gogoPath: %s   gogoImportExists: %t\n", gogoPath, gogoImportExists)
 
-		if !gogoImportExists {
-			fmt.Printf("github.com/gogo/protobuf/gogoproto/gogo.proto does not exist on path - install with\ngo get -u github.com/gogo/protobuf/proto\n\n")
-			return nil, fmt.Errorf("github.com/gogo/protobuf/gogoproto/gogo.proto does not exist")
-		}
-
-		*gogoProtoImport = srcPath
+	if !gogoImportExists {
+		fmt.Print(au.Red("github.com/gogo/protobuf/gogoproto/gogo.proto does not exist on path - install with\ngo get -u github.com/gogo/protobuf/proto\n\n"))
+		return nil, fmt.Errorf("github.com/gogo/protobuf/gogoproto/gogo.proto does not exist")
 	}
 
+	*gogoProtoImport = srcPath
+
 	fmt.Printf("----------------------------\n")
 
 	args := []string{fmt.Sprintf("-I%s", *gogoProtoImport),