From b691de474cc2292b2ebc551a97307533b401ff9b Mon Sep 17 00:00:00 2001
From: letusfly85 <letusfly95@gmail.com>
Date: Fri, 12 Feb 2016 19:21:10 +0900
Subject: [PATCH] added optional import for other package structure

---
 gen.go       | 6 ++++--
 gen_funcs.go | 4 +++-
 2 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/gen.go b/gen.go
index 3f62cb7..e6f41ed 100644
--- a/gen.go
+++ b/gen.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"os"
+	"strings"
 
 	"github.com/drone/sqlgen/parse"
 	"github.com/drone/sqlgen/schema"
@@ -20,6 +21,7 @@ var (
 	genSchema  = flag.Bool("schema", true, "generate sql schema and queries")
 	genFuncs   = flag.Bool("funcs", true, "generate sql helper functions")
 	extraFuncs = flag.Bool("extras", true, "generate extra sql helper functions")
+	importPath = ""
 )
 
 func main() {
@@ -36,7 +38,7 @@ func main() {
 	// if the code is generated in a different folder
 	// that the struct we need to import the struct
 	if tree.Pkg != *pkgName && *pkgName != "main" {
-		// TODO
+		importPath = strings.Join([]string{*pkgName, tree.Pkg}, "/")
 	}
 
 	// load the Tree into a schema Object
@@ -47,7 +49,7 @@ func main() {
 
 	if *genFuncs {
 		writePackage(&buf, *pkgName)
-		writeImports(&buf, tree, "database/sql")
+		writeImports(&buf, tree, "database/sql", importPath)
 		writeRowFunc(&buf, tree)
 		writeRowsFunc(&buf, tree)
 		writeSliceFunc(&buf, tree)
diff --git a/gen_funcs.go b/gen_funcs.go
index b7eebdc..1df99f2 100644
--- a/gen_funcs.go
+++ b/gen_funcs.go
@@ -43,7 +43,9 @@ func writeImports(w io.Writer, tree *parse.Node, pkgs ...string) {
 	// encoder package that was specified.
 	fmt.Fprintln(w, "\nimport (")
 	for pkg, _ := range pmap {
-		fmt.Fprintf(w, "\t%q\n", pkg)
+		if pkg != "" {
+			fmt.Fprintf(w, "\t%q\n", pkg)
+		}
 	}
 	fmt.Fprintln(w, ")")
 }