add DISTINCT ON

mgutz · Jul 8, 2015 · f4f89d4 · f4f89d4
1 parent f679e19
commit f4f89d4
Show file tree

Hide file tree

Showing 6 changed files with 133 additions and 3 deletions.
diff --git a/select.go b/select.go
@@ -5,6 +5,7 @@ type SelectBuilder struct {
 	Execer
 
 	isDistinct      bool
+	distinctColumns []string
 	isInterpolated  bool
 	columns         []string
 	table           string
@@ -34,6 +35,13 @@ func (b *SelectBuilder) Distinct() *SelectBuilder {
 	return b
 }
 
+// DistinctOn sets the columns for DISTINCT ON
+func (b *SelectBuilder) DistinctOn(columns ...string) *SelectBuilder {
+	b.isDistinct = true
+	b.distinctColumns = columns
+	return b
+}
+
 // From sets the table to SELECT FROM. JOINs may also be defined here.
 func (b *SelectBuilder) From(from string) *SelectBuilder {
 	b.table = from
@@ -118,7 +126,18 @@ func (b *SelectBuilder) ToSQL() (string, []interface{}) {
 	buf.WriteString("SELECT ")
 
 	if b.isDistinct {
-		buf.WriteString("DISTINCT ")
+		if len(b.distinctColumns) == 0 {
+			buf.WriteString("DISTINCT ")
+		} else {
+			buf.WriteString("DISTINCT ON (")
+			for i, s := range b.distinctColumns {
+				if i > 0 {
+					buf.WriteString(", ")
+				}
+				buf.WriteString(s)
+			}
+			buf.WriteString(") ")
+		}
 	}
 
 	for i, s := range b.columns {

diff --git a/select_doc.go b/select_doc.go
@@ -100,7 +100,18 @@ func (b *SelectDocBuilder) ToSQL() (string, []interface{}) {
 	}
 
 	if b.isDistinct {
-		buf.WriteString("DISTINCT ")
+		if len(b.distinctColumns) == 0 {
+			buf.WriteString("DISTINCT ")
+		} else {
+			buf.WriteString("DISTINCT ON (")
+			for i, s := range b.distinctColumns {
+				if i > 0 {
+					buf.WriteString(", ")
+				}
+				buf.WriteString(s)
+			}
+			buf.WriteString(") ")
+		}
 	}
 
 	for i, s := range b.columns {
@@ -224,6 +235,13 @@ func (b *SelectDocBuilder) Distinct() *SelectDocBuilder {
 	return b
 }
 
+// DistinctOn sets the columns for DISTINCT ON
+func (b *SelectDocBuilder) DistinctOn(columns ...string) *SelectDocBuilder {
+	b.isDistinct = true
+	b.distinctColumns = columns
+	return b
+}
+
 // From sets the table to SELECT FROM
 func (b *SelectDocBuilder) From(from string) *SelectDocBuilder {
 	b.table = from

diff --git a/select_doc_test.go b/select_doc_test.go
@@ -130,3 +130,29 @@ func TestDocScopeWhere(t *testing.T) {
 	assert.Equal(t, stripWS(expected), stripWS(sql))
 	assert.Exactly(t, args, []interface{}{1, "published"})
 }
+
+func TestDocDistinctOn(t *testing.T) {
+	published := `
+		INNER JOIN posts p on (p.author_id = u.id)
+		WHERE
+			p.state = $1
+	`
+	sql, args := SelectDoc("u.*, p.*").
+		DistinctOn("aa", "bb").
+		From(`users u`).
+		Scope(published, "published").
+		Where(`u.id = $1`, 1).
+		ToSQL()
+	expected := `
+		SELECT row_to_json(dat__item.*)
+		FROM (
+			SELECT DISTINCT ON (aa, bb)
+			u.*, p.*
+			FROM users u
+				INNER JOIN posts p on (p.author_id = u.id)
+			WHERE (u.id = $1) AND ( p.state = $2 )
+		) as dat__item
+	`
+	assert.Equal(t, stripWS(expected), stripWS(sql))
+	assert.Exactly(t, args, []interface{}{1, "published"})
+}
diff --git a/select_test.go b/select_test.go
@@ -256,3 +256,22 @@ func TestScopeJoinOnly(t *testing.T) {
 	assert.Equal(t, "SELECT u.*, p.* FROM users u INNER JOIN posts p on (p.author_id = u.id) WHERE (u.id = $1)", sql)
 	assert.Exactly(t, args, []interface{}{1})
 }
+
+func TestDistinctOn(t *testing.T) {
+	published := `
+		INNER JOIN posts p on (p.author_id = u.id)
+	`
+
+	sql, args := Select("u.*, p.*").
+		DistinctOn("foo", "bar").
+		From(`users u`).
+		Scope(published).
+		Where(`u.id = $1`, 1).
+		ToSQL()
+	assert.Equal(t, stripWS(`
+		SELECT DISTINCT ON (foo, bar) u.*, p.*
+		FROM users u
+			INNER JOIN posts p on (p.author_id = u.id)
+		WHERE (u.id = $1)`), stripWS(sql))
+	assert.Exactly(t, args, []interface{}{1})
+}
diff --git a/sqlx-runner/select_doc_test.go b/sqlx-runner/select_doc_test.go
@@ -154,7 +154,6 @@ func TestSelectDocRowsNil(t *testing.T) {
 	assert.Equal(sql.ErrNoRows, err)
 }
 
-// Not efficient but it's doable
 func TestSelectDoc(t *testing.T) {
 	assert := assert.New(t)
 
@@ -176,6 +175,28 @@ func TestSelectDoc(t *testing.T) {
 	assert.Equal(1, person.ID)
 }
 
+func TestSelectDocDistinctOn(t *testing.T) {
+	assert := assert.New(t)
+
+	type Person struct {
+		ID    int
+		Name  string
+		Posts []*Post
+	}
+
+	var person Person
+	err := testDB.
+		SelectDoc("id", "name").
+		DistinctOn("id").
+		From("people").
+		Where("id = $1", 1).
+		QueryStruct(&person)
+
+	assert.NoError(err)
+	assert.Equal("Mario", person.Name)
+	assert.Equal(1, person.ID)
+}
+
 func TestSelectQueryEmbeddedJSON(t *testing.T) {
 	s := beginTxWithFixtures()
 	defer s.AutoRollback()

diff --git a/sqlx-runner/select_exec_test.go b/sqlx-runner/select_exec_test.go
@@ -105,6 +105,33 @@ func TestSelectQueryStruct(t *testing.T) {
 	assert.Contains(t, err.Error(), "no rows")
 }
 
+func TestSelectQueryDistinctOn(t *testing.T) {
+	s := beginTxWithFixtures()
+	defer s.AutoRollback()
+
+	// Found:
+	var person Person
+	err := s.
+		Select("id", "name", "email").
+		DistinctOn("id").
+		From("people").
+		Where("email = $1", "john@acme.com").
+		QueryStruct(&person)
+	assert.NoError(t, err)
+	assert.True(t, person.ID > 0)
+	assert.Equal(t, person.Name, "John")
+	assert.True(t, person.Email.Valid)
+	assert.Equal(t, person.Email.String, "john@acme.com")
+
+	// Not found:
+	var person2 Person
+	err = s.
+		Select("id", "name", "email").
+		From("people").Where("email = $1", "dontexist@acme.com").
+		QueryStruct(&person2)
+	assert.Contains(t, err.Error(), "no rows")
+}
+
 func TestSelectBySqlQueryStructs(t *testing.T) {
 	s := beginTxWithFixtures()
 	defer s.AutoRollback()