From 4aec9c7a20e5d27f1c96ff7dd42cbb38f08e9b80 Mon Sep 17 00:00:00 2001
From: Rafael Jesus <rafaelljesus86@gmail.com>
Date: Wed, 21 Feb 2018 23:32:50 +0100
Subject: [PATCH] proper support for inner join query creation

---
 pkg/config/config.go      |  4 +++-
 pkg/dumper/generic/sql.go |  2 ++
 pkg/reader/generic/sql.go | 13 +++++++++----
 pkg/reader/reader.go      |  5 ++++-
 4 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/pkg/config/config.go b/pkg/config/config.go
index 94936cb..4cf9a83 100644
--- a/pkg/config/config.go
+++ b/pkg/config/config.go
@@ -13,6 +13,7 @@ type (
 	// Table represents a klepto table definition
 	Table struct {
 		Name          string
+		PrimaryKey    string
 		IgnoreData    bool
 		Filter        Filter
 		Anonymise     map[string]string
@@ -28,9 +29,10 @@ type (
 
 	// Relationship represents a relationship definition
 	Relationship struct {
+		Table           string
+		ForeignKey      string
 		ReferencedTable string
 		ReferencedKey   string
-		ForeignKey      string
 	}
 )
 
diff --git a/pkg/dumper/generic/sql.go b/pkg/dumper/generic/sql.go
index 60fd285..6a0b43c 100644
--- a/pkg/dumper/generic/sql.go
+++ b/pkg/dumper/generic/sql.go
@@ -94,6 +94,7 @@ func (p *sqlDumper) readAndDumpTables(done chan<- struct{}, configTables config.
 			}
 
 			opts = reader.ReadTableOpt{
+				PrimaryKey:    tableConfig.PrimaryKey,
 				Match:         tableConfig.Filter.Match,
 				Limit:         tableConfig.Filter.Limit,
 				Relationships: p.relationshipConfigToOptions(tableConfig.Relationships),
@@ -144,6 +145,7 @@ func (p *sqlDumper) relationshipConfigToOptions(relationshipsConfig []*config.Re
 
 	for _, r := range relationshipsConfig {
 		opts = append(opts, &reader.RelationshipOpt{
+			Table:           r.Table,
 			ReferencedTable: r.ReferencedTable,
 			ReferencedKey:   r.ReferencedKey,
 			ForeignKey:      r.ForeignKey,
diff --git a/pkg/reader/generic/sql.go b/pkg/reader/generic/sql.go
index 1ab1a45..b46d625 100644
--- a/pkg/reader/generic/sql.go
+++ b/pkg/reader/generic/sql.go
@@ -117,20 +117,25 @@ func (s *SqlReader) buildQuery(tableName string, opts reader.ReadTableOpt) (sq.S
 	var query sq.SelectBuilder
 
 	query = sq.Select(opts.Columns...).From(s.QuoteIdentifier(tableName))
-
 	for _, r := range opts.Relationships {
+		if r.Table == "" {
+			r.Table = tableName
+		}
 		query = query.Join(fmt.Sprintf(
 			"%s ON %s.%s = %s.%s",
 			r.ReferencedTable,
-			tableName,
-			r.ForeignKey,
 			r.ReferencedTable,
 			r.ReferencedKey,
+			r.Table,
+			r.ForeignKey,
 		))
 	}
 
 	if len(opts.Relationships) > 0 {
-		query = query.GroupBy(fmt.Sprintf("%s.id", tableName))
+		if opts.PrimaryKey == "" {
+			opts.PrimaryKey = "id"
+		}
+		query = query.GroupBy(fmt.Sprintf("%s.%s", tableName, opts.PrimaryKey))
 	}
 
 	if opts.Match != "" {
diff --git a/pkg/reader/reader.go b/pkg/reader/reader.go
index 55d3889..1c9caed 100644
--- a/pkg/reader/reader.go
+++ b/pkg/reader/reader.go
@@ -31,6 +31,8 @@ type (
 
 	// ReadTableOpt represents the read table options
 	ReadTableOpt struct {
+		// PrimaryKey is the primary key
+		PrimaryKey string
 		// Columns contains the (quoted) column of the table
 		Columns []string
 		// Match is a condition field to dump only certain amount data
@@ -43,9 +45,10 @@ type (
 
 	// RelationshipOpt represents the relationships options
 	RelationshipOpt struct {
+		Table           string
+		ForeignKey      string
 		ReferencedTable string
 		ReferencedKey   string
-		ForeignKey      string
 	}
 
 	// ConnOpts are the options to create a connection