o
    g                     @   sL  d dl Z d dlZd dlZd dlmZmZmZmZmZm	Z	m
Z
mZ d dlmZ d dlmZ d dlmZmZ d dlmZ d dlmZmZmZmZ d dlmZmZmZ d d	lmZm Z m!Z!m"Z" d d
l#m$Z$ d dl%m&Z&m'Z'm(Z( d dl)m*Z*m+Z+m,Z,m-Z- d dl.m/Z/ d dl0m1Z1m2Z2m3Z3m4Z4m5Z5m6Z6m7Z7m8Z8m9Z9m:Z: e;e<Z=G dd de e$Z>dS )    N)OptionalSequenceAnyTuplecastDictUnionSet)UUID)override)TableColumn)groupby)CollectionConfigurationInternalConfigurationParameterHNSWConfigurationInternalInvalidConfigurationError)DEFAULT_DATABASEDEFAULT_TENANTSystem)CursorSqlDBParameterValueget_sql)SysDB)InvalidCollectionExceptionNotFoundErrorUniqueConstraintError)add_attributes_to_current_spanOpenTelemetryClientOpenTelemetryGranularitytrace_method)Producer)
CollectionAndSegmentsDatabaseOptionalArgumentSegmentMetadata
CollectionSegmentScopeTenantUnspecifiedUpdateMetadatac                       s  e Zd ZU eed< def fddZedej	e
dK fdd	Ze
efd
edededdfddZe
efdededefddZe
efdededdfddZe
ddefdee dee dedee fddZe
deddfddZe
dedefddZdededdfddZedej	e
deddfddZed ej	e
ddd!eefd
eded"ed#ee d$ee  d%ee d&e!ded'ede"e#e!f fd(d Z$ed)ej	e
			dLd*ed
ee d+ee d,ee% dee f
d-d)Z&ed.ej	e
ddeeddfd
ee dee ded'edee dee dee# fd/d.Z'e
d0ede(fd1d2Z)ed3ej	e
d*ed
eddfd4d3Z*ded*eddfd5d6Z+ed7ej	e
eefd
eded'eddfd8d7Z,ed9ej	e
e- fd*ed
ed$e.ee/  ddfd:d9Z0ed;ej	e
e- e- e- fd
ede.e d%e.ee  d$e.ee/  ddf
d<d;Z1ed=ej	d>ee"e2d?f  dee  fd@d=Z3edAej		dMdedBe4dCe5d
ed$e/dDee6e  ddfdEdAZ7d0edFedefdGdHZ8d0e2d$ee  defdIdJZ9  Z:S )NSqlSysDB	_producersystemc                    s   t  | |t| _d S N)super__init__requirer   _opentelemetry_client)selfr/   	__class__ O/var/www/visachat/venv/lib/python3.10/site-packages/chromadb/db/mixins/sysdb.pyr2   4   s   zSqlSysDB.__init__zSqlSysDB.create_segmentreturnNc                    s   t    | jt| _d S r0   )r1   start_systeminstancer"   r.   )r5   r6   r8   r9   r;   8   s   
zSqlSysDB.startidnametenantc                 C   s   |   h}td}td}|  ||j|j|jt	| 
|t	||  |j||jt	|k}t||  \}}	z|||	 W n |  yc }
 ztd| d| |
d }
~
ww W d    d S 1 sow   Y  d S )N	databasestenants	Database z already exists for tenant )txr   querybuilderintocolumnsr>   r?   	tenant_idinsertr   
uuid_to_dbselectfrom_wherer   parameter_formatexecuteunique_constraint_errorr   )r5   r>   r?   r@   currA   rB   insert_databasesqlparamser8   r8   r9   create_database>   s8   
"zSqlSysDB.create_databasec           
      C   s   |   k}td}|  ||j|j|jt|k|j	t|k}t
||  \}}||| }|sDtd| d| d|d d u rUtd| d| dtt| |d }	t|	|d |dW  d    S 1 srw   Y  d S )NrA   rC    not found for tenant z. Are you sure it exists?r      r>   r?   r@   )rD   r   rE   rL   rK   r>   r?   rM   r   rH   r   rN   rO   fetchoner   r   r
   uuid_from_dbr$   )
r5   r?   r@   rQ   rA   qrS   rT   rowr>   r8   r8   r9   get_database[   s2   
$zSqlSysDB.get_databasec           
      C   s   |   m}td}|  ||jt|k|jt|k }t	|| 
 \}}|d }||| }|sCtd| d| td}	|  |	|	jt|d k }t	|| 
 \}}||| W d    d S 1 stw   Y  d S )NrA    RETURNING idrC   rW   collectionsr   )rD   r   rE   rL   rM   r?   r   rH   deleter   rN   rO   rZ   r   database_id)
r5   r?   r@   rQ   rA   r\   rS   rT   resultr`   r8   r8   r9   delete_databasew   s,   
"zSqlSysDB.delete_databaselimitoffsetc           
         s      M}td}  ||j|j|jt	k
||d u r(tjn||j}t|  \}}||| }	 fdd|	D W  d    S 1 sTw   Y  d S )NrA   c              	      s.   g | ]}t tt |d  |d dqS )r   rX   rY   )r$   r   r
   r[   ).0r]   r5   r@   r8   r9   
<listcomp>   s    z+SqlSysDB.list_databases.<locals>.<listcomp>)rD   r   rE   rL   rK   r>   r?   rM   rH   r   rf   re   sysmaxsizeorderby
created_atr   rN   rO   fetchall)
r5   re   rf   r@   rQ   rA   r\   rS   rT   rowsr8   rh   r9   list_databases   s$   
$zSqlSysDB.list_databasesc                 C   s   |   F}td}|  ||jt|}t|| 	 \}}z|
|| W n |  yA } z	td| d|d }~ww W d    d S 1 sMw   Y  d S )NrB   Tenant  already exists)rD   r   rE   rF   rG   r>   rI   r   r   rN   rO   rP   r   )r5   r?   rQ   rB   insert_tenantrS   rT   rU   r8   r8   r9   create_tenant   s    

"zSqlSysDB.create_tenantc                 C   s   |   ?}td}|  ||j|jt|k}t|| 	 \}}|
|| }|s7td| dt|dW  d    S 1 sFw   Y  d S )NrB   rq   
 not found)r?   )rD   r   rE   rL   rK   r>   rM   r   r   rN   rO   rZ   r   r*   )r5   r?   rQ   rB   r\   rS   rT   r]   r8   r8   r9   
get_tenant   s   
$zSqlSysDB.get_tenantrQ   segmentc           	   
   C   sH  t t|d |d |d jt|d d td}|  ||j|j|j	|j
t| |d t|d t|d jt| |d }t||  \}}z||| W n |  yr } ztd|d  d|d }~ww td	}|d
 rz| |||j|d |d
  W d S  ty } z	td|   d }~ww d S )Nr>   typescope
collection
segment_idsegment_typesegment_scoperz   segmentsSegment rr   segment_metadatametadataz"Error inserting segment metadata: )r   strvaluer   rE   rF   rG   r>   rx   ry   rz   rI   r   rJ   r   rN   rO   rP   r   _insert_metadatar|   	Exceptionloggererror)	r5   rQ   rw   r   insert_segmentrS   rT   rU   
metadata_tr8   r8   r9   create_segment_with_tx   sd   

	

zSqlSysDB.create_segment_with_txc                 C   s:   |   }| || W d    d S 1 sw   Y  d S r0   )rD   r   )r5   rw   rQ   r8   r8   r9   create_segment  s   
"zSqlSysDB.create_collectionFconfigurationr   r   	dimensionget_or_createdatabasec
                 C   s  |d u r
|s
t dtt||d | j|||	d}
|
r8|r0|
d }| j|j||	dd dfS td| dt|||||||	dd	}|  }td
}td}| 	 
||j|j|j|j|jt| |d t|d t| t|d | 	 |j||jt|	k|jt|k}t||  \}}z||| W n |  y } ztd|d  d|d }~ww td}|d r| |||j|j|d  |D ]}| || qW d    |dfS 1 sw   Y  |dfS )Nz.id must be specified if get_or_create is Falsecollection_idcollection_name)r?   r@   r   r   )r>   r@   r   FCollection rr   r>   r?   r   r   r   r@   r   versionr`   rA   r>   r?   r   collection_metadatar   T)
ValueErrorr   r   get_collectionsr>   r   r(   rD   r   rE   rF   rG   r?   config_json_strr   rb   rI   r   rJ   to_json_strrK   rL   rM   rH   r   rN   rO   rP   r   r   r   )r5   r>   r?   r   r   r   r   r   r@   r   existingrz   rQ   r`   rA   insert_collectionrS   rT   rU   r   rw   r8   r8   r9   create_collection  s   




//zSqlSysDB.get_segmentsrz   rx   ry   c                 C   s  t t||r|nd|r|jndt|d td}td}|  ||j|j|j	|j
|j|j|j|j|j	||j|jk|j}|rV||jt| |k}|rb||jt|k}|ro||j	t|jk}|r~||j
t| |k}|  f}t||  \}	}
||	|
 }t|dd }g }|D ]?\}}| t|}t|}t|d d }tt|d d	 }| |d d
 }|  |}|!t"t#t$|||||i d q|W  d    S 1 sw   Y  d S )N r{   r   r   c                 S      t t| d S Nr   r   objectrr8   r8   r9   <lambda>      z'SqlSysDB.get_segments.<locals>.<lambda>r   rX         )r>   rx   ry   rz   r   
file_paths)%r   r   r   r   rE   rL   rK   r>   rx   ry   rz   key	str_value	int_valuefloat_value
bool_value	left_joinonr|   rl   rM   r   rJ   rD   r   rN   rO   rn   r   r[   listr)   _metadata_from_rowsappendr&   r   r
   )r5   rz   r>   rx   ry   
segments_tr   r\   rQ   rS   rT   ro   
by_segmentr   r|   segment_rowsr   r8   r8   r9   get_segmentso  sx   	


$zSqlSysDB.get_collectionsc                 C   s  |dur|du s|du rt dtt||r|ndd td}td}td}	|  ||j|j|j	|j
|	j|	j|j|j|j|j|j||j|jk|	|j|	jk|j}
|ro|
|jt| |k}
|r{|
|jt|k}
|du r|r|rtd}	|
|j|  |	j|	|	jt|k|	jt|kk}
|  }t|
|  \}}||| }t|dd	 }g }|D ]f\}}| t|}t |}t|d
 d }| !|}|d
 d rt"|d
 d nd}|d
 d dur| #t||d
 d }n| $||}|%t&t't(|||||t|d
 d t|d
 d d
d q|durE|du r<d
}||||  }n||d }|W  d   S 1 sXw   Y  dS )z;Get collections by name, embedding function and/or metadataNzmIf name is specified, tenant and database must also be specified in order to uniquely identify the collectionr   r   r`   r   rA   c                 S   r   r   r   r   r8   r8   r9   r      r   z*SqlSysDB.get_collections.<locals>.<lambda>r   rX   r   r         r   ))r   r   r   r   rE   rL   rK   r>   r?   r   r   rH   r   r   r   r   r   r   r   r   rb   rl   rM   r   rJ   rD   r   rN   rO   rn   r   r[   r   r   int&_load_config_from_json_str_and_migrate!_insert_config_from_legacy_paramsr   r(   r   r
   )r5   r>   r?   r@   r   re   rf   collections_tr   databases_tr\   rQ   rS   rT   ro   by_collectionr`   r   collection_rowsr   r   r   r8   r8   r9   r     s   



 

&r   c                 C   s@   | j |d}t|dkrtd| dt|d | j|ddS )N)r>   r   r   z does not exist.)rz   )rz   r   )r   lenr   r#   r   )r5   r   r`   r8   r8   r9   get_collection_with_segments,  s   

z%SqlSysDB.get_collection_with_segmentszSqlSysDB.delete_segmentc           	      C   s   t dt|i td}|  ||jt| |k	 }| 
 (}t||  \}}|d }||| }|sDtd| dW d   dS 1 sOw   Y  dS )zDelete a segment from the SysDBr|   r   r_   r   ru   N)r   r   r   rE   rL   rM   r>   r   rJ   ra   rD   r   rN   rO   rZ   r   )	r5   rz   r>   tr\   rQ   rS   rT   rc   r8   r8   r9   delete_segment:  s$   
"c                 C   sR   t d}|  ||jt| |k }t|| 	 \}}|
|| d S )Nr   )r   rE   rL   rM   rz   r   rJ   ra   r   rN   rO   )r5   rQ   rz   r   r\   rS   rT   r8   r8   r9   delete_segments_for_collectionT  s   z'SqlSysDB.delete_segments_for_collectionzSqlSysDB.delete_collectionc                 C   s  t dt|i td}td}|  ||jt| |k|j	|  
|j||jt|k|jt|kk }|  -}t||  \}}	|d }|||	 }
|
sgtd| d| || W d   n1 sww   Y  | j|
d  dS )	z{Delete a collection and all associated segments from the SysDB. Deletes
        the log stream for this collection as well.r   r`   rA   r_   r   ru   Nr   )r   r   r   rE   rL   rM   r>   r   rJ   rb   rK   r?   rH   ra   rD   r   rN   rO   rZ   r   r   r.   
delete_log)r5   r>   r@   r   r   r   r\   rQ   rS   rT   rc   r8   r8   r9   delete_collection_  s:   


zSqlSysDB.update_segmentc           
      C   sZ  t t|t|d td}td}|  ||jt| |k	|j
t| |}|  p}t||  \}}	|rE|||	 |d u rm|  ||jt| |k }t||  \}}	|||	 n.|t krtt|}tt|}| |||j||t	|  W d    d S W d    d S W d    d S 1 sw   Y  d S )N)r|   rz   r   r   )r   r   r   rE   updaterM   r>   r   rJ   setrz   rD   r   rN   rO   rL   r|   ra   r+   r   r,   r   keys)
r5   rz   r>   r   r   r   r\   rQ   rS   rT   r8   r8   r9   update_segment  sP   




"zSqlSysDB.update_collectionc                 C   s  t dt|i td}td}|  ||jt| |k}|t	 ks0|
|jt|}|t	 ks>|
|jt|}|  z}t||  \}	}
|	rd|	d }	||	|
}| sdtd| d|t	 kr|  ||jt| |k }t||  \}	}
||	|
 |d urtt|}| |||j||t
|  W d    d S W d    d S W d    d S 1 sw   Y  d S )Nr   r`   r   r_   r   ru   )r   r   r   rE   r   rM   r>   r   rJ   r+   r   r?   r   rD   r   rN   rO   rZ   r   rL   r   ra   r   r,   r   r   )r5   r>   r?   r   r   r   r   r\   rQ   rS   rT   rc   r8   r8   r9   update_collection  s\   	





"zSqlSysDB._metadata_from_rowsro   .c                 C   s   t dt|i i }|D ]C}t|d }|d dur#t|d ||< q|d dur2t|d ||< q|d durAt|d ||< q|d durOt|d ||< q|pSdS )zGiven SQL rows, return a metadata map (assuming that the last four columns
        are the key, str_value, int_value & float_value)num_rowsN)r   r   r   r   floatbool)r5   ro   r   r]   r   r8   r8   r9   r     s"   zSqlSysDB._insert_metadatatableid_col
clear_keysc                 C   s  t dt|i |r8|  ||t| |k|jdd |D 	 }t
||  \}}	|||	 |  |||j|j|j|j|j}| |}
| D ]c\}}t|trq|t|
t|d d d tt|}qUt|tr|t|
t|t|d d d }qUt|tr|t|
t|d t|d d }qUt|tr|t|
t|d d t|d }qU|d u rqUqUt
||  \}}	|r|||	 d S d S )Nnum_keysc                 S   s   g | ]}t |qS r8   )r   )rg   kr8   r8   r9   ri   *  s    z-SqlSysDB._insert_metadata.<locals>.<listcomp>)r   r   rE   rL   rM   r   rJ   r   isinra   r   rN   rO   rF   rG   r   r   r   r   items
isinstancer   rI   r   r   r   )r5   rQ   r   r   r>   r   r   r\   rS   rT   sql_idr   vr8   r8   r9   r     s   





json_strc                 C   sB  zt |}W n t jy   td| w zt|W S  ty } zw|d}|r|d}|d}|r|r||krdd | D }|	d|i t
|}td}	|  	|	|	jt| |	jt|k}
t|
|  \}}|  }||| W d    n1 sw   Y  |W  Y d }~S |d }~ww )Nz1Unable to decode configuration from JSON string: hnsw_configuration
batch_sizesync_thresholdc                 S   s   i | ]\}}|d vr||qS ))r   r   r8   )rg   r   r   r8   r8   r9   
<dictcomp>  s
    zCSqlSysDB._load_config_from_json_str_and_migrate.<locals>.<dictcomp>r`   )jsonloadsJSONDecodeErrorr   r   from_json_strr   getr   r   	from_jsonr   rE   r   r   r   r   rM   r>   r   rN   rD   rO   )r5   r   r   config_jsonr   r   r   r   r   r   r\   rS   rT   rQ   r8   r8   r9   r   k  sN   



	
z/SqlSysDB._load_config_from_json_str_and_migratec                 C   s   ddl m} td}||pi }t|}ttd|dgd}| }| 	 
||jt||jt|k}	t|	|  \}
}|  }||
| W d   |S 1 sYw   Y  |S )zuInsert the configuration from legacy metadata params into the collections table, and return the configuration object.r   )PersistentHnswParamsr`   r   )r?   r   )
parametersN)(chromadb.segment.impl.vector.hnsw_paramsr   r   extractr   from_legacy_paramsr   r   r   rE   r   r   r   r   rM   r>   r   rN   rD   rO   )r5   r   r   r   r   hnsw_metadata_paramsr   r   configuration_json_strr\   rS   rT   rQ   r8   r8   r9   r     s8   	

z*SqlSysDB._insert_config_from_legacy_params)r:   N)NNNr0   );__name__
__module____qualname__r"   __annotations__r   r2   r!   r    ALLr   r;   r   r
   r   rV   r$   r^   rd   r   r   r   rp   rt   r*   rv   r   r&   r   r   r   r   r'   r   r   r(   r   r)   r   r   r#   r   r   r   r   r+   r%   r,   r   r   r   r   r   r   r	   r   r   r   __classcell__r8   r8   r6   r9   r-   /   s  
 

3
	


_
Gr


)

0


=

W
2r-   )?r   loggingrj   typingr   r   r   r   r   r   r   r	   uuidr
   	overridesr   pypikar   r   	itertoolsr   chromadb.api.configurationr   r   r   r   chromadb.configr   r   r   chromadb.db.baser   r   r   r   chromadb.db.systemr   chromadb.errorsr   r   r    chromadb.telemetry.opentelemetryr   r   r    r!   chromadb.ingestr"   chromadb.typesr#   r$   r%   r&   r'   r(   r)   r*   r+   r,   	getLoggerr   r   r-   r8   r8   r8   r9   <module>   s$    (0
