Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add total count to search responses #1232

Merged
merged 2 commits into from
Nov 1, 2023
Merged

feat: add total count to search responses #1232

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
d08fa21
feat: add total count to search responses
eemmiillyy Oct 31, 2023
ffcf6cb
refactor: nest search.byTable
eemmiillyy Oct 31, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/tidy-readers-dance.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@xata.io/client': minor
---

Add support for totalCount on search responses
22 changes: 14 additions & 8 deletions packages/client/src/schema/repository.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ import {
TransactionOperation
} from '../api/schemas';
import { XataPluginOptions } from '../plugins';
import { SearchXataRecord } from '../search';
import { SearchXataRecord, TotalCount } from '../search';
import { Boosters } from '../search/boosters';
import { TargetColumn } from '../search/target';
import { chunk, compact, isDefined, isNumber, isObject, isString, promiseMap } from '../util/lang';
Expand Down Expand Up @@ -748,7 +748,7 @@ export abstract class Repository<Record extends XataRecord> extends Query<
page?: SearchPageConfig;
target?: TargetColumn<Record>[];
}
): Promise<SearchXataRecord<SelectedPick<Record, ['*']>>[]>;
): Promise<{ records: SearchXataRecord<SelectedPick<Record, ['*']>>[] } & TotalCount>;

/**
* Search for vectors in the table.
Expand Down Expand Up @@ -777,7 +777,7 @@ export abstract class Repository<Record extends XataRecord> extends Query<
size?: number;
filter?: Filter<Record>;
}
): Promise<SearchXataRecord<SelectedPick<Record, ['*']>>[]>;
): Promise<{ records: SearchXataRecord<SelectedPick<Record, ['*']>>[] } & TotalCount>;

/**
* Aggregates records in the table.
Expand Down Expand Up @@ -1761,7 +1761,7 @@ export class RestRepository<Record extends XataRecord>
} = {}
) {
return this.#trace('search', async () => {
const { records } = await searchTable({
const { records, totalCount } = await searchTable({
pathParams: {
workspace: '{workspaceId}',
dbBranchName: '{dbBranch}',
Expand All @@ -1784,7 +1784,10 @@ export class RestRepository<Record extends XataRecord>
const schemaTables = await this.#getSchemaTables();

// TODO - Column selection not supported by search endpoint yet
return records.map((item) => initObject(this.#db, schemaTables, this.#table, item, ['*'])) as any;
return {
records: records.map((item) => initObject(this.#db, schemaTables, this.#table, item, ['*'])) as any,
totalCount
};
});
}

Expand All @@ -1798,9 +1801,9 @@ export class RestRepository<Record extends XataRecord>
filter?: Filter<Record> | undefined;
}
| undefined
): Promise<SearchXataRecord<SelectedPick<Record, ['*']>>[]> {
): Promise<{ records: SearchXataRecord<SelectedPick<Record, ['*']>>[] } & TotalCount> {
return this.#trace('vectorSearch', async () => {
const { records } = await vectorSearchTable({
const { records, totalCount } = await vectorSearchTable({
pathParams: {
workspace: '{workspaceId}',
dbBranchName: '{dbBranch}',
Expand All @@ -1820,7 +1823,10 @@ export class RestRepository<Record extends XataRecord>
const schemaTables = await this.#getSchemaTables();

// TODO - Column selection not supported by search endpoint yet
return records.map((item) => initObject(this.#db, schemaTables, this.#table, item, ['*'])) as any;
return {
records: records.map((item) => initObject(this.#db, schemaTables, this.#table, item, ['*'])),
totalCount
} as any;
});
}

Expand Down
69 changes: 40 additions & 29 deletions packages/client/src/search/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import { getBranchDetails, searchBranch } from '../api';
import { Responses, getBranchDetails, searchBranch } from '../api';
import { FuzzinessExpression, HighlightExpression, PrefixExpression, SearchPageConfig, Table } from '../api/schemas';
import { XataPlugin, XataPluginOptions } from '../plugins';
import { SchemaPluginResult } from '../schema';
Expand Down Expand Up @@ -28,32 +28,40 @@ export type SearchOptions<Schemas extends Record<string, BaseData>, Tables exten
page?: SearchPageConfig;
};

export type TotalCount = Pick<Responses.SearchResponse, 'totalCount'>;

export type SearchPluginResult<Schemas extends Record<string, BaseData>> = {
all: <Tables extends StringKeys<Schemas>>(
query: string,
options?: SearchOptions<Schemas, Tables>
) => Promise<
Values<{
[Model in ExtractTables<
Schemas,
Tables,
GetArrayInnerType<NonNullable<NonNullable<typeof options>['tables']>>
>]: {
table: Model;
record: Awaited<SearchXataRecord<SelectedPick<Schemas[Model] & XataRecord, ['*']>>>;
};
}>[]
TotalCount & {
records: Values<{
[Model in ExtractTables<
Schemas,
Tables,
GetArrayInnerType<NonNullable<NonNullable<typeof options>['tables']>>
>]: {
table: Model;
record: Awaited<SearchXataRecord<SelectedPick<Schemas[Model] & XataRecord, ['*']>>>;
};
}>[];
}
>;
byTable: <Tables extends StringKeys<Schemas>>(
query: string,
options?: SearchOptions<Schemas, Tables>
) => Promise<{
[Model in ExtractTables<
Schemas,
Tables,
GetArrayInnerType<NonNullable<NonNullable<typeof options>['tables']>>
>]?: Awaited<SearchXataRecord<SelectedPick<Schemas[Model] & XataRecord, ['*']>>[]>;
}>;
) => Promise<
TotalCount & {
records: {
[Model in ExtractTables<
Schemas,
Tables,
GetArrayInnerType<NonNullable<NonNullable<typeof options>['tables']>>
>]?: Awaited<SearchXataRecord<SelectedPick<Schemas[Model] & XataRecord, ['*']>>[]>;
};
}
>;
};

export class SearchPlugin<Schemas extends Record<string, XataRecord>> extends XataPlugin {
Expand All @@ -67,24 +75,26 @@ export class SearchPlugin<Schemas extends Record<string, XataRecord>> extends Xa
build(pluginOptions: XataPluginOptions): SearchPluginResult<Schemas> {
return {
all: async <Tables extends StringKeys<Schemas>>(query: string, options: SearchOptions<Schemas, Tables> = {}) => {
const records = await this.#search(query, options, pluginOptions);
const { records, totalCount } = await this.#search(query, options, pluginOptions);
const schemaTables = await this.#getSchemaTables(pluginOptions);

return records.map((record) => {
const { table = 'orphan' } = record.xata;

// TODO: Search endpoint doesn't support column selection
return { table, record: initObject(this.db, schemaTables, table, record, ['*']) } as any;
});
return {
totalCount,
records: records.map((record) => {
const { table = 'orphan' } = record.xata;
// TODO: Search endpoint doesn't support column selection
return { table, record: initObject(this.db, schemaTables, table, record, ['*']) } as any;
})
};
},
byTable: async <Tables extends StringKeys<Schemas>>(
query: string,
options: SearchOptions<Schemas, Tables> = {}
) => {
const records = await this.#search(query, options, pluginOptions);
const { records: rawRecords, totalCount } = await this.#search(query, options, pluginOptions);
const schemaTables = await this.#getSchemaTables(pluginOptions);

return records.reduce((acc, record) => {
const records = rawRecords.reduce((acc, record) => {
const { table = 'orphan' } = record.xata;

const items = acc[table] ?? [];
Expand All @@ -93,6 +103,7 @@ export class SearchPlugin<Schemas extends Record<string, XataRecord>> extends Xa

return { ...acc, [table]: [...items, item] };
}, {} as any);
return { totalCount, records };
}
};
}
Expand All @@ -104,14 +115,14 @@ export class SearchPlugin<Schemas extends Record<string, XataRecord>> extends Xa
) {
const { tables, fuzziness, highlight, prefix, page } = options ?? {};

const { records } = await searchBranch({
const { records, totalCount } = await searchBranch({
pathParams: { workspace: '{workspaceId}', dbBranchName: '{dbBranch}', region: '{region}' },
// @ts-ignore https://github.com/xataio/client-ts/issues/313
body: { tables, query, fuzziness, prefix, highlight, page },
...pluginOptions
});

return records;
return { records, totalCount };
}

async #getSchemaTables(pluginOptions: XataPluginOptions): Promise<Table[]> {
Expand Down
111 changes: 74 additions & 37 deletions test/integration/search.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,32 +60,38 @@ afterEach(async (ctx) => {

describe('search', () => {
test('search in table', async () => {
const owners = await xata.db.users.search('Owner');
expect(owners.length).toBeGreaterThan(0);

expect(owners.length).toBe(2);
expect(owners[0].id).toBeDefined();
expect(owners[0].full_name?.includes('Owner')).toBeTruthy();
expect(owners[0].read).toBeDefined();
expect(owners[0].getMetadata().score).toBeDefined();
expect(owners[0].getMetadata().table).toBe('users');
const { records, totalCount } = await xata.db.users.search('Owner');
expect(totalCount).toBe(2);
expect(records.length).toBeGreaterThan(0);

expect(records.length).toBe(2);
expect(records[0].id).toBeDefined();
expect(records[0].full_name?.includes('Owner')).toBeTruthy();
expect(records[0].read).toBeDefined();
expect(records[0].getMetadata().score).toBeDefined();
expect(records[0].getMetadata().table).toBe('users');
});

test('search in table with filtering', async () => {
const owners = await xata.db.users.search('Owner', {
const { records, totalCount } = await xata.db.users.search('Owner', {
filter: { full_name: 'Owner of team animals' }
});

expect(owners.length).toBe(1);
expect(owners[0].id).toBeDefined();
expect(owners[0].full_name?.includes('Owner of team animals')).toBeTruthy();
expect(owners[0].read).toBeDefined();
expect(owners[0].getMetadata().score).toBeDefined();
expect(totalCount).toBe(1);
expect(records.length).toBe(1);
expect(records[0].id).toBeDefined();
expect(records[0].full_name?.includes('Owner of team animals')).toBeTruthy();
expect(records[0].read).toBeDefined();
expect(records[0].getMetadata().score).toBeDefined();
});

test('search by tables with multiple tables', async () => {
const { users = [], teams = [] } = await xata.search.byTable('fruits', { tables: ['teams', 'users'] });
const {
records: { users = [], teams = [] },
totalCount
} = await xata.search.byTable('fruits', { tables: ['teams', 'users'] });

expect(totalCount).toBeGreaterThan(0);
expect(users.length).toBeGreaterThan(0);
expect(teams.length).toBeGreaterThan(0);

Expand All @@ -101,8 +107,12 @@ describe('search', () => {
});

test('search by table with all tables', async () => {
const { users = [], teams = [] } = await xata.search.byTable('fruits');
const {
records: { users = [], teams = [] },
totalCount
} = await xata.search.byTable('fruits');

expect(totalCount).toBeGreaterThan(0);
expect(users.length).toBeGreaterThan(0);
expect(teams.length).toBeGreaterThan(0);

Expand All @@ -118,9 +128,11 @@ describe('search', () => {
});

test('search all with multiple tables', async () => {
const results = await xata.search.all('fruits', { tables: ['teams', 'users'] });
const { records, totalCount } = await xata.search.all('fruits', { tables: ['teams', 'users'] });
expect(records).toBeDefined();

for (const result of results) {
expect(totalCount).toBeGreaterThan(0);
for (const result of records) {
if (result.table === 'teams') {
expect(result.record.id).toBeDefined();
expect(result.record.read).toBeDefined();
Expand All @@ -138,9 +150,11 @@ describe('search', () => {
});

test('search all with one table', async () => {
const results = await xata.search.all('fruits', { tables: ['teams'] });
const { records, totalCount } = await xata.search.all('fruits', { tables: ['teams'] });
expect(records).toBeDefined();

for (const result of results) {
expect(totalCount).toBeGreaterThan(0);
for (const result of records) {
expect(result.record.id).toBeDefined();
expect(result.record.read).toBeDefined();
expect(result.record.name?.includes('fruits')).toBeTruthy();
Expand All @@ -152,9 +166,11 @@ describe('search', () => {
});

test('search all with all tables', async () => {
const results = await xata.search.all('fruits');
const { records, totalCount } = await xata.search.all('fruits');
expect(records).toBeDefined();

for (const result of results) {
expect(totalCount).toBeGreaterThan(0);
for (const result of records) {
if (result.table === 'teams') {
expect(result.record.id).toBeDefined();
expect(result.record.read).toBeDefined();
Expand All @@ -175,26 +191,33 @@ describe('search', () => {
});

test('search all with filters', async () => {
const results = await xata.search.all('fruits', {
const { records, totalCount } = await xata.search.all('fruits', {
tables: [{ table: 'teams', filter: { name: 'Team fruits' } }]
});
expect(records).toBeDefined();

expect(results.length).toBe(1);
expect(results[0].table).toBe('teams');
expect(totalCount).toBe(1);
expect(records.length).toBe(1);
expect(records[0].table).toBe('teams');

if (results[0].table === 'teams') {
expect(results[0].record.id).toBeDefined();
expect(results[0].record.read).toBeDefined();
expect(results[0].record.name?.includes('fruits')).toBeTruthy();
expect(results[0].record.getMetadata().score).toBeDefined();
if (records[0].table === 'teams') {
expect(records[0].record.id).toBeDefined();
expect(records[0].record.read).toBeDefined();
expect(records[0].record.name?.includes('fruits')).toBeTruthy();
expect(records[0].record.getMetadata().score).toBeDefined();
}
});

test('search with page and offset', async () => {
const owners = await xata.db.users.search('Owner');
const page1 = await xata.db.users.search('Owner', { page: { size: 1 } });
const page2 = await xata.db.users.search('Owner', { page: { size: 1, offset: 1 } });
const { records: owners, totalCount } = await xata.db.users.search('Owner');
const { records: page1, totalCount: page1Count } = await xata.db.users.search('Owner', { page: { size: 1 } });
const { records: page2, totalCount: page2Count } = await xata.db.users.search('Owner', {
page: { size: 1, offset: 1 }
});

expect(totalCount).toBe(2);
expect(page1Count).toBe(2);
expect(page2Count).toBe(2);
expect(page1.length).toBe(1);
expect(page2.length).toBe(1);

Expand All @@ -205,10 +228,24 @@ describe('search', () => {
});

test('global search with page and offset', async () => {
const { users: owners = [] } = await xata.search.byTable('Owner');
const { users: page1 = [] } = await xata.search.byTable('Owner', { page: { size: 1 } });
const { users: page2 = [] } = await xata.search.byTable('Owner', { page: { size: 1, offset: 1 } });
const {
records: { users: owners = [] },
totalCount
} = await xata.search.byTable('Owner');
const {
records: { users: page1 = [] },
totalCount: page1Count
} = await xata.search.byTable('Owner', { page: { size: 1 } });
const {
records: { users: page2 = [] },
totalCount: page2Count
} = await xata.search.byTable('Owner', {
page: { size: 1, offset: 1 }
});

expect(totalCount).toBe(2);
expect(page1Count).toBe(2);
expect(page2Count).toBe(2);
expect(page1.length).toBe(1);
expect(page2.length).toBe(1);

Expand Down
Loading
Loading