From cfc76eb88fad1927b0c2057b359c88e841e043b6 Mon Sep 17 00:00:00 2001
From: Jim <jlambert@hashicorp.com>
Date: Tue, 22 Sep 2020 10:41:42 -0400
Subject: [PATCH] add support for WithSessionIds in ListSesions (#395)

---
 internal/session/options.go                 |  9 +++++++++
 internal/session/options_test.go            |  7 +++++++
 internal/session/repository_session.go      | 11 ++++++++++-
 internal/session/repository_session_test.go | 18 ++++++++++++++++++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/internal/session/options.go b/internal/session/options.go
index 17aecc4ec..161037f60 100644
--- a/internal/session/options.go
+++ b/internal/session/options.go
@@ -25,6 +25,7 @@ type options struct {
 	withExpirationTime *timestamp.Timestamp
 	withTestTofu       []byte
 	withListingConvert bool
+	withSessionIds     []string
 }
 
 func getDefaultOptions() options {
@@ -75,6 +76,14 @@ func WithTestTofu(tofu []byte) Option {
 	}
 }
 
+// WithSessionIds allows the specification of the session ids to use for the
+// operation.
+func WithSessionIds(ids ...string) Option {
+	return func(o *options) {
+		o.withSessionIds = ids
+	}
+}
+
 func withListingConvert(withListingConvert bool) Option {
 	return func(o *options) {
 		o.withListingConvert = withListingConvert
diff --git a/internal/session/options_test.go b/internal/session/options_test.go
index 6e37efffd..6595bd266 100644
--- a/internal/session/options_test.go
+++ b/internal/session/options_test.go
@@ -70,4 +70,11 @@ func Test_GetOpts(t *testing.T) {
 		copy(testOpts.withTestTofu, tofu)
 		assert.Equal(opts, testOpts)
 	})
+	t.Run("WithSessionIds", func(t *testing.T) {
+		assert := assert.New(t)
+		opts := getOpts(WithSessionIds("s_1", "s_2", "s_3"))
+		testOpts := getDefaultOptions()
+		testOpts.withSessionIds = []string{"s_1", "s_2", "s_3"}
+		assert.Equal(opts, testOpts)
+	})
 }
diff --git a/internal/session/repository_session.go b/internal/session/repository_session.go
index 59ff74229..8a4b7bc57 100644
--- a/internal/session/repository_session.go
+++ b/internal/session/repository_session.go
@@ -162,7 +162,7 @@ func (r *Repository) LookupSession(ctx context.Context, sessionId string, opt ..
 	return &session, authzSummary, nil
 }
 
-// ListSessions will sessions.  Supports the WithLimit, and WithScopeId options.
+// ListSessions will sessions.  Supports the WithLimit, WithScopeId and WithSessionIds options.
 func (r *Repository) ListSessions(ctx context.Context, opt ...Option) ([]*Session, error) {
 	opts := getOpts(opt...)
 	var where []string
@@ -177,6 +177,15 @@ func (r *Repository) ListSessions(ctx context.Context, opt ...Option) ([]*Sessio
 		inClauseCnt += 1
 		where, args = append(where, fmt.Sprintf("user_id = $%d", inClauseCnt)), append(args, opts.withUserId)
 	}
+	if len(opts.withSessionIds) > 0 {
+		idsInClause := make([]string, 0, len(opts.withSessionIds))
+		for _, id := range opts.withSessionIds {
+			inClauseCnt += 1
+			idsInClause, args = append(idsInClause, fmt.Sprintf("$%d", inClauseCnt)), append(args, id)
+		}
+		where = append(where, fmt.Sprintf("s.public_id in(%s)", strings.Join(idsInClause, ",")))
+	}
+
 	var limit string
 	switch {
 	case opts.withLimit < 0: // any negative number signals unlimited results
diff --git a/internal/session/repository_session_test.go b/internal/session/repository_session_test.go
index 9931d72bb..fcbf8775c 100644
--- a/internal/session/repository_session_test.go
+++ b/internal/session/repository_session_test.go
@@ -138,6 +138,24 @@ func TestRepository_ListSession(t *testing.T) {
 		assert.Equal(1, len(got))
 		assert.Equal(got[0].UserId, s.UserId)
 	})
+	t.Run("WithSessionIds", func(t *testing.T) {
+		assert, require := assert.New(t), require.New(t)
+		require.NoError(conn.Where("1=1").Delete(AllocSession()).Error)
+		testSessions := []*Session{}
+		for i := 0; i < 10; i++ {
+			s := TestSession(t, conn, wrapper, composedOf)
+			_ = TestState(t, conn, s.PublicId, StatusActive)
+			testSessions = append(testSessions, s)
+		}
+		assert.Equal(10, len(testSessions))
+		withIds := []string{testSessions[0].PublicId, testSessions[1].PublicId}
+		conn.LogMode(true)
+		got, err := repo.ListSessions(context.Background(), WithSessionIds(withIds...), WithOrder("create_time asc"))
+		require.NoError(err)
+		assert.Equal(2, len(got))
+		assert.Equal(StatusActive, got[0].States[0].Status)
+		assert.Equal(StatusPending, got[0].States[1].Status)
+	})
 }
 
 func TestRepository_CreateSession(t *testing.T) {