
    <i                        S SK r S SKJr  S SKJrJr  S SKJr  S SKJ	r	    SS\S\
S\\	R                     S	\S
\SS4S jjr   SS\S\S\\
   S-  S\S\SS4S jjrS\S\
S\4S jr SS\S\\
   S-  S\\
   4S jjrS\S\\
   S\\
   4S jrS\S\S\
SS4S jrS\S\
S\\
\	R.                  4   SS4S jr SS\S\S\
S\SS4
S jjrg)    N)Iterable)to_dictmodel_fields)
QdrantBase)modelsclientcollection_namepointsmax_attemptspausereturnc           	         SnXS::  a   U R                  UUSS9  g [        S	U S
35      e! [         aM  n[        SU SU SU 35        XS:  a%  [        SU S35        [        R                  " U5        US-  n S nAOS nAff = fXS::  a  M  Nq)N   T)r	   r
   waitzException: z
, attempt /zNext attempt in z secondszFailed to upload points after z	 attempts)upload_points	Exceptionprinttimesleep)r   r	   r
   r   r   attemptses          W/home/james-whalen/.local/lib/python3.13/site-packages/qdrant_client/migrate/migrate.pyupload_with_retryr   	   s     H

"	   / ! 
  4\N)L
MM  	Ks*XJa~FG&(x89

5!MH	 
"s   * 
BAA<<Bsource_clientdest_clientcollection_namesrecreate_on_collision
batch_sizec                 l  ^  [        T U5      n[        U 4S jU 5       5      (       a  [        S5      e[        X5      n[	        U5      [	        U5      -
  nU(       a  U(       d  [        SU 35      eU H  n[        T X5        [        T XU5        M     U H  n[        T X5        [        T XU5        M     g)a  
Migrate collections from source client to destination client

Args:
    source_client (QdrantBase): Source client
    dest_client (QdrantBase): Destination client
    collection_names (list[str], optional): List of collection names to migrate.
        If None - migrate all source client collections. Defaults to None.
    recreate_on_collision (bool, optional): If True - recreate collection if it exists, otherwise
        raise ValueError.
    batch_size (int, optional): Batch size for scrolling and uploading vectors. Defaults to 100.
c              3   <   >#    U  H  n[        TU5      v   M     g 7fN)_has_custom_shards).0r	   r   s     r   	<genexpr>migrate.<locals>.<genexpr>7   s       Rb=/::Rbs   z@Migration of collections with custom shards is not supported yetz*Collections already exist in dest_client: N)_select_source_collectionsany
ValueError_find_collisionsset_recreate_collection_migrate_collection)r   r   r   r   r   
collisionsabsent_dest_collectionsr	   s   `       r   migrater0   #   s    & 2-AQR
 Rb   [\\!+@J!"23c*oE/Ej\RSS2]KIM;T 3 &]KIM;T &    c                     U R                  U5      n[        UR                  R                  SS 5      [        R
                  R                  :H  $ )Nsharding_method)get_collectiongetattrconfigparamsr   ShardingMethodCUSTOM)r   r	   collection_infos      r   r#   r#   K   sD    #22?CO&&--/@$G  ''	(r1   c                    ^ U R                  5       R                  nU Vs/ s H  o3R                  PM     snmUb<  [        U4S jU 5       5      (       d   S[	        U5      [	        T5      -
   35       e U$ TnU$ s  snf )Nc              3   ,   >#    U  H	  oT;   v   M     g 7fr"    )r$   r	   source_collection_namess     r   r%   -_select_source_collections.<locals>.<genexpr>Z   s      
N^?66N^s   z)Source client does not have collections: )get_collectionscollectionsnameallr+   )r   r   source_collections
collectionr>   s       @r   r'   r'   S   s     '668DDASTAS:AST# 
N^
 
 
 	n6s;K7LsSjOk7k6lm	n 
  3 Us   A;c                     U R                  5       R                  nU Vs1 s H  o3R                  iM     nnU[        U5      -  n[	        U5      $ s  snf r"   )r@   rA   rB   r+   list)r   r   dest_collectionsrE   dest_collection_namesexisting_dest_collectionss         r   r*   r*   c   sS    "224@@?OP?O__?OP 5<L8M M)** Qs   Ac                    U R                  U5      nUR                  nUR                  nUR                  U5      (       a  UR	                  U5        S nUR
                  bk  [        R                  " S0 [        UR
                  5      R                  5        VVs0 s H&  u  pxU[        [        R                  5      ;   d  M$  Xx_M(     snnD6nUR                  UUR                  R                  UR                  R                  UR                  R                  UR                  R                   UR                  R"                  UR                  R$                  [        R&                  " S0 [        UR(                  5      D6[        R*                  " S0 [        UR,                  5      D6[        R.                  " S0 [        UR0                  5      D6UR2                  US9  [5        XU5        g s  snnf )N)vectors_configsparse_vectors_configshard_numberreplication_factorwrite_consistency_factoron_disk_payloadhnsw_configoptimizers_config
wal_configquantization_configstrict_mode_configr=   )r4   r6   payload_schemacollection_existsdelete_collectionrV   r   StrictModeConfigr   itemsr   create_collectionr7   vectorssparse_vectorsrN   rO   rP   rQ   HnswConfigDiffrR   OptimizersConfigDiffoptimizer_configWalConfigDiffrT   rU   _recreate_payload_schema)	r   r   r	   src_collection_info
src_configsrc_payload_schemarV   kvs	            r   r,   r,   j   s   
 (66G$++J,;;$$_55%%o69=$$0#44 
 $J$A$ABHHJJDAV%<%<== J
 !!!((00(//>>&&33%,,??!+!2!2!K!K"))99))LGJ4J4J,KL 55]
@[@[8\]''I'*2G2G*HI&::- "  [;MN+s   #G)?G)rW   c                     UR                  5        H9  u  p4U R                  UUUR                  c  UR                  OUR                  S9  M;     g )N)
field_namefield_schema)r[   create_payload_indexr7   	data_type)r   r	   rW   rj   
field_infos        r   rc   rc      sP    
 #1"6"6"8
((!1;1B1B1J--PZPaPa 	) 	
 #9r1   c                    U R                  USSS9u  pE[        XUS9  Ub"  U R                  X%USS9u  pE[        XUS9  Ub  M"  U R                  U5      R                  nUR                  U5      R                  nXg:X  d   SU SU 35       eg)	a.  Migrate collection from source client to destination client

Args:
    collection_name (str): Collection name
    source_client (QdrantBase): Source client
    dest_client (QdrantBase): Destination client
    batch_size (int, optional): Batch size for scrolling and uploading vectors. Defaults to 100.
   T)limitwith_vectors)r   r	   r
   N)offsetrq   rr   zCMigration failed, vectors count are not equal: source vector count z, dest vector count )scrollr   count)r   r   r	   r   recordsnext_offsetsource_client_vectors_countdest_client_vectors_counts           r   r-   r-      s     )//qW[/\G[RYZ

!,33zPT  4  
 	V]^	 
!
 #0"5"5o"F"L"L + 1 1/ B H H#@Z	LMhLii}  X  ~Y  ZZ@r1   )   g      @)NFd   r"   )r{   )r   typingr   qdrant_client._pydantic_compatr   r   qdrant_client.client_baser   qdrant_client.httpr   strPointStructintfloatr   rG   boolr0   r#   r'   r*   r,   dictPayloadIndexInforc   r-   r=   r1   r   <module>r      s     @ 0 % NNN V''(N 	N
 N 
N: *."'%U%U%U 3i$&%U  	%U
 %U 
%UPj 3 4  EI15cT1A	#Y +* +S	 +dSVi +#O#O#O #O 
	#OL





 f5556

 
	

" 	ZZZ Z 	Z
 
Zr1   