diff --git a/src/collections/collection/index.ts b/src/collections/collection/index.ts index ea89267f..fe69a1c6 100644 --- a/src/collections/collection/index.ts +++ b/src/collections/collection/index.ts @@ -125,7 +125,7 @@ const collection = ( name: name, query: queryCollection, sort: sort(), - tenants: tenants(connection, capitalizedName), + tenants: tenants(connection, capitalizedName, dbVersionSupport), exists: () => new ClassExists(connection).withClassName(capitalizedName).do(), iterator: (opts?: IteratorOptions) => new Iterator((limit: number, after?: string) => diff --git a/src/collections/tenants/index.ts b/src/collections/tenants/index.ts index 82205e93..1e3994da 100644 --- a/src/collections/tenants/index.ts +++ b/src/collections/tenants/index.ts @@ -1,33 +1,91 @@ -import Connection from '../../connection/index.js'; +import { ConnectionGRPC } from '../../connection/index.js'; +import { WeaviateUnsupportedFeatureError } from '../../errors.js'; +import { TenantActivityStatus, TenantsGetReply } from '../../proto/v1/tenants.js'; import { TenantsCreator, TenantsDeleter, TenantsGetter, TenantsUpdater } from '../../schema/index.js'; +import { DbVersionSupport } from '../../utils/dbVersion.js'; export type Tenant = { name: string; activityStatus?: 'COLD' | 'HOT'; }; -const tenants = (connection: Connection, name: string): Tenants => { - const parseTenants = (tenants: Tenant | Tenant[]) => (Array.isArray(tenants) ? tenants : [tenants]); +export type TenantsGetOptions = { + tenants?: string; +}; + +class ActivityStatusMapper { + static from(status: TenantActivityStatus): 'COLD' | 'HOT' { + switch (status) { + case TenantActivityStatus.TENANT_ACTIVITY_STATUS_COLD: + return 'COLD'; + case TenantActivityStatus.TENANT_ACTIVITY_STATUS_HOT: + return 'HOT'; + default: + throw new Error(`Unsupported tenant activity status: ${status}`); + } + } +} + +const mapReply = (reply: TenantsGetReply): Record => { + const tenants: Record = {}; + reply.tenants.forEach((t) => { + tenants[t.name] = { + name: t.name, + activityStatus: ActivityStatusMapper.from(t.activityStatus), + }; + }); + return tenants; +}; + +const checkSupportForGRPCTenantsGetEndpoint = async (dbVersionSupport: DbVersionSupport) => { + const check = await dbVersionSupport.supportsTenantsGetGRPCMethod(); + if (!check.supports) throw new WeaviateUnsupportedFeatureError(check.message); +}; + +const parseTenantOrTenantArray = (tenants: Tenant | Tenant[]) => + Array.isArray(tenants) ? tenants : [tenants]; + +const parseStringOrTenant = (tenant: string | Tenant) => (typeof tenant === 'string' ? tenant : tenant.name); + +const tenants = ( + connection: ConnectionGRPC, + collection: string, + dbVersionSupport: DbVersionSupport +): Tenants => { + const getGRPC = (names?: string[]) => + checkSupportForGRPCTenantsGetEndpoint(dbVersionSupport) + .then(() => connection.tenants(collection)) + .then((builder) => builder.withGet({ names })) + .then(mapReply); + const getREST = () => + new TenantsGetter(connection, collection).do().then((tenants) => { + const result: Record = {}; + tenants.forEach((tenant) => { + if (!tenant.name) return; + result[tenant.name] = tenant as Tenant; + }); + return result; + }); return { create: (tenants: Tenant | Tenant[]) => - new TenantsCreator(connection, name, parseTenants(tenants)).do() as Promise, - get: () => - new TenantsGetter(connection, name).do().then((tenants) => { - const result: Record = {}; - tenants.forEach((tenant) => { - if (!tenant.name) return; - result[tenant.name] = tenant as Tenant; - }); - return result; - }), + new TenantsCreator(connection, collection, parseTenantOrTenantArray(tenants)).do() as Promise, + get: async function () { + const check = await dbVersionSupport.supportsTenantsGetGRPCMethod(); + return check.supports ? getGRPC() : getREST(); + }, + getByNames: (tenants: (string | Tenant)[]) => getGRPC(tenants.map(parseStringOrTenant)), + getByName: (tenant: string | Tenant) => { + const tenantName = parseStringOrTenant(tenant); + return getGRPC([tenantName]).then((tenants) => tenants[tenantName] || null); + }, remove: (tenants: Tenant | Tenant[]) => new TenantsDeleter( connection, - name, - parseTenants(tenants).map((t) => t.name) + collection, + parseTenantOrTenantArray(tenants).map((t) => t.name) ).do(), update: (tenants: Tenant | Tenant[]) => - new TenantsUpdater(connection, name, parseTenants(tenants)).do() as Promise, + new TenantsUpdater(connection, collection, parseTenantOrTenantArray(tenants)).do() as Promise, }; }; @@ -58,6 +116,24 @@ export interface Tenants { * @returns {Promise>} A list of tenants as an object of Tenant types, where the key is the tenant name. */ get: () => Promise>; + /** + * Return the specified tenants from a collection in Weaviate. + * + * The collection must have been created with multi-tenancy enabled. + * + * @param {(string | Tenant)[]} names The tenants to retrieve. + * @returns {Promise} The list of tenants. If the tenant does not exist, it will not be included in the list. + */ + getByNames: (names: (string | Tenant)[]) => Promise>; + /** + * Return the specified tenant from a collection in Weaviate. + * + * The collection must have been created with multi-tenancy enabled. + * + * @param {string | Tenant} name The name of the tenant to retrieve. + * @returns {Promise} The tenant as a Tenant type, or null if the tenant does not exist. + */ + getByName: (name: string | Tenant) => Promise; /** * Remove the specified tenants from a collection in Weaviate. * diff --git a/src/collections/tenants/integration.test.ts b/src/collections/tenants/integration.test.ts index ce70c92a..4bd2d34c 100644 --- a/src/collections/tenants/integration.test.ts +++ b/src/collections/tenants/integration.test.ts @@ -1,8 +1,9 @@ /* eslint-disable @typescript-eslint/no-non-null-assertion */ +import { WeaviateUnsupportedFeatureError } from '../../errors.js'; import weaviate, { WeaviateClient } from '../../index.js'; import { Collection } from '../collection/index.js'; -describe('Testing of the collection.data methods', () => { +describe('Testing of the collection.tenants methods', () => { let client: WeaviateClient; let collection: Collection; const collectionName = 'TestCollectionTenants'; @@ -75,4 +76,83 @@ describe('Testing of the collection.data methods', () => { expect(result[0].name).toBe('cold'); expect(result[0].activityStatus).toBe('HOT'); }); + + describe('getByName and getByNames', () => { + it('should be able to get a tenant by name string', async () => { + const query = () => collection.tenants.getByName('hot'); + if (await client.getWeaviateVersion().then((ver) => ver.isLowerThan(1, 25, 0))) { + await expect(query()).rejects.toThrow(WeaviateUnsupportedFeatureError); + return; + } + const result = await query(); + expect(result).toHaveProperty('name', 'hot'); + expect(result).toHaveProperty('activityStatus', 'HOT'); + }); + + it('should be able to get a tenant by tenant object', async () => { + const query = () => collection.tenants.getByName({ name: 'hot' }); + if (await client.getWeaviateVersion().then((ver) => ver.isLowerThan(1, 25, 0))) { + await expect(query()).rejects.toThrow(WeaviateUnsupportedFeatureError); + return; + } + const result = await query(); + expect(result).toHaveProperty('name', 'hot'); + expect(result).toHaveProperty('activityStatus', 'HOT'); + }); + + it('should fail to get a non-existing tenant', async () => { + const query = () => collection.tenants.getByName('non-existing'); + if (await client.getWeaviateVersion().then((ver) => ver.isLowerThan(1, 25, 0))) { + await expect(query()).rejects.toThrow(WeaviateUnsupportedFeatureError); + return; + } + const result = await query(); + expect(result).toBeNull(); + }); + + it('should be able to get tenants by name strings', async () => { + const query = () => collection.tenants.getByNames(['hot', 'cold']); + if (await client.getWeaviateVersion().then((ver) => ver.isLowerThan(1, 25, 0))) { + await expect(query()).rejects.toThrow(WeaviateUnsupportedFeatureError); + return; + } + const result = await query(); + expect(result).toHaveProperty('hot'); + expect(result).toHaveProperty('cold'); + }); + + it('should be able to get tenants by tenant objects', async () => { + const query = () => collection.tenants.getByNames([{ name: 'hot' }, { name: 'cold' }]); + if (await client.getWeaviateVersion().then((ver) => ver.isLowerThan(1, 25, 0))) { + await expect(query()).rejects.toThrow(WeaviateUnsupportedFeatureError); + return; + } + const result = await query(); + expect(result).toHaveProperty('hot'); + expect(result).toHaveProperty('cold'); + }); + + it('should be able to get tenants by mixed name strings and tenant objects', async () => { + const query = () => collection.tenants.getByNames(['hot', { name: 'cold' }]); + if (await client.getWeaviateVersion().then((ver) => ver.isLowerThan(1, 25, 0))) { + await expect(query()).rejects.toThrow(WeaviateUnsupportedFeatureError); + return; + } + const result = await query(); + expect(result).toHaveProperty('hot'); + expect(result).toHaveProperty('cold'); + }); + + it('should be able to get partial tenants', async () => { + const query = () => collection.tenants.getByNames(['hot', 'non-existing']); + if (await client.getWeaviateVersion().then((ver) => ver.isLowerThan(1, 25, 0))) { + await expect(query()).rejects.toThrow(WeaviateUnsupportedFeatureError); + return; + } + const result = await query(); + expect(result).toHaveProperty('hot'); + expect(result).not.toHaveProperty('cold'); + expect(result).not.toHaveProperty('non-existing'); + }); + }); }); diff --git a/src/connection/grpc.ts b/src/connection/grpc.ts index 2268c34b..73c53cd1 100644 --- a/src/connection/grpc.ts +++ b/src/connection/grpc.ts @@ -12,6 +12,7 @@ import { HealthDefinition, HealthCheckResponse_ServingStatus } from '../proto/go import Batcher, { Batch } from '../grpc/batcher.js'; import Searcher, { Search } from '../grpc/searcher.js'; import { DbVersionSupport, initDbVersionProvider } from '../utils/dbVersion.js'; +import TenantsManager, { Tenants } from '../grpc/tenantsManager.js'; import { WeaviateGRPCUnavailableError } from '../errors.js'; @@ -63,20 +64,29 @@ export default class ConnectionGRPC extends ConnectionGQL { } } - search = (name: string, consistencyLevel?: ConsistencyLevel, tenant?: string) => { + search = (collection: string, consistencyLevel?: ConsistencyLevel, tenant?: string) => { if (this.authEnabled) { return this.login().then((token) => - this.grpc.search(name, consistencyLevel, tenant, `Bearer ${token}`) + this.grpc.search(collection, consistencyLevel, tenant, `Bearer ${token}`) ); } - return new Promise((resolve) => resolve(this.grpc.search(name, consistencyLevel, tenant))); + return new Promise((resolve) => resolve(this.grpc.search(collection, consistencyLevel, tenant))); }; - batch = (name: string, consistencyLevel?: ConsistencyLevel, tenant?: string) => { + batch = (collection: string, consistencyLevel?: ConsistencyLevel, tenant?: string) => { if (this.authEnabled) { - return this.login().then((token) => this.grpc.batch(name, consistencyLevel, tenant, `Bearer ${token}`)); + return this.login().then((token) => + this.grpc.batch(collection, consistencyLevel, tenant, `Bearer ${token}`) + ); + } + return new Promise((resolve) => resolve(this.grpc.batch(collection, consistencyLevel, tenant))); + }; + + tenants = (collection: string) => { + if (this.authEnabled) { + return this.login().then((token) => this.grpc.tenants(collection, `Bearer ${token}`)); } - return new Promise((resolve) => resolve(this.grpc.batch(name, consistencyLevel, tenant))); + return new Promise((resolve) => resolve(this.grpc.tenants(collection))); }; close = () => { @@ -87,14 +97,20 @@ export default class ConnectionGRPC extends ConnectionGQL { export interface GrpcClient { close: () => void; - batch: (name: string, consistencyLevel?: ConsistencyLevel, tenant?: string, bearerToken?: string) => Batch; + batch: ( + collection: string, + consistencyLevel?: ConsistencyLevel, + tenant?: string, + bearerToken?: string + ) => Batch; health: () => Promise; search: ( - name: string, + collection: string, consistencyLevel?: ConsistencyLevel, tenant?: string, bearerToken?: string ) => Search; + tenants: (collection: string, bearerToken?: string) => Tenants; } export const grpcClient = (config: GrpcConnectionParams): GrpcClient => { @@ -117,10 +133,10 @@ export const grpcClient = (config: GrpcConnectionParams): GrpcClient => { const health = clientFactory.create(HealthDefinition, channel); return { close: () => channel.close(), - batch: (name: string, consistencyLevel?: ConsistencyLevel, tenant?: string, bearerToken?: string) => + batch: (collection: string, consistencyLevel?: ConsistencyLevel, tenant?: string, bearerToken?: string) => Batcher.use( client, - name, + collection, new Metadata(bearerToken ? { ...config.headers, authorization: bearerToken } : config.headers), consistencyLevel, tenant @@ -129,13 +145,24 @@ export const grpcClient = (config: GrpcConnectionParams): GrpcClient => { health .check({ service: '/grpc.health.v1.Health/Check' }) .then((res) => res.status === HealthCheckResponse_ServingStatus.SERVING), - search: (name: string, consistencyLevel?: ConsistencyLevel, tenant?: string, bearerToken?: string) => + search: ( + collection: string, + consistencyLevel?: ConsistencyLevel, + tenant?: string, + bearerToken?: string + ) => Searcher.use( client, - name, + collection, new Metadata(bearerToken ? { ...config.headers, authorization: bearerToken } : config.headers), consistencyLevel, tenant ), + tenants: (collection: string, bearerToken?: string) => + TenantsManager.use( + client, + collection, + new Metadata(bearerToken ? { ...config.headers, authorization: bearerToken } : config.headers) + ), }; }; diff --git a/src/grpc/base.ts b/src/grpc/base.ts index 2355e695..f09a1318 100644 --- a/src/grpc/base.ts +++ b/src/grpc/base.ts @@ -6,20 +6,20 @@ import { Metadata } from 'nice-grpc'; export default class Base { protected connection: WeaviateClient; - protected name: string; + protected collection: string; protected consistencyLevel?: ConsistencyLevelGRPC; protected tenant?: string; protected metadata?: Metadata; protected constructor( connection: WeaviateClient, - name: string, + collection: string, metadata: Metadata, consistencyLevel?: ConsistencyLevel, tenant?: string ) { this.connection = connection; - this.name = name; + this.collection = collection; this.consistencyLevel = this.mapConsistencyLevel(consistencyLevel); this.tenant = tenant; this.metadata = metadata; diff --git a/src/grpc/batcher.ts b/src/grpc/batcher.ts index 9ffccddf..e74afda5 100644 --- a/src/grpc/batcher.ts +++ b/src/grpc/batcher.ts @@ -28,12 +28,12 @@ export interface BatchDeleteArgs { export default class Batcher extends Base implements Batch { public static use( connection: WeaviateClient, - name: string, + collection: string, metadata: Metadata, consistencyLevel?: ConsistencyLevel, tenant?: string ): Batch { - return new Batcher(connection, name, metadata, consistencyLevel, tenant); + return new Batcher(connection, collection, metadata, consistencyLevel, tenant); } public withDelete = (args: BatchDeleteArgs) => this.callDelete(BatchDeleteRequest.fromPartial(args)); @@ -44,7 +44,7 @@ export default class Batcher extends Base implements Batch { .batchDelete( { ...message, - collection: this.name, + collection: this.collection, consistencyLevel: this.consistencyLevel, tenant: this.tenant, }, diff --git a/src/grpc/searcher.ts b/src/grpc/searcher.ts index c73d3ab7..b1b3372e 100644 --- a/src/grpc/searcher.ts +++ b/src/grpc/searcher.ts @@ -114,12 +114,12 @@ export interface Search { export default class Searcher extends Base implements Search { public static use( connection: WeaviateClient, - name: string, + collection: string, metadata: Metadata, consistencyLevel?: ConsistencyLevel, tenant?: string ): Search { - return new Searcher(connection, name, metadata, consistencyLevel, tenant); + return new Searcher(connection, collection, metadata, consistencyLevel, tenant); } public withFetch = (args: SearchFetchArgs) => this.call(SearchRequest.fromPartial(args)); @@ -140,7 +140,7 @@ export default class Searcher extends Base implements Search { .search( { ...message, - collection: this.name, + collection: this.collection, consistencyLevel: this.consistencyLevel, tenant: this.tenant, uses123Api: true, diff --git a/src/grpc/tenantsManager.ts b/src/grpc/tenantsManager.ts new file mode 100644 index 00000000..8b156fbf --- /dev/null +++ b/src/grpc/tenantsManager.ts @@ -0,0 +1,33 @@ +import Base from './base.js'; +import { Metadata } from 'nice-grpc'; +import { WeaviateClient } from '../proto/v1/weaviate.js'; +import { TenantsGetReply, TenantsGetRequest } from '../proto/v1/tenants.js'; + +export type TenantsGetArgs = { + names?: string[]; +}; + +export interface Tenants { + withGet: (args: TenantsGetArgs) => Promise; +} + +export default class TenantsManager extends Base implements TenantsManager { + public static use(connection: WeaviateClient, collection: string, metadata: Metadata): Tenants { + return new TenantsManager(connection, collection, metadata); + } + + public withGet = (args: TenantsGetArgs) => + this.call(TenantsGetRequest.fromPartial({ names: args.names ? { values: args.names } : undefined })); + + private call(message: TenantsGetRequest) { + return this.connection.tenantsGet( + { + ...message, + collection: this.collection, + }, + { + metadata: this.metadata, + } + ); + } +} diff --git a/src/utils/dbVersion.ts b/src/utils/dbVersion.ts index 997d8453..834ba6c0 100644 --- a/src/utils/dbVersion.ts +++ b/src/utils/dbVersion.ts @@ -119,6 +119,16 @@ export class DbVersionSupport { }; }); }; + + supportsTenantsGetGRPCMethod = () => { + return this.dbVersionProvider.getVersion().then((version) => { + return { + version: version, + supports: version.isAtLeast(1, 25, 0), + message: this.errorMessage('Tenants get method', version.show(), '1.25.0'), + }; + }); + }; } const EMPTY_VERSION = '';