Skip to content

Commit

Permalink
[data.search] Move search method inside session service and add tests (…
Browse files Browse the repository at this point in the history 
…#84715)
  • Loading branch information
@lukasolson
lukasolson authored Dec 3, 2020
1 parent 9f09c84 commit 726424a
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 39 deletions.
35 changes: 11 additions & 24 deletions src/plugins/data/server/search/search_service.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
* under the License.
*/

import { BehaviorSubject, from, Observable } from 'rxjs';
import { BehaviorSubject, Observable } from 'rxjs';
import { pick } from 'lodash';
import {
CoreSetup,
Expand All @@ -29,7 +29,7 @@ import {
SharedGlobalConfig,
StartServicesAccessor,
} from 'src/core/server';
import { catchError, first, map, switchMap } from 'rxjs/operators';
import { catchError, first, map } from 'rxjs/operators';
import { BfetchServerSetup } from 'src/plugins/bfetch/server';
import { ExpressionsServerSetup } from 'src/plugins/expressions/server';
import {
Expand All @@ -50,7 +50,11 @@ import { DataPluginStart } from '../plugin';
import { UsageCollectionSetup } from '../../../usage_collection/server';
import { registerUsageCollector } from './collectors/register';
import { usageProvider } from './collectors/usage';
import { BACKGROUND_SESSION_TYPE, searchTelemetry } from '../saved_objects';
import {
BACKGROUND_SESSION_TYPE,
backgroundSessionMapping,
searchTelemetry,
} from '../saved_objects';
import {
IEsSearchRequest,
IEsSearchResponse,
Expand All @@ -74,8 +78,6 @@ import { aggShardDelay } from '../../common/search/aggs/buckets/shard_delay_fn';
import { ConfigSchema } from '../../config';
import { BackgroundSessionService, ISearchSessionClient } from './session';
import { registerSessionRoutes } from './routes/session';
import { backgroundSessionMapping } from '../saved_objects';
import { tapFirst } from '../../common/utils';

declare module 'src/core/server' {
interface RequestHandlerContext {
Expand Down Expand Up @@ -297,32 +299,17 @@ export class SearchService implements Plugin<ISearchSetup, ISearchStart> {
SearchStrategyRequest extends IKibanaSearchRequest = IEsSearchRequest,
SearchStrategyResponse extends IKibanaSearchResponse = IEsSearchResponse
>(
searchRequest: SearchStrategyRequest,
request: SearchStrategyRequest,
options: ISearchOptions,
deps: SearchStrategyDependencies
) => {
const strategy = this.getSearchStrategy<SearchStrategyRequest, SearchStrategyResponse>(
options.strategy
);

// If this is a restored background search session, look up the ID using the provided sessionId
const getSearchRequest = async () =>
!options.isRestore || searchRequest.id
? searchRequest
: {
...searchRequest,
id: await this.sessionService.getId(searchRequest, options, deps),
};

return from(getSearchRequest()).pipe(
switchMap((request) => strategy.search(request, options, deps)),
tapFirst((response) => {
if (searchRequest.id || !options.sessionId || !response.id || options.isRestore) return;
this.sessionService.trackId(searchRequest, response.id, options, {
savedObjectsClient: deps.savedObjectsClient,
});
})
);
return options.sessionId
? this.sessionService.search(strategy, request, options, deps)
: strategy.search(request, options, deps);
};

private cancel = (id: string, options: ISearchOptions, deps: SearchStrategyDependencies) => {
Expand Down
89 changes: 79 additions & 10 deletions src/plugins/data/server/search/session/session_service.test.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@
* under the License.
*/

import { of } from 'rxjs';
import type { SavedObject, SavedObjectsClientContract } from 'kibana/server';
import type { SearchStrategyDependencies } from '../types';
import { savedObjectsClientMock } from '../../../../../core/server/mocks';
import { BackgroundSessionStatus } from '../../../common';
import { BACKGROUND_SESSION_TYPE } from '../../saved_objects';
Expand All @@ -28,6 +30,7 @@ describe('BackgroundSessionService', () => {
let savedObjectsClient: jest.Mocked<SavedObjectsClientContract>;
let service: BackgroundSessionService;

const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
const mockSavedObject: SavedObject = {
id: 'd7170a35-7e2c-48d6-8dec-9a056721b489',
type: BACKGROUND_SESSION_TYPE,
Expand All @@ -45,9 +48,13 @@ describe('BackgroundSessionService', () => {
service = new BackgroundSessionService();
});

it('save throws if `name` is not provided', () => {
const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
it('search throws if `name` is not provided', () => {
expect(() => service.save(sessionId, {}, { savedObjectsClient })).rejects.toMatchInlineSnapshot(
`[Error: Name is required]`
);
});

it('save throws if `name` is not provided', () => {
expect(() => service.save(sessionId, {}, { savedObjectsClient })).rejects.toMatchInlineSnapshot(
`[Error: Name is required]`
);
Expand All @@ -56,7 +63,6 @@ describe('BackgroundSessionService', () => {
it('get calls saved objects client', async () => {
savedObjectsClient.get.mockResolvedValue(mockSavedObject);

const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
const response = await service.get(sessionId, { savedObjectsClient });

expect(response).toBe(mockSavedObject);
Expand Down Expand Up @@ -93,7 +99,6 @@ describe('BackgroundSessionService', () => {
};
savedObjectsClient.update.mockResolvedValue(mockUpdateSavedObject);

const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
const attributes = { name: 'new_name' };
const response = await service.update(sessionId, attributes, { savedObjectsClient });

Expand All @@ -108,19 +113,87 @@ describe('BackgroundSessionService', () => {
it('delete calls saved objects client', async () => {
savedObjectsClient.delete.mockResolvedValue({});

const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
const response = await service.delete(sessionId, { savedObjectsClient });

expect(response).toEqual({});
expect(savedObjectsClient.delete).toHaveBeenCalledWith(BACKGROUND_SESSION_TYPE, sessionId);
});

describe('search', () => {
const mockSearch = jest.fn().mockReturnValue(of({}));
const mockStrategy = { search: mockSearch };
const mockDeps = {} as SearchStrategyDependencies;

beforeEach(() => {
mockSearch.mockClear();
});

it('searches using the original request if not restoring', async () => {
const searchRequest = { params: {} };
const options = { sessionId, isStored: false, isRestore: false };

await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();

expect(mockSearch).toBeCalledWith(searchRequest, options, mockDeps);
});

it('searches using the original request if `id` is provided', async () => {
const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
const searchRequest = { id: searchId, params: {} };
const options = { sessionId, isStored: true, isRestore: true };

await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();

expect(mockSearch).toBeCalledWith(searchRequest, options, mockDeps);
});

it('searches by looking up an `id` if restoring and `id` is not provided', async () => {
const searchRequest = { params: {} };
const options = { sessionId, isStored: true, isRestore: true };
const spyGetId = jest.spyOn(service, 'getId').mockResolvedValueOnce('my_id');

await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();

expect(mockSearch).toBeCalledWith({ ...searchRequest, id: 'my_id' }, options, mockDeps);

spyGetId.mockRestore();
});

it('calls `trackId` once if the response contains an `id` and not restoring', async () => {
const searchRequest = { params: {} };
const options = { sessionId, isStored: false, isRestore: false };
const spyTrackId = jest.spyOn(service, 'trackId').mockResolvedValue();
mockSearch.mockReturnValueOnce(of({ id: 'my_id' }, { id: 'my_id' }));

await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();

expect(spyTrackId).toBeCalledTimes(1);
expect(spyTrackId).toBeCalledWith(searchRequest, 'my_id', options, mockDeps);

spyTrackId.mockRestore();
});

it('does not call `trackId` if restoring', async () => {
const searchRequest = { params: {} };
const options = { sessionId, isStored: true, isRestore: true };
const spyGetId = jest.spyOn(service, 'getId').mockResolvedValueOnce('my_id');
const spyTrackId = jest.spyOn(service, 'trackId').mockResolvedValue();
mockSearch.mockReturnValueOnce(of({ id: 'my_id' }));

await service.search(mockStrategy, searchRequest, options, mockDeps).toPromise();

expect(spyTrackId).not.toBeCalled();

spyGetId.mockRestore();
spyTrackId.mockRestore();
});
});

describe('trackId', () => {
it('stores hash in memory when `isStored` is `false` for when `save` is called', async () => {
const searchRequest = { params: {} };
const requestHash = createRequestHash(searchRequest.params);
const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
const isStored = false;
const name = 'my saved background search session';
const appId = 'my_app_id';
Expand Down Expand Up @@ -164,7 +237,6 @@ describe('BackgroundSessionService', () => {
const searchRequest = { params: {} };
const requestHash = createRequestHash(searchRequest.params);
const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
const isStored = true;

await service.trackId(
Expand All @@ -191,7 +263,6 @@ describe('BackgroundSessionService', () => {

it('throws if there is not a saved object', () => {
const searchRequest = { params: {} };
const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';

expect(() =>
service.getId(searchRequest, { sessionId, isStored: false }, { savedObjectsClient })
Expand All @@ -202,7 +273,6 @@ describe('BackgroundSessionService', () => {

it('throws if not restoring a saved session', () => {
const searchRequest = { params: {} };
const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';

expect(() =>
service.getId(
Expand All @@ -219,7 +289,6 @@ describe('BackgroundSessionService', () => {
const searchRequest = { params: {} };
const requestHash = createRequestHash(searchRequest.params);
const searchId = 'FnpFYlBpeXdCUTMyZXhCLTc1TWFKX0EbdDFDTzJzTE1Sck9PVTBIcW1iU05CZzo4MDA0';
const sessionId = 'd7170a35-7e2c-48d6-8dec-9a056721b489';
const mockSession = {
id: 'd7170a35-7e2c-48d6-8dec-9a056721b489',
type: BACKGROUND_SESSION_TYPE,
Expand Down
37 changes: 32 additions & 5 deletions src/plugins/data/server/search/session/session_service.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,19 @@
*/

import { CoreStart, KibanaRequest, SavedObjectsClientContract } from 'kibana/server';
import { from } from 'rxjs';
import { switchMap } from 'rxjs/operators';
import {
BackgroundSessionSavedObjectAttributes,
BackgroundSessionStatus,
IKibanaSearchRequest,
IKibanaSearchResponse,
ISearchOptions,
SearchSessionFindOptions,
BackgroundSessionStatus,
tapFirst,
} from '../../../common';
import { BACKGROUND_SESSION_TYPE } from '../../saved_objects';
import { ISearchStrategy, SearchStrategyDependencies } from '../types';
import { createRequestHash } from './utils';

const DEFAULT_EXPIRATION = 7 * 24 * 60 * 60 * 1000;
Expand Down Expand Up @@ -59,6 +64,32 @@ export class BackgroundSessionService {
this.sessionSearchMap.clear();
};

public search = <Request extends IKibanaSearchRequest, Response extends IKibanaSearchResponse>(
strategy: ISearchStrategy<Request, Response>,
searchRequest: Request,
options: ISearchOptions,
deps: SearchStrategyDependencies
) => {
// If this is a restored background search session, look up the ID using the provided sessionId
const getSearchRequest = async () =>
!options.isRestore || searchRequest.id
? searchRequest
: {
...searchRequest,
id: await this.getId(searchRequest, options, deps),
};

return from(getSearchRequest()).pipe(
switchMap((request) => strategy.search(request, options, deps)),
tapFirst((response) => {
if (searchRequest.id || !options.sessionId || !response.id || options.isRestore) return;
this.trackId(searchRequest, response.id, options, {
savedObjectsClient: deps.savedObjectsClient,
});
})
);
};

// TODO: Generate the `userId` from the realm type/realm name/username
public save = async (
sessionId: string,
Expand Down Expand Up @@ -208,10 +239,6 @@ export class BackgroundSessionService {
update: (sessionId: string, attributes: Partial<BackgroundSessionSavedObjectAttributes>) =>
this.update(sessionId, attributes, deps),
delete: (sessionId: string) => this.delete(sessionId, deps),
trackId: (searchRequest: IKibanaSearchRequest, searchId: string, options: ISearchOptions) =>
this.trackId(searchRequest, searchId, options, deps),
getId: (searchRequest: IKibanaSearchRequest, options: ISearchOptions) =>
this.getId(searchRequest, options, deps),
};
};
};
Expand Down

0 comments on commit 726424a

Please sign in to comment.