Skip to content

Commit 5f60e13

Browse files
viambotnjooma
viambot
and
njooma
authored
AI update based on proto changes (#641)
Co-authored-by: Naveed Jooma <[email protected]>
1 parent ad3f7cd commit 5f60e13

File tree

3 files changed

+308
-7
lines changed

3 files changed

+308
-7
lines changed

‎src/app/data-client.spec.ts‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,14 @@ import {
2323
CaptureInterval,
2424
ConfigureDatabaseUserRequest,
2525
ConfigureDatabaseUserResponse,
26+
CreateIndexRequest,
27+
CreateIndexResponse,
2628
DataRequest,
2729
DeleteBinaryDataByFilterRequest,
2830
DeleteBinaryDataByFilterResponse,
2931
DeleteBinaryDataByIDsResponse,
32+
DeleteIndexRequest,
33+
DeleteIndexResponse,
3034
DeleteTabularDataResponse,
3135
ExportTabularDataRequest,
3236
ExportTabularDataResponse,
@@ -35,6 +39,11 @@ import {
3539
GetDatabaseConnectionResponse,
3640
GetLatestTabularDataRequest,
3741
GetLatestTabularDataResponse,
42+
Index,
43+
IndexableCollection,
44+
IndexCreator,
45+
ListIndexesRequest,
46+
ListIndexesResponse,
3847
RemoveBinaryDataFromDatasetByIDsRequest,
3948
RemoveBinaryDataFromDatasetByIDsResponse,
4049
RemoveBoundingBoxFromImageByIDRequest,
@@ -46,6 +55,7 @@ import {
4655
TabularData,
4756
TabularDataByFilterRequest,
4857
TabularDataByFilterResponse,
58+
TabularDataByMQLRequest,
4959
TabularDataByMQLResponse,
5060
TabularDataBySQLResponse,
5161
TabularDataSourceType,
@@ -325,6 +335,36 @@ describe('DataClient tests', () => {
325335
expect(result[0]?.key1).toBeInstanceOf(Date);
326336
expect(promise).toEqual(data);
327337
});
338+
339+
it('get tabular data from MQL with queryPrefixName', async () => {
340+
const expectedRequest = new TabularDataByMQLRequest({
341+
organizationId: 'some_org_id',
342+
mqlBinary: [BSON.serialize({ query: 'some_mql_query' })],
343+
queryPrefixName: 'my_prefix',
344+
});
345+
let capReq: TabularDataByMQLRequest | undefined = undefined;
346+
mockTransport = createRouterTransport(({ service }) => {
347+
service(DataService, {
348+
tabularDataByMQL: (req) => {
349+
capReq = req;
350+
return new TabularDataByMQLResponse({
351+
rawData: data.map((x) => BSON.serialize(x)),
352+
});
353+
},
354+
});
355+
});
356+
const promise = await subject().tabularDataByMQL(
357+
'some_org_id',
358+
[{ query: 'some_mql_query' }],
359+
false,
360+
undefined,
361+
'my_prefix'
362+
);
363+
expect(capReq).toStrictEqual(expectedRequest);
364+
const result = promise as typeof data;
365+
expect(result[0]?.key1).toBeInstanceOf(Date);
366+
expect(promise).toEqual(data);
367+
});
328368
});
329369

330370
describe('tabularDataByFilter tests', () => {
@@ -1040,6 +1080,155 @@ describe('DataClient tests', () => {
10401080
});
10411081
});
10421082

1083+
describe('createIndex tests', () => {
1084+
let capReq: CreateIndexRequest;
1085+
beforeEach(() => {
1086+
mockTransport = createRouterTransport(({ service }) => {
1087+
service(DataService, {
1088+
createIndex: (req) => {
1089+
capReq = req;
1090+
return new CreateIndexResponse();
1091+
},
1092+
});
1093+
});
1094+
});
1095+
it('creates an index', async () => {
1096+
const organizationId = 'orgId';
1097+
const collectionType = IndexableCollection.PIPELINE_SINK;
1098+
const indexSpec = { keys: { field: 1 }, options: { priority: 1 } };
1099+
const pipelineName = 'pipeline1';
1100+
await subject().createIndex(
1101+
organizationId,
1102+
collectionType,
1103+
indexSpec,
1104+
pipelineName
1105+
);
1106+
expect(capReq.organizationId).toBe(organizationId);
1107+
expect(capReq.collectionType).toBe(collectionType);
1108+
expect(
1109+
capReq.indexSpec.map((spec) => BSON.deserialize(spec))[0]
1110+
).toStrictEqual(indexSpec);
1111+
expect(capReq.pipelineName).toBe(pipelineName);
1112+
});
1113+
it('creates an index without pipeline name', async () => {
1114+
const organizationId = 'orgId';
1115+
const collectionType = IndexableCollection.HOT_STORE;
1116+
const indexSpec = { keys: { field: 2 }, options: { priority: 2 } };
1117+
await subject().createIndex(organizationId, collectionType, indexSpec);
1118+
expect(capReq.organizationId).toBe(organizationId);
1119+
expect(capReq.collectionType).toBe(collectionType);
1120+
expect(
1121+
capReq.indexSpec.map((spec) => BSON.deserialize(spec))[0]
1122+
).toStrictEqual(indexSpec);
1123+
});
1124+
});
1125+
describe('listIndexes tests', () => {
1126+
let capReq: ListIndexesRequest;
1127+
const index1 = new Index({
1128+
collectionType: IndexableCollection.HOT_STORE,
1129+
indexName: 'index1',
1130+
indexSpec: [new TextEncoder().encode(JSON.stringify({ field: 1 }))],
1131+
createdBy: IndexCreator.CUSTOMER,
1132+
});
1133+
const index2 = new Index({
1134+
collectionType: IndexableCollection.PIPELINE_SINK,
1135+
pipelineName: 'pipeline1',
1136+
indexName: 'index2',
1137+
indexSpec: [
1138+
new TextEncoder().encode(JSON.stringify({ another_field: -1 })),
1139+
],
1140+
createdBy: IndexCreator.VIAM,
1141+
});
1142+
const indexes = [index1, index2];
1143+
beforeEach(() => {
1144+
mockTransport = createRouterTransport(({ service }) => {
1145+
service(DataService, {
1146+
listIndexes: (req) => {
1147+
capReq = req;
1148+
return new ListIndexesResponse({
1149+
indexes,
1150+
});
1151+
},
1152+
});
1153+
});
1154+
});
1155+
it('lists indexes', async () => {
1156+
const organizationId = 'orgId';
1157+
const collectionType = IndexableCollection.HOT_STORE;
1158+
const pipelineName = 'pipeline1';
1159+
const expectedRequest = new ListIndexesRequest({
1160+
organizationId,
1161+
collectionType,
1162+
pipelineName,
1163+
});
1164+
const result = await subject().listIndexes(
1165+
organizationId,
1166+
collectionType,
1167+
pipelineName
1168+
);
1169+
expect(capReq).toStrictEqual(expectedRequest);
1170+
expect(result).toEqual(indexes);
1171+
});
1172+
it('lists indexes without pipeline name', async () => {
1173+
const organizationId = 'orgId';
1174+
const collectionType = IndexableCollection.HOT_STORE;
1175+
const expectedRequest = new ListIndexesRequest({
1176+
organizationId,
1177+
collectionType,
1178+
});
1179+
const result = await subject().listIndexes(
1180+
organizationId,
1181+
collectionType
1182+
);
1183+
expect(capReq).toStrictEqual(expectedRequest);
1184+
expect(result).toEqual(indexes);
1185+
});
1186+
});
1187+
describe('deleteIndex tests', () => {
1188+
let capReq: DeleteIndexRequest;
1189+
beforeEach(() => {
1190+
mockTransport = createRouterTransport(({ service }) => {
1191+
service(DataService, {
1192+
deleteIndex: (req) => {
1193+
capReq = req;
1194+
return new DeleteIndexResponse();
1195+
},
1196+
});
1197+
});
1198+
});
1199+
it('deletes an index', async () => {
1200+
const organizationId = 'orgId';
1201+
const collectionType = IndexableCollection.HOT_STORE;
1202+
const indexName = 'my_index';
1203+
const pipelineName = 'pipeline1';
1204+
const expectedRequest = new DeleteIndexRequest({
1205+
organizationId,
1206+
collectionType,
1207+
indexName,
1208+
pipelineName,
1209+
});
1210+
await subject().deleteIndex(
1211+
organizationId,
1212+
collectionType,
1213+
indexName,
1214+
pipelineName
1215+
);
1216+
expect(capReq).toStrictEqual(expectedRequest);
1217+
});
1218+
it('deletes an index without pipeline name', async () => {
1219+
const organizationId = 'orgId';
1220+
const collectionType = IndexableCollection.HOT_STORE;
1221+
const indexName = 'my_index';
1222+
const expectedRequest = new DeleteIndexRequest({
1223+
organizationId,
1224+
collectionType,
1225+
indexName,
1226+
});
1227+
await subject().deleteIndex(organizationId, collectionType, indexName);
1228+
expect(capReq).toStrictEqual(expectedRequest);
1229+
});
1230+
});
1231+
10431232
describe('createFilter tests', () => {
10441233
it('create empty filter', () => {
10451234
const testFilter = DataClient.createFilter({});

‎src/app/data-client.ts‎

Lines changed: 114 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@ import {
1212
CaptureInterval,
1313
CaptureMetadata,
1414
Filter,
15+
Index,
16+
IndexableCollection,
1517
Order,
1618
TabularDataSource,
1719
TabularDataSourceType,
@@ -117,7 +119,7 @@ export class DataClient {
117119
private datasetClient: Client<typeof DatasetService>;
118120
private dataSyncClient: Client<typeof DataSyncService>;
119121
private dataPipelinesClient: Client<typeof DataPipelinesService>;
120-
static readonly UPLOAD_CHUNK_SIZE = 8;
122+
static readonly UPLOAD_CHUNK_SIZE = 1024 * 64;
121123

122124
constructor(transport: Transport) {
123125
this.dataClient = createClient(DataService, transport);
@@ -273,13 +275,15 @@ export class DataClient {
273275
* store. Defaults to false. Deprecated - use dataSource instead.
274276
* @param dataSource The data source to query. Defaults to the standard data
275277
* source.
278+
* @param queryPrefixName Optional name of the query prefix.
276279
* @returns An array of data objects
277280
*/
278281
async tabularDataByMQL(
279282
organizationId: string,
280283
query: Uint8Array[] | Record<string, Date | JsonValue>[],
281284
useRecentData?: boolean,
282-
tabularDataSource?: TabularDataSource
285+
tabularDataSource?: TabularDataSource,
286+
queryPrefixName?: string
283287
) {
284288
const binary: Uint8Array[] =
285289
query[0] instanceof Uint8Array
@@ -301,6 +305,7 @@ export class DataClient {
301305
organizationId,
302306
mqlBinary: binary,
303307
dataSource,
308+
queryPrefixName,
304309
});
305310
return resp.rawData.map((value) => BSON.deserialize(value));
306311
}
@@ -1544,6 +1549,7 @@ export class DataClient {
15441549
* @param query The MQL query to run as a list of BSON documents
15451550
* @param schedule The schedule to run the query on (cron expression)
15461551
* @param dataSourceType The type of data source to use for the data pipeline
1552+
* @param enableBackfill Whether to enable backfill for the data pipeline
15471553
* @returns The ID of the created data pipeline
15481554
*/
15491555
async createDataPipeline(
@@ -1630,6 +1636,107 @@ export class DataClient {
16301636
resp.nextPageToken
16311637
);
16321638
}
1639+
1640+
/**
1641+
* CreateIndex starts a custom index build
1642+
*
1643+
* @example
1644+
*
1645+
* ```ts
1646+
* await dataClient.createIndex(
1647+
* '123abc45-1234-5678-90ab-cdef12345678',
1648+
* IndexableCollection.HOT_STORE,
1649+
* [new TextEncoder().encode(JSON.stringify({ field: 1 }))]
1650+
* );
1651+
* ```
1652+
*
1653+
* @param organizationId The ID of the organization
1654+
* @param collectionType The type of collection to create the index on
1655+
* @param indexSpec The MongoDB index specification in JSON format, as a
1656+
* Uint8Array
1657+
* @param pipelineName Optional name of the pipeline if collectionType is
1658+
* PIPELINE_SINK
1659+
*/
1660+
async createIndex(
1661+
organizationId: string,
1662+
collectionType: IndexableCollection,
1663+
indexSpec: {
1664+
keys: Record<string, number>;
1665+
options?: Record<string, unknown>;
1666+
},
1667+
pipelineName?: string
1668+
) {
1669+
await this.dataClient.createIndex({
1670+
organizationId,
1671+
collectionType,
1672+
indexSpec: [BSON.serialize(indexSpec)],
1673+
pipelineName,
1674+
});
1675+
}
1676+
1677+
/**
1678+
* ListIndexes returns all the indexes for a given collection
1679+
*
1680+
* @example
1681+
*
1682+
* ```ts
1683+
* const indexes = await dataClient.listIndexes(
1684+
* '123abc45-1234-5678-90ab-cdef12345678',
1685+
* IndexableCollection.HOT_STORE
1686+
* );
1687+
* ```
1688+
*
1689+
* @param organizationId The ID of the organization
1690+
* @param collectionType The type of collection to list indexes for
1691+
* @param pipelineName Optional name of the pipeline if collectionType is
1692+
* PIPELINE_SINK
1693+
* @returns An array of indexes
1694+
*/
1695+
async listIndexes(
1696+
organizationId: string,
1697+
collectionType: IndexableCollection,
1698+
pipelineName?: string
1699+
): Promise<Index[]> {
1700+
const resp = await this.dataClient.listIndexes({
1701+
organizationId,
1702+
collectionType,
1703+
pipelineName,
1704+
});
1705+
return resp.indexes;
1706+
}
1707+
1708+
/**
1709+
* DeleteIndex drops the specified custom index from a collection
1710+
*
1711+
* @example
1712+
*
1713+
* ```ts
1714+
* await dataClient.deleteIndex(
1715+
* '123abc45-1234-5678-90ab-cdef12345678',
1716+
* IndexableCollection.HOT_STORE,
1717+
* 'my_index'
1718+
* );
1719+
* ```
1720+
*
1721+
* @param organizationId The ID of the organization
1722+
* @param collectionType The type of collection to delete the index from
1723+
* @param indexName The name of the index to delete
1724+
* @param pipelineName Optional name of the pipeline if collectionType is
1725+
* PIPELINE_SINK
1726+
*/
1727+
async deleteIndex(
1728+
organizationId: string,
1729+
collectionType: IndexableCollection,
1730+
indexName: string,
1731+
pipelineName?: string
1732+
) {
1733+
await this.dataClient.deleteIndex({
1734+
organizationId,
1735+
collectionType,
1736+
indexName,
1737+
pipelineName,
1738+
});
1739+
}
16331740
}
16341741

16351742
export class ListDataPipelineRunsPage {
@@ -1682,5 +1789,9 @@ export class ListDataPipelineRunsPage {
16821789
}
16831790
}
16841791

1685-
export { type BinaryID, type Order } from '../gen/app/data/v1/data_pb';
1792+
export {
1793+
type BinaryID,
1794+
type IndexableCollection,
1795+
type Order,
1796+
} from '../gen/app/data/v1/data_pb';
16861797
export { type UploadMetadata } from '../gen/app/datasync/v1/data_sync_pb';

0 commit comments

Comments
 (0)