@@ -225,105 +225,115 @@ async def _migrate_kvstore_to_sql(self) -> None:
225225
226226 sql_store = self .metadata_store .sql_store
227227
228- existing = await sql_store .fetch_all (table = TABLE_VECTOR_STORES , limit = 1 )
229- if existing .data :
230- return
231-
232228 stores_data = await self .kvstore .values_in_range (
233229 OPENAI_VECTOR_STORES_PREFIX , f"{ OPENAI_VECTOR_STORES_PREFIX } \xff "
234230 )
235231 if not stores_data :
236232 return
237233
238- logger .info ("Starting KVStore to SQL migration for vector store metadata" , store_count = len (stores_data ))
239-
240234 migrated_stores = 0
241235 migrated_files = 0
242236 migrated_chunks = 0
243237 migrated_batches = 0
244238
245- for raw in stores_data :
246- info = json .loads (raw )
247- store_id = info ["id" ]
248- await sql_store .insert (
249- table = TABLE_VECTOR_STORES ,
250- data = {
251- "id" : store_id ,
252- "store_data" : info ,
253- "owner_principal" : "" ,
254- "access_attributes" : None ,
255- },
256- )
257- migrated_stores += 1
258-
259- file_keys = await self .kvstore .keys_in_range (
260- f"{ OPENAI_VECTOR_STORES_FILES_PREFIX } { store_id } :" ,
261- f"{ OPENAI_VECTOR_STORES_FILES_PREFIX } { store_id } :\xff " ,
239+ # Per-table migration: each table is checked independently so a crash
240+ # mid-migration doesn't skip remaining tables on the next boot.
241+ existing_stores = await sql_store .fetch_all (table = TABLE_VECTOR_STORES , limit = 1 )
242+ if not existing_stores .data :
243+ logger .info (
244+ "Starting KVStore to SQL migration for vector store metadata" ,
245+ store_count = len (stores_data ),
262246 )
263- for file_key in file_keys :
264- suffix = file_key [len (OPENAI_VECTOR_STORES_FILES_PREFIX ) :]
265- file_id = suffix .split (":" , 1 )[1 ] if ":" in suffix else suffix
266- raw_file = await self .kvstore .get (file_key )
267- if not raw_file :
268- continue
269- file_info = json .loads (raw_file )
247+ for raw in stores_data :
248+ info = json .loads (raw )
249+ store_id = info ["id" ]
270250 await sql_store .insert (
271- table = TABLE_VECTOR_STORE_FILES ,
251+ table = TABLE_VECTOR_STORES ,
272252 data = {
273- "id" : f"{ store_id } :{ file_id } " ,
274- "store_id" : store_id ,
275- "file_id" : file_id ,
276- "file_data" : file_info ,
253+ "id" : store_id ,
254+ "store_data" : info ,
277255 "owner_principal" : "" ,
278256 "access_attributes" : None ,
279257 },
280258 )
281- migrated_files += 1
282-
283- chunk_prefix = f"{ OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX } { store_id } :{ file_id } :"
284- chunk_values = await self .kvstore .values_in_range (chunk_prefix , f"{ chunk_prefix } \xff " )
285- for idx , raw_chunk in enumerate (chunk_values ):
286- chunk = json .loads (raw_chunk )
259+ migrated_stores += 1
260+
261+ existing_files = await sql_store .fetch_all (table = TABLE_VECTOR_STORE_FILES , limit = 1 )
262+ if not existing_files .data :
263+ for raw in stores_data :
264+ info = json .loads (raw )
265+ store_id = info ["id" ]
266+ file_keys = await self .kvstore .keys_in_range (
267+ f"{ OPENAI_VECTOR_STORES_FILES_PREFIX } { store_id } :" ,
268+ f"{ OPENAI_VECTOR_STORES_FILES_PREFIX } { store_id } :\xff " ,
269+ )
270+ for file_key in file_keys :
271+ suffix = file_key [len (OPENAI_VECTOR_STORES_FILES_PREFIX ) :]
272+ file_id = suffix .split (":" , 1 )[1 ] if ":" in suffix else suffix
273+ raw_file = await self .kvstore .get (file_key )
274+ if not raw_file :
275+ continue
276+ file_info = json .loads (raw_file )
287277 await sql_store .insert (
288- table = TABLE_VECTOR_STORE_FILE_CONTENTS ,
278+ table = TABLE_VECTOR_STORE_FILES ,
289279 data = {
290- "id" : f"{ store_id } :{ file_id } : { idx } " ,
280+ "id" : f"{ store_id } :{ file_id } " ,
291281 "store_id" : store_id ,
292282 "file_id" : file_id ,
293- "chunk_index" : idx ,
294- "chunk_data" : chunk ,
283+ "file_data" : file_info ,
295284 "owner_principal" : "" ,
296285 "access_attributes" : None ,
297286 },
298287 )
299- migrated_chunks += 1
288+ migrated_files += 1
289+
290+ chunk_prefix = f"{ OPENAI_VECTOR_STORES_FILES_CONTENTS_PREFIX } { store_id } :{ file_id } :"
291+ chunk_values = await self .kvstore .values_in_range (chunk_prefix , f"{ chunk_prefix } \xff " )
292+ for idx , raw_chunk in enumerate (chunk_values ):
293+ chunk = json .loads (raw_chunk )
294+ await sql_store .insert (
295+ table = TABLE_VECTOR_STORE_FILE_CONTENTS ,
296+ data = {
297+ "id" : f"{ store_id } :{ file_id } :{ idx } " ,
298+ "store_id" : store_id ,
299+ "file_id" : file_id ,
300+ "chunk_index" : idx ,
301+ "chunk_data" : chunk ,
302+ "owner_principal" : "" ,
303+ "access_attributes" : None ,
304+ },
305+ )
306+ migrated_chunks += 1
300307
301- batch_data = await self .kvstore .values_in_range (
302- OPENAI_VECTOR_STORES_FILE_BATCHES_PREFIX , f"{ OPENAI_VECTOR_STORES_FILE_BATCHES_PREFIX } \xff "
303- )
304- for raw_batch in batch_data :
305- batch_info = json .loads (raw_batch )
306- batch_id = batch_info ["id" ]
307- await sql_store .insert (
308- table = TABLE_VECTOR_STORE_FILE_BATCHES ,
309- data = {
310- "id" : batch_id ,
311- "store_id" : batch_info .get ("vector_store_id" , "" ),
312- "batch_data" : batch_info ,
313- "expires_at" : batch_info .get ("expires_at" , 0 ),
314- "owner_principal" : "" ,
315- "access_attributes" : None ,
316- },
308+ existing_batches = await sql_store .fetch_all (table = TABLE_VECTOR_STORE_FILE_BATCHES , limit = 1 )
309+ if not existing_batches .data :
310+ batch_data = await self .kvstore .values_in_range (
311+ OPENAI_VECTOR_STORES_FILE_BATCHES_PREFIX , f"{ OPENAI_VECTOR_STORES_FILE_BATCHES_PREFIX } \xff "
312+ )
313+ for raw_batch in batch_data :
314+ batch_info = json .loads (raw_batch )
315+ batch_id = batch_info ["id" ]
316+ await sql_store .insert (
317+ table = TABLE_VECTOR_STORE_FILE_BATCHES ,
318+ data = {
319+ "id" : batch_id ,
320+ "store_id" : batch_info .get ("vector_store_id" , "" ),
321+ "batch_data" : batch_info ,
322+ "expires_at" : batch_info .get ("expires_at" , 0 ),
323+ "owner_principal" : "" ,
324+ "access_attributes" : None ,
325+ },
326+ )
327+ migrated_batches += 1
328+
329+ if migrated_stores or migrated_files or migrated_chunks or migrated_batches :
330+ logger .info (
331+ "KVStore to SQL migration complete" ,
332+ stores = migrated_stores ,
333+ files = migrated_files ,
334+ chunks = migrated_chunks ,
335+ batches = migrated_batches ,
317336 )
318- migrated_batches += 1
319-
320- logger .info (
321- "KVStore to SQL migration complete" ,
322- stores = migrated_stores ,
323- files = migrated_files ,
324- chunks = migrated_chunks ,
325- batches = migrated_batches ,
326- )
327337
328338 async def _save_openai_vector_store (self , store_id : str , store_info : dict [str , Any ]) -> None :
329339 """Save vector store metadata to persistent storage."""
0 commit comments