id"dZddlmZddlZddlmZmZmZmZm Z m Z m Z m Z ddl Z ddlmZddlmZddl mZmZmZmZmZmZddlmZmZmZdd lmZmZdd lm Z dd l!m"Z"dd l#m$Z$d dZ%e%d!dZ&ddd"dZ'd#dZ(GddZ)dS)$%SQLAlchemy wrapper around a database.) annotationsN)AnyDictIterableListLiteralOptionalSequenceUnion) deprecated) get_from_env)MetaDataTable create_engineinspectselecttext)URLEngineResult)ProgrammingErrorSQLAlchemyError) CreateTable) Executable)NullTypereturnNonec ddl}t|drEt|jdr2|jj}t|dr|jd d}||_dSdSdSdS#t $rYdSwxYw) z?Patch DuckDBPyType to be hashable for SQLAlchemy compatibility.rNtyping DuckDBPyType__hash__rintc:tt|SN)hashstrselfs C:\Users\Dell Inspiron 16\Desktop\tws\AgrotaPowerBi\back-agrota-powerbi\mcp-client-agrota\venv\Lib\site-packages\langchain_community/utilities/sql_database.pyr"z%_patch_duckdb_types..__hash__&sD ??*)rr#)duckdbhasattrr r!r" ImportError)r, duckdb_typer"s r*_patch_duckdb_typesr0s   68 $ $ 0)O)O 0 -4K; 33 0{7K7S++++(0 $$$ 0 0 0 07S7S      sAA!! A/.A/index+sqlalchemy.engine.interfaces.ReflectedIndexr'c Vd|dd|ddt|dS)NzName: namez , Unique: uniquez , Columns: column_namesr')r1s r* _format_indexr83sF 2v 2 2%/ 2 2~.// 2 2r+z...)suffixcontentrlengthr#r9ct|tr|dkr|St||kr|S|d|t|z ddd|zS)z] Truncate a string to a certain number of words, based on the max string length. rN ) isinstancer'lenrsplit)r:r;r9s r* truncate_wordrB:so gs # #v{{ 7||v )Vc&kk)) * 1 1#q 9 9! K|]}|v|j|fVdSr%)_custom_table_info).0table intersectionr)s r* z'SQLDatabase.__init__..sG++L((/67((((++r+viewsbindonlyrCN)_engine_schemarGr _inspectorsetlistget_table_namesget_view_names _all_tables_include_tables_ignore_tablesget_usable_table_names_usable_tablesr?r# TypeError_sample_rows_in_table_info_indexes_in_table_infor`dictrc_max_string_length _view_supportr _metadatareflect)r)rNrCrPrRrTrUrVrXrZr[r\missing_tables usable_tablesrcs` @r*__init__zSQLDatabase.__init__Vs   Um USTT T!$,// 000?? @ @@LTt--V-<<222#%%   !1D4DDN  LnLLL5BLc-000suu   !043CCN  K^KKK3355 4AWc-000tGW3S99 LJKK K*C'&;#"3  " d5t<< 3 t677DDTEUVVL&*+++++!4+++''D # #4)!/XZZ$  N " ""\$-..| #       r+ database_uriUnion[str, URL] engine_argskwargsrrc 2|pi}|t|fi|fi|S)z'Construct a SQLAlchemy engine from URI.)r)clsrrr _engine_argss r*from_urizSQLDatabase.from_uris6#(b s=>>>>II&IIIr+z0.3.18zFor performing structured retrieval using Databricks SQL, see the latest best practices and recommended APIs at https://docs.unitycatalog.io/ai/integrations/langchain/ insteadz1.0)messageremovalcatalogr'host api_token warehouse_id cluster_idc  ddlm} n#t$rtdwxYwd} ddlm} | } | j} n#tt f$rd} YnwxYw|tdd| }| r| jnd} |tdd | }||| r| j }ntd |r|rtd |rd |}nd |}d|d|d|d|d| }|j d||d|S)a Class method to create an SQLDatabase instance from a Databricks connection. This method requires the 'databricks-sql-connector' package. If not installed, it can be added using `pip install databricks-sql-connector`. Args: catalog (str): The catalog name in the Databricks database. schema (str): The schema name in the catalog. host (Optional[str]): The Databricks workspace hostname, excluding 'https://' part. If not provided, it attempts to fetch from the environment variable 'DATABRICKS_HOST'. If still unavailable and if running in a Databricks notebook, it defaults to the current workspace hostname. Defaults to None. api_token (Optional[str]): The Databricks personal access token for accessing the Databricks SQL warehouse or the cluster. If not provided, it attempts to fetch from 'DATABRICKS_TOKEN'. If still unavailable and running in a Databricks notebook, a temporary token for the current user is generated. Defaults to None. warehouse_id (Optional[str]): The warehouse ID in the Databricks SQL. If provided, the method configures the connection to use this warehouse. Cannot be used with 'cluster_id'. Defaults to None. cluster_id (Optional[str]): The cluster ID in the Databricks Runtime. If provided, the method configures the connection to use this cluster. Cannot be used with 'warehouse_id'. If running in a Databricks notebook and both 'warehouse_id' and 'cluster_id' are None, it uses the ID of the cluster the notebook is attached to. Defaults to None. engine_args (Optional[dict]): The arguments to be used when connecting Databricks. Defaults to None. **kwargs (Any): Additional keyword arguments for the `from_uri` method. Returns: SQLDatabase: An instance of SQLDatabase configured with the provided Databricks connection details. Raises: ValueError: If 'databricks-sql-connector' is not found, or if both 'warehouse_id' and 'cluster_id' are provided, or if neither 'warehouse_id' nor 'cluster_id' are provided and it's not executing inside a Databricks notebook. r)sqlzfdatabricks-sql-connector package not found, please install with `pip install databricks-sql-connector`N) get_contextrDATABRICKS_HOSTrDATABRICKS_TOKENz6Need to provide either 'warehouse_id' or 'cluster_id'.z/Can't have both 'warehouse_id' or 'cluster_id'.z/sql/1.0/warehouses/z/sql/protocolv1/o/0/zdatabricks://token:@z ?http_path=z &catalog=z&schema=)rr) databricksrr.!dbruntime.databricks_repl_contextrbrowserHostNameAttributeErrorrapiToken clusterIdrGr)rrrCrrrrrrrcontextr default_hostdefault_api_token http_pathuris r*from_databrickszSQLDatabase.from_databrickssx  & & & & & & &   :    E E E E E E!kmmG"2LL^,   LLL  <(9<HHD07AG,,T  $[2DFWXXI  J$6 $.  L  PJ PNOO O  <=|==II;z;;I G) G Gd G G" G G-4 G G>D G G s|P+PPPPPs #AAA127.0.0.1:8902rootcnosdbpublicurluserpasswordtenantdatabasec ddlm}||||||}||S#t$rtdwxYw)a Class method to create an SQLDatabase instance from a CnosDB connection. This method requires the 'cnos-connector' package. If not installed, it can be added using `pip install cnos-connector`. Args: url (str): The HTTP connection host name and port number of the CnosDB service, excluding "http://" or "https://", with a default value of "127.0.0.1:8902". user (str): The username used to connect to the CnosDB service, with a default value of "root". password (str): The password of the user connecting to the CnosDB service, with a default value of "". tenant (str): The name of the tenant used to connect to the CnosDB service, with a default value of "cnosdb". database (str): The name of the database in the CnosDB tenant. Returns: SQLDatabase: An instance of SQLDatabase configured with the provided CnosDB connection details. r)make_cnosdb_langchain_uri)rzRcnos-connector package not found, please install with `pip install cnos-connector`)cnosdb_connectorrrr.)rrrrrrrrs r* from_cnosdbzSQLDatabase.from_cnosdbsw<  B B B B B B++CxRRC< |jD]:\} } t!| jt"ur|j| ;nq#t($rdt+|jD]:\} } t!| jt"ur|j| ;YnwxYwt-t/|j} | } |rL i} |jD]}|jr|j| |j <| r | d| dz } n#t6$rtd wxYwjpj}|r| d z } jr| d |d z } jr| d |d z } |r| d z } | | | d !|}|S)fGet information about specified tables. Follows best practices as specified in: Rajkumar et al, 2022 (https://arxiv.org/abs/2204.00498) If `sample_rows_in_table_info`, the specified number of sample rows will be appended to each table description. This can increase performance as demonstrated in the paper. Nz table_names r^cg|] }|j Srr4)ratbls r* z.SQLDatabase.get_table_info..ksQQQSQQQr+recg|]?}|jtvr'jdkr|jd=|@S)sqlitesqlite_)r4rlr startswith)rarall_table_namesr)s r*rz.SQLDatabase.get_table_info..us\   x3////\X--#(2E2Ei2P2P- .--r+z /* Column Comments: z */z:Column comments are available on PostgreSQL, MySQL, Oraclez /* z*/z )"rsrl differencerGr{ sorted_tablesr|rzrirmrjr`r4appendcolumnsitemstyper_columnsremoverrxr'rcompilerstripcomment Exceptionrwrv_get_table_indexes_get_sample_rowssortjoin)r)rrr}metadata_table_names to_reflect meta_tablestablesrbkv create_tablercolumn_comments_dictcolumnhas_extra_info final_strrs` @r*rzSQLDatabase.get_table_infoXs5577  " --88IIN X !V!V!V!VWWW)OQQDN4PQQQ))C0D,E,EE   N " "(\*%%| #        ~3    1 &1 &E& 5:9P+P+P d5ejABBB 1!M//1111DAqAF||x//--a0001" 1 1 1 //557711DAqAF||x//--a00011 1{51199$,GGHHL(//113J   +-("'-OO!>O@F0=+"R8LRRR !$T +Nt/N  'h& * FE4#:#:5#A#AEEEE . DC4#8#8#?#?CCCC  #d" MM* % % % % KK'' s AE44A+G"!G"1-II9rbrc|j|j}dt t |}d|S)NrzTable Indexes: )rk get_indexesr4rmapr8)r)rbindexesindexes_formatteds r*rzSQLDatabase._get_table_indexessD/--ej99 IIc-&A&ABB5"3555r+ct||j}dd|jD} |j5}||}ttd|}dddn #1swxYwYdd|D}n#t$rd}YnwxYw|jd|j d|d|S) N cg|] }|j Srr)racols r*rz0SQLDatabase._get_sample_rows..s C C Cc C C Cr+cd|DS)Nc<g|]}t|ddS)Ndr7)rais r*rzBSQLDatabase._get_sample_rows....s&#=#=#=QCFF4C4L#=#=#=r+r)lss r*z.SQLDatabase._get_sample_rows..s#=#="#=#=#=r+rc8g|]}d|S)r)rrarows r*rz0SQLDatabase._get_sample_rows..s"(O(O(OC3(O(O(Or+rz rows from z table: ) rlimitrvrrriconnectexecutermrrr4)r)rbcommand columns_str connectionsample_rows_result sample_rowssample_rows_strs r*rzSQLDatabase._get_sample_rowssk--%%d&EFFii C CU] C C CDD  !%%'' :%/%7%7%@%@""==?QRR                #ii(O(O;(O(O(OPPOO  ! ! ! OOO !. ! !5: ! ! ! ! ! ! s6 C&4B& C&B**C-B*."C C C all parametersexecution_optionsrUnion[str, Executable]fetchLiteral['all', 'one', 'cursor']rOptional[Dict[str, Any]]r'Union[Sequence[Dict[str, Any]], Result]c,|pi}|pi}|j5}|jN|jdkr |d|jf|n#|jdkr|d|jf|n|jdkrn|jdkr|d |jf|n|jd kr |d |j|n|jd kr |d |j|nm|jdkrna|jdkr|d|jf|n7|jdkr,|dt |j|t |trt|}n5t |trntdt|| |||}|j r|dkrd|D}nW|dkr.|}|gn|g}n#|dkr|cdddSt#d|cdddS dddn #1swxYwYgS)z Executes SQL command through underlying engine. If the statement returns no rows, an empty list is returned. N snowflakez"ALTER SESSION SET search_path = %s)rbigqueryzSET @@dataset_id=?mssqltrinozUSE ?r,zSET search_path TO oraclez#ALTER SESSION SET CURRENT_SCHEMA = sqlany postgresqlzSET search_path TO %shanaz SET SCHEMA z#Query expression has unknown type: rc6g|]}|Sr)_asdict)raxs r*rz(SQLDatabase._execute..'s EEEaaiikkEEEr+onecursorz8Fetch parameter must be either 'one', 'all', or 'cursor')ribeginrjrexec_driver_sqlrIr?r'rrrurr returns_rowsfetchallfetchonerrG) r)rrrrrrresult first_results r*_executezSQLDatabase._executes %2 -3 \   ! !L Z|'<;....<*;/ \Z//..,*;/ \W,,\W,,..*;/ \X-- ..>EE6??3D3DEEEFFe^^#)??#4#4L#/#7RRl>R>R>T>T=UFFh&&!OL L L L L L L L R%RYL L L L L L L L @ AL L L L L L L L L L L L L L L Z sH0J J  J J include_columns1Union[str, Sequence[Dict[str, Any]], Result[Any]]c||||}|dkr|Sfd|D}|s d|D}|sdSt|S)zExecute a SQL command and return a string representing the results. If the statement returns rows, a string of the results is returned. If the statement returns no rows, an empty string is returned. rrcPg|]"}fd|D#S)cDi|]\}}|t|jS))r;)rBry)rarvaluer)s r* z.SQLDatabase.run...Js?   !FE eD4KLLL   r+)r)rarr)s r*rz#SQLDatabase.run..IsT        %&WWYY      r+cPg|]#}t|$Sr)tuplevaluesrs r*rz#SQLDatabase.run..Rs(66635&&666r+r)rr')r)rrrrrrress` r*runzSQLDatabase.run4s UzEV   H  M         766#666C 2s88Or+ch ||S#t$r} d|cYd}~Sd}~wwxYw)rError: N)rrG)r)res r*get_table_info_no_throwz#SQLDatabase.get_table_info_no_throwYsX !&&{33 3 ! ! ! * Q==  !s 1,11Literal['all', 'one']cr ||||||S#t$r} d|cYd}~Sd}~wwxYw)a*Execute a SQL command and return a string representing the results. If the statement returns rows, a string of the results is returned. If the statement returns no rows, an empty string is returned. If the statement throws an error, the error message is returned. )rrrr,N)r*r)r)rrrrrr-s r* run_no_throwzSQLDatabase.run_no_throwisp !88%"3 /   ! ! ! * Q==  !s 6166Dict[str, Any]ct|}|}|d|dS)z4Return db context that you may want in agent prompt.z, )rr)rmrsr.r)r)rrs r*rzSQLDatabase.get_contextsF4668899 1133 (;9O9OPPPr+) NNNNrLFNFrMF)rNrrCrOrPrQrRrSrTrSrUr#rVrWrXrYrZrWr[r#r\rWr%)rrrrYrrrrK)NNNNN)rr'rCr'rrOrrOrrOrrOrrYrrrrK)rrrrr) rr'rr'rr'rr'rr'rrK)rr')rr)NF)rrSrrWrr')rbrrr')r) rrrrrrrrrr)rF) rrrrrrWrrrrrr)rrSrr') rr'rr/rrWrrrrrr)rr2)__name__ __module__ __qualname____doc__r classmethodrr rrpropertyrrsrnrrrrrr*r.r1rrr+r*rKrKSs// !%'+-1.2)*&+,0"!$&+NNNNN`'+JJJJ[JZ ##'&*$(&*^Q^Q^Q^Q[^Q@$ &&&&[&P)))X)>>>> Z%=uMMM---NM-%%%X% QVYYYYYv6666     B27\ 046: \\\\\\B27 % # 046:######J!!!!!&(- % ! 046:!!!!!!8QQQQQQr+rK)rr)r1r2rr')r:rr;r#r9r'rr')rCr'rr')*r7 __future__rrEr rrrrr r r r sqlalchemylangchain_core._apir langchain_core.utilsrrrrrrrsqlalchemy.enginerrrsqlalchemy.excrrsqlalchemy.schemarsqlalchemy.sql.expressionrsqlalchemy.typesrr0r8rBrIrKrr+r*rCs ++"""""" PPPPPPPPPPPPPPPPPPPP******------2111111111<<<<<<<<))))))000000%%%%%%     (?D F F F F F FvQvQvQvQvQvQvQvQvQvQr+