Skip to content

Commit 70ee2c3

Browse files
committed
feat: add volume manager api
Signed-off-by: ryjiang <jiangruiyi@gmail.com>
1 parent 8dd59d1 commit 70ee2c3

6 files changed

Lines changed: 411 additions & 1 deletion

File tree

milvus/bulkwriter/VolumeManager.ts

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,195 @@
1+
import {
2+
DEFAULT_HTTP_TIMEOUT,
3+
VolumeApplyReq,
4+
VolumeCreateReq,
5+
VolumeListReq,
6+
VolumeManagerConfig,
7+
VolumeNameReq,
8+
VolumeRequestOptions,
9+
VolumeResponse,
10+
} from '../';
11+
12+
export class VolumeManager {
13+
private readonly cloudEndpoint: string;
14+
private readonly apiKey: string;
15+
private readonly fetchImpl: typeof fetch;
16+
private readonly timeout: number;
17+
18+
constructor(config: VolumeManagerConfig);
19+
constructor(cloudEndpoint: string, apiKey: string, fetchImpl?: typeof fetch);
20+
constructor(
21+
configOrEndpoint: VolumeManagerConfig | string,
22+
apiKey?: string,
23+
fetchImpl?: typeof fetch
24+
) {
25+
if (typeof configOrEndpoint === 'string') {
26+
this.cloudEndpoint = configOrEndpoint;
27+
this.apiKey = apiKey || '';
28+
this.fetchImpl = fetchImpl || fetch;
29+
this.timeout = DEFAULT_HTTP_TIMEOUT;
30+
} else {
31+
this.cloudEndpoint = configOrEndpoint.cloudEndpoint;
32+
this.apiKey = configOrEndpoint.apiKey;
33+
this.fetchImpl = configOrEndpoint.fetch || fetch;
34+
this.timeout = configOrEndpoint.timeout || DEFAULT_HTTP_TIMEOUT;
35+
}
36+
}
37+
38+
async createVolume(
39+
data: VolumeCreateReq,
40+
options?: VolumeRequestOptions
41+
): Promise<VolumeResponse> {
42+
const { timeout, ...params } = data;
43+
return this.POST('/v2/volumes/create', params, {
44+
...options,
45+
timeout: options?.timeout ?? timeout ?? this.timeout,
46+
});
47+
}
48+
49+
async listVolumes(
50+
data: VolumeListReq,
51+
options?: VolumeRequestOptions
52+
): Promise<VolumeResponse> {
53+
const { timeout, ...params } = data;
54+
return this.GET('/v2/volumes', params, {
55+
...options,
56+
timeout: options?.timeout ?? timeout ?? this.timeout,
57+
});
58+
}
59+
60+
async deleteVolume(
61+
data: VolumeNameReq,
62+
options?: VolumeRequestOptions
63+
): Promise<VolumeResponse> {
64+
const { timeout, volumeName } = data;
65+
return this.DELETE(`/v2/volumes/${encodeURIComponent(volumeName)}`, {
66+
...options,
67+
timeout: options?.timeout ?? timeout ?? this.timeout,
68+
});
69+
}
70+
71+
async describeVolume(
72+
data: VolumeNameReq,
73+
options?: VolumeRequestOptions
74+
): Promise<VolumeResponse> {
75+
const { timeout, volumeName } = data;
76+
return this.GET(
77+
`/v2/volumes/${encodeURIComponent(volumeName)}`,
78+
{},
79+
{
80+
...options,
81+
timeout: options?.timeout ?? timeout ?? this.timeout,
82+
}
83+
);
84+
}
85+
86+
async applyVolume(
87+
data: VolumeApplyReq,
88+
options?: VolumeRequestOptions
89+
): Promise<VolumeResponse> {
90+
const { timeout, ...params } = data;
91+
return this.POST('/v2/volumes/apply', params, {
92+
...options,
93+
timeout: options?.timeout ?? timeout ?? this.timeout,
94+
});
95+
}
96+
97+
private get headers() {
98+
return {
99+
Authorization: `Bearer ${this.apiKey}`,
100+
Accept: 'application/json',
101+
ContentType: 'application/json',
102+
};
103+
}
104+
105+
private get baseURL() {
106+
return this.cloudEndpoint.replace(/\/+$/, '');
107+
}
108+
109+
private async handleResponse<T>(response: Response, url: string): Promise<T> {
110+
if (!response.ok) {
111+
const errorText = await response.text().catch(() => '');
112+
throw new Error(
113+
`HTTP ${response.status} ${response.statusText}: ${url}${
114+
errorText ? ` - ${errorText}` : ''
115+
}`
116+
);
117+
}
118+
return response.json() as T;
119+
}
120+
121+
private async request<T>(
122+
method: 'GET' | 'POST' | 'DELETE',
123+
path: string,
124+
data?: Record<string, any>,
125+
options?: VolumeRequestOptions
126+
): Promise<T> {
127+
let id: ReturnType<typeof setTimeout> | undefined;
128+
const timeout = options?.timeout ?? this.timeout;
129+
const abortController = options?.abortController ?? new AbortController();
130+
131+
try {
132+
id = setTimeout(() => abortController.abort(), timeout);
133+
const url = this.buildUrl(path, method === 'GET' ? data : undefined);
134+
const init: RequestInit = {
135+
method,
136+
headers: this.headers,
137+
signal: abortController.signal,
138+
};
139+
140+
if (method === 'POST') {
141+
init.body = JSON.stringify(data || {});
142+
}
143+
144+
const response = await this.fetchImpl(url, init);
145+
return this.handleResponse<T>(response, url);
146+
} finally {
147+
if (id !== undefined) {
148+
clearTimeout(id);
149+
}
150+
}
151+
}
152+
153+
private buildUrl(path: string, params?: Record<string, any>) {
154+
const url = `${this.baseURL}${path}`;
155+
if (!params || Object.keys(params).length === 0) {
156+
return url;
157+
}
158+
159+
const searchParams = new URLSearchParams();
160+
Object.entries(params).forEach(([key, value]) => {
161+
if (value !== undefined && value !== null) {
162+
searchParams.append(key, String(value));
163+
}
164+
});
165+
166+
const query = searchParams.toString();
167+
return query ? `${url}?${query}` : url;
168+
}
169+
170+
private GET<T>(
171+
path: string,
172+
params?: Record<string, any>,
173+
options?: VolumeRequestOptions
174+
) {
175+
return this.request<T>('GET', path, params, options);
176+
}
177+
178+
private POST<T>(
179+
path: string,
180+
data?: Record<string, any>,
181+
options?: VolumeRequestOptions
182+
) {
183+
return this.request<T>('POST', path, data, options);
184+
}
185+
186+
private DELETE<T>(path: string, options?: VolumeRequestOptions) {
187+
return this.request<T>('DELETE', path, undefined, options);
188+
}
189+
190+
create_volume = this.createVolume;
191+
list_volumes = this.listVolumes;
192+
delete_volume = this.deleteVolume;
193+
describe_volume = this.describeVolume;
194+
apply_volume = this.applyVolume;
195+
}

milvus/bulkwriter/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
export { BulkWriter } from './BulkWriter';
2+
export { VolumeManager } from './VolumeManager';
23
export { ColumnBuffer } from './ColumnBuffer';
34
export { JsonFormatter } from './JsonFormatter';
45
export { ParquetFormatter } from './ParquetFormatter';

milvus/types.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,4 +13,5 @@ export * from './types/Http';
1313
export * from './types/Segments';
1414
export * from './types/Insert';
1515
export * from './types/Search';
16-
export * from './types/DataTypes';
16+
export * from './types/DataTypes';
17+
export * from './types/Volume';

milvus/types/Volume.ts

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
import { FetchOptions, HttpBaseResponse } from './Http';
2+
3+
export enum VolumeType {
4+
MANAGED = 'MANAGED',
5+
EXTERNAL = 'EXTERNAL',
6+
}
7+
8+
export interface VolumeBaseReq {
9+
timeout?: number;
10+
}
11+
12+
export interface VolumeCreateReq extends VolumeBaseReq {
13+
projectId: string;
14+
regionId: string;
15+
volumeName: string;
16+
type?: VolumeType | keyof typeof VolumeType | string;
17+
storageIntegrationId?: string;
18+
path?: string;
19+
}
20+
21+
export interface VolumeListReq extends VolumeBaseReq {
22+
projectId: string;
23+
currentPage?: number;
24+
pageSize?: number;
25+
type?: VolumeType | keyof typeof VolumeType | string;
26+
}
27+
28+
export interface VolumeNameReq extends VolumeBaseReq {
29+
volumeName: string;
30+
}
31+
32+
export interface VolumeApplyReq extends VolumeNameReq {
33+
path: string;
34+
}
35+
36+
export type VolumeResponse<T = Record<string, any>> = HttpBaseResponse<T>;
37+
38+
export interface VolumeManagerConfig {
39+
cloudEndpoint: string;
40+
apiKey: string;
41+
fetch?: typeof fetch;
42+
timeout?: number;
43+
}
44+
45+
export interface VolumeRequestOptions extends Partial<FetchOptions> {}

milvus/types/index.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,3 +17,4 @@ export * from './Insert';
1717
export * from './Search';
1818
export * from './DataTypes';
1919
export * from './GlobalCluster';
20+
export * from './Volume';

0 commit comments

Comments
 (0)