Merge pull request #4095 from hashicorp/mikemountain-adjust-session-pagination-queries

Update session query with optimized indexes

internal/db/schema/migrations/oss/postgres/28/01_session_create_time_idx.up.sql

@@ -4,6 +4,7 @@
 begin;
 -- Session list queries always use an order by on create_time.
 -- This index can aide in performing this order by.
+-- Dropped in 81/05_session_base_table_updates.up.sql.
 create index
   session_create_time
 on

internal/db/schema/migrations/oss/postgres/50/01_session_list_index.up.sql

@@ -9,6 +9,7 @@ begin;
   -- it will include where clauses that:
   --  * include a project_id paired with a user_id
   --  * and where termination_reason is null
+  -- Dropped in 81/05_session_base_table_updates.up.sql.
   create index session_list_pix on session (project_id, user_id, termination_reason) where termination_reason is null;
   analyze session;
 end;

internal/db/schema/migrations/oss/postgres/81/05_session_base_table_updates.up.sql

@@ -3,11 +3,33 @@
 
 begin;
 
-  -- Add new indexes for the create time and update time queries.
-  create index session_create_time_public_id_idx
-      on session (create_time desc, public_id desc);
-  create index session_update_time_public_id_idx
-      on session (update_time desc, public_id desc);
+  -- Drop old list indexes that interfere with the new list requests.
+  drop index session_create_time, session_list_pix;
+
+  -- Add new indexes for the list queries. These indexes have been
+  -- carefully tested to cover the different expected queries.
+  create index session_project_id_create_time_list_idx
+            on session (project_id,
+                        create_time desc,
+                        public_id   desc,
+                        termination_reason);
+  create index session_project_id_update_time_list_idx
+            on session (project_id,
+                        update_time desc,
+                        public_id   desc,    
+                        termination_reason);
+  create index session_user_id_project_id_create_time_list_idx
+            on session (user_id,
+                        project_id,
+                        create_time desc,
+                        public_id   desc,
+                        termination_reason);
+  create index session_user_id_project_id_update_time_list_idx
+            on session (user_id,
+                        project_id,         
+                        update_time desc,
+                        public_id   desc,
+                        termination_reason);
 
   analyze session;
 

internal/db/sqltest/tests/pagination/session.sql

@@ -2,11 +2,13 @@
 -- SPDX-License-Identifier: BUSL-1.1
 
 begin;
-  select plan(2);
+  select plan(4);
 
-  select has_index('session', 'session_create_time_public_id_idx', array['create_time', 'public_id']);
-  select has_index('session', 'session_update_time_public_id_idx', array['update_time', 'public_id']);
+  select has_index('session', 'session_project_id_create_time_list_idx', array['project_id', 'create_time', 'public_id', 'termination_reason']);
+  select has_index('session', 'session_project_id_update_time_list_idx', array['project_id', 'update_time', 'public_id', 'termination_reason']);
+  select has_index('session', 'session_user_id_project_id_create_time_list_idx', array['user_id', 'project_id', 'create_time', 'public_id', 'termination_reason']);
+  select has_index('session', 'session_user_id_project_id_update_time_list_idx', array['user_id', 'project_id', 'update_time', 'public_id', 'termination_reason']);
 
   select * from finish();
 
-rollback;
+rollback;

internal/session/query.go

@@ -124,27 +124,6 @@ from
 	session_connection_limit, session_connection_count;
 `
 
-	sessionList = `
-with
-session_ids as (
-	select public_id
-	from session as s
-	-- where clause is constructed
-	%s
-	-- order by clause is constructed
-	%s
-	-- limit is constructed
-	%s
-)
-select *
-from session_list
-where
-	session_list.public_id in (select * from session_ids)
--- order by clause again since order from cte is not guaranteed to be preserved
-%s
-;
-`
-
 	terminateSessionIfPossible = `
     -- is terminate_session_id in a canceling state
     with session_version as (
@@ -427,8 +406,7 @@ where session_id = ?
 with session_ids as (
     select public_id
       from session
-           -- search condition for applying permissions is constructed
-     where %s
+     where %s -- search condition for applying permissions is constructed
   order by create_time desc, public_id desc
      limit %d
 )
@@ -442,8 +420,7 @@ with session_ids as (
     select public_id
       from session
      where (create_time, public_id) < (@last_item_create_time, @last_item_id)
-           -- search condition for applying permissions is constructed
-       and %s
+       and %s -- search condition for applying permissions is constructed
   order by create_time desc, public_id desc
      limit %d
 )
@@ -457,8 +434,7 @@ with session_ids as (
     select public_id
       from session
      where update_time > @updated_after_time
-           -- search condition for applying permissions is constructed
-       and %s
+       and %s -- search condition for applying permissions is constructed
   order by update_time desc, public_id desc
      limit %d
 )
@@ -473,8 +449,7 @@ with session_ids as (
       from session
      where update_time > @updated_after_time
        and (update_time, public_id) < (@last_item_update_time, @last_item_id)
-           -- search condition for applying permissions is constructed
-       and %s
+       and %s -- search condition for applying permissions is constructed
   order by update_time desc, public_id desc
      limit %d
 )

internal/session/repository_session.go

@@ -290,16 +290,9 @@ func (r *Repository) listSessions(ctx context.Context, opt ...Option) ([]*Sessio
 
 	opts := getOpts(opt...)
 
-	var permissionWhereClause string
-	if len(where) > 0 {
-		permissionWhereClause = "(" + strings.Join(where, " or ") + ")"
-		if !opts.withTerminated {
-			permissionWhereClause += " and termination_reason is null"
-		}
-	} else {
-		if !opts.withTerminated {
-			permissionWhereClause = " where termination_reason is null"
-		}
+	permissionWhereClause := "(" + strings.Join(where, " or ") + ")"
+	if !opts.withTerminated {
+		permissionWhereClause += " and termination_reason is null"
 	}
 
 	limit := r.defaultLimit
@@ -339,16 +332,9 @@ func (r *Repository) listSessionsRefresh(ctx context.Context, updatedAfter time.
 
 	opts := getOpts(opt...)
 
-	var permissionWhereClause string
-	if len(where) > 0 {
-		permissionWhereClause = "(" + strings.Join(where, " or ") + ")"
-		if !opts.withTerminated {
-			permissionWhereClause += " and termination_reason is null"
-		}
-	} else {
-		if !opts.withTerminated {
-			permissionWhereClause = " where termination_reason is null"
-		}
+	permissionWhereClause := "(" + strings.Join(where, " or ") + ")"
+	if !opts.withTerminated {
+		permissionWhereClause += " and termination_reason is null"
 	}
 
 	limit := r.defaultLimit

internal/session/repository_test.go

@@ -132,7 +132,7 @@ func TestRepository_convertToSessions(t *testing.T) {
 	sess, err = repo.CreateSession(ctx, sessionWrapper, sess, []string{"0.0.0.0"})
 	require.NoError(t, err)
 
-	query := fmt.Sprintf(sessionList, "", "", "", "")
+	query := fmt.Sprintf(listSessionsTemplate, "termination_reason is null", 1000)
 	rows, err := rw.Query(ctx, query, nil)
 	require.NoError(t, err)
 	t.Cleanup(func() {