q-PhcddlmZddlZddlZddlZddlZddlmZmZm Z ddl m Z ddl m Z ddlmZmZmZmZer8ddlmZddlmZddlZdd lmZdd lmZmZdd lmZdd lmZmZdd l m!Z!ddl"m#Z#ej$e%5ddl&m'Z'e'j(eeeedddn #1swxYwYGddZ)GddZ*ddZ+dS)) annotationsN) TYPE_CHECKINGAnyLiteral)issue_unstable_warning)wrap_ldf) CatalogInfo ColumnInfo NamespaceInfo TableInfo) Generator)datetime) SchemaDict)DataSourceFormat TableType) DataFrame)CredentialProviderFunction CredentialProviderFunctionReturnCredentialProviderBuilder) LazyFrame)PyCatalogClient)catalog_info_clsnamespace_info_clstable_info_clscolumn_info_clsceZdZdZddddQd ZdRdZdSdZdTdZdUdZdVdZ dddddddWd*Z d+ddddd,dXd5Z ddd6dYd:Z d;d<dZd>Z ddd6d[d@Zd;d<d\dAZdddddBd]dJZd^dKZd_dOZed`dPZdS)aCatalogz Unity catalog client. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. autoT) bearer_token require_https workspace_urlstrr str | Noner!boolreturnNonectd|r*|dsd|d}t||dksk|dkrydtjvrk|drV|ddd d d r|}|dkrd }tj |||_ d S) a Initialize a catalog client. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- workspace_url URL of the workspace, or alternatively the URL of the Unity catalog API endpoint. bearer_token Bearer token to authenticate with. This can also be set to: * "auto": Automatically retrieve bearer tokens from the environment. * "databricks-sdk": Use the Databricks SDK to retrieve and use the bearer token from the environment. require_https Require the `workspace_url` to use HTTPS. z/`Catalog` functionality is considered unstable.zhttps://z%a non-HTTPS workspace_url was given (z5). To allow non-HTTPS URLs, pass require_https=False.zdatabricks-sdkrDATABRICKS_RUNTIME_VERSION/rz.cloud.databricks.comN) r startswith ValueErrorosenviron removeprefixsplitendswith_get_databricks_tokenrnew_client)selfr"r r!msgs [/var/www/html/test/jupyter/venv/lib/python3.11/site-packages/polars/catalog/unity/client.py__init__zCatalog.__init__6s8 PQQQ  "!9!9*!E!E "B BBB S// ! + + + F " "- ::((44;**:66sAq"122 ; 5577L 6 ! !L&*=,GG list[CatalogInfo]c4|jS)z List the available catalogs. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. )r5 list_catalogs)r6s r8r=zCatalog.list_catalogsos|))+++r: catalog_namelist[NamespaceInfo]c6|j|S)ac List the available namespaces (unity schema) under the specified catalog. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. )r5list_namespaces)r6r>s r8rAzCatalog.list_namespacesys|++L999r: namespacelist[TableInfo]c8|j||S)a List the available tables under the specified schema. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. namespace Name of the namespace (unity schema). )r5 list_tables)r6r>rBs r8rEzCatalog.list_tabless|'' i@@@r: table_namer c:|j|||S)a Retrieve the metadata of the specified table. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. namespace Name of the namespace (unity schema). table_name Name of the table. )r5get_table_infor6r>rBrFs r8rHzCatalog.get_table_infos&|**<JOOOr:table_idwrite1tuple[dict[str, str] | None, dict[str, str], int]c:|j||S)N)rJrK)r5get_table_credentials)r6rJrKs r8_get_table_credentialszCatalog._get_table_credentialss|11851QQQr:N)delta_table_versiondelta_table_optionsstorage_optionscredential_providerretriesrQint | str | datetime | NonerRdict[str, Any] | NonerSrT3CredentialProviderFunction | Literal['auto'] | NonerUintrc ||||} t| d\} } |||| dd\}}| dvrddlm} | | ||||S|d | } t | |d | } t | |r"t |}nd }t|j |||||| S) a Retrieve the metadata of the specified table. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. namespace Name of the namespace (unity schema). table_name Name of the table. delta_table_version Version of the table to scan (Deltalake only). delta_table_options Additional keyword arguments while reading a Deltalake table. storage_options Options that indicate how to connect to a cloud provider. The cloud providers currently supported are AWS, GCP, and Azure. See supported keys here: * `aws `_ * `gcp `_ * `azure `_ * Hugging Face (`hf://`): Accepts an API key under the `token` parameter: `{'token': '...'}`, or by setting the `HF_TOKEN` environment variable. If `storage_options` is not provided, Polars will try to infer the information from environment variables. credential_provider Provide a function that can be called to provide cloud storage credentials. The function is expected to return a dictionary of credential keys along with an optional credential expiry time. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. retries Number of retries if accessing a cloud instance fails. scan tableFzCatalog.scan_tablerK caller_nameDELTA DELTASHARINGr) scan_delta)versionrRrSrTNz3cannot apply delta_table_version for table of type z3cannot apply delta_table_options for table of type )rT cloud_optionsrU) rH!_extract_location_and_data_format_init_credentialspolars.io.deltarar-listitemsrr5 scan_table)r6r>rBrFrQrRrSrTrU table_infostorage_locationdata_source_formatrar7s r8rizCatalog.scan_tablesqv((y*MM /P  0 0 ,,04/E/E   , 0F0 0 ,_ !: : : 2 2 2 2 2 2: +$7 /$7    *(%(( S// !  *(%(( S// !  #"?#8#8#:#:;;OO#O L # #$7- $     r:error) delta_modedelta_write_optionsdelta_merge_optionsrSrTdfrrn:Literal['error', 'append', 'overwrite', 'ignore', 'merge']rorpdict[str, str] | None"None | deltalake.table.TableMergerc  ||||} t| d\} } || || dd\} }| dvr|| || |||Sd|d|d|d | d } t | ) a Write a DataFrame to a catalog table. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- df DataFrame to write. catalog_name Name of the catalog. namespace Name of the namespace (unity schema). table_name Name of the table. delta_mode : {'error', 'append', 'overwrite', 'ignore', 'merge'} (For delta tables) How to handle existing data. - If 'error', throw an error if the table already exists (default). - If 'append', will add new data. - If 'overwrite', will replace table with new data. - If 'ignore', will not write anything if table already exists. - If 'merge', return a `TableMerger` object to merge data from the DataFrame with the existing data. delta_write_options (For delta tables) Additional keyword arguments while writing a Delta lake Table. See a list of supported write options `here `__. delta_merge_options (For delta tables) Keyword arguments which are required to `MERGE` a Delta lake Table. See a list of supported merge options `here `__. storage_options Options that indicate how to connect to a cloud provider. The cloud providers currently supported are AWS, GCP, and Azure. See supported keys here: * `aws `_ * `gcp `_ * `azure `_ * Hugging Face (`hf://`): Accepts an API key under the `token` parameter: `{'token': '...'}`, or by setting the `HF_TOKEN` environment variable. If `storage_options` is not provided, Polars will try to infer the information from environment variables. credential_provider Provide a function that can be called to provide cloud storage credentials. The function is expected to return a dictionary of credential keys along with an optional credential expiry time. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. r[TzCatalog.write_tabler\r^)rSrTmoderorpzwrite_table: table format of .z (z) is unsupported.)rHrdre write_deltaNotImplementedError)r6rqr>rBrFrnrorprSrTrjrkrlr7s r8 write_tablezCatalog.write_table%sT((y*MM /P  0 0 ,,04/E/E   - 0F0 0 ,_ !: : :>> /$7$7$7 " :::"+::.8::&:::  &c** *r:)comment storage_rootr{r|r c<|j|||S)a Create a catalog. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. comment Leaves a comment about the catalog. storage_root Base location at which to store the catalog. )r>r{r|)r5create_catalog)r6r>r{r|s r8r~zCatalog.create_catalogs*.|**%w\+   r:F)forcerc>|j||dS)a_ Delete a catalog. Note that depending on the table type and catalog server, this may not delete the actual data files from storage. For more details, please consult the documentation of the catalog provider you are using. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. force Forcibly delete the catalog even if it is not empty. )r>rN)r5delete_catalog)r6r>rs r8rzCatalog.delete_catalogs%0 ##U#KKKKKr:r c>|j||||S)a Create a namespace (unity schema) in the catalog. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. namespace Name of the namespace (unity schema). comment Leaves a comment about the table. storage_root Base location at which to store the namespace. )r>rBr{r|)r5create_namespace)r6r>rBr{r|s r8rzCatalog.create_namespaces/4|,,%% -   r:c@|j|||dS)a Delete a namespace (unity schema) in the catalog. Note that depending on the table type and catalog server, this may not delete the actual data files from storage. For more details, please consult the documentation of the catalog provider you are using. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. namespace Name of the namespace (unity schema). force Forcibly delete the namespace even if it is not empty. )r>rBrN)r5delete_namespace)r6r>rBrs r8rzCatalog.delete_namespaces46 %%%% &     r:)rlr{r| propertiesschemaSchemaDict | None table_typerrlDataSourceFormat | Nonerc|j||||||||t| pi S)a Create a table in the catalog. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. namespace Name of the namespace (unity schema). table_name Name of the table. schema Schema of the table. table_type Type of the table data_source_format Storage format of the table. comment Leaves a comment about the table. storage_root Base location at which to store the table. properties Extra key-value metadata to store. ) r>rBrFrrrlr{r|r)r5 create_tablergrh) r6r>rBrFrrrlr{r|rs r8rzCatalog.create_tablesXR|((%!!1%Z-2446677)   r:c@|j|||dS)a% Delete the table stored at this location. Note that depending on the table type and catalog server, this may not delete the actual data files from storage. For more details, please consult the documentation of the catalog provider you are using. If you would like to perform manual deletions, the storage location of the files can be found using `get_table_info`. .. warning:: This functionality is considered **unstable**. It may be changed at any point without it being considered a breaking change. Parameters ---------- catalog_name Name of the catalog. namespace Name of the namespace (unity schema). table_name Name of the table. )r>rBrFN)r5 delete_tablerIs r8rzCatalog.delete_table8s6: !!%! "     r:rjr]>tuple[CredentialProviderBuilder | None, dict[str, Any] | None]c tddlm}|dkr|r|j||fSd|fStjddk}t ||j|} |} t| } | ri|pi| }| D]}  |r2|j } |j} d| d| d }t|tj |j||fS#t$rD}|r8|j } |j} d |d | d| d |d }t|tj Yd}~nd}~wwxYwddlm}|d|j|||fS)NrrrPOLARS_VERBOSE1rKz9auto-selected CatalogCredentialProvider for table_name = z (table_id = ))filez3error auto-initializing CatalogCredentialProvider: z table_name = z ) (write = )!_init_credential_provider_builder),polars.io.cloud.credential_provider._builderrfrom_initialized_providerr.getenvCatalogCredentialProviderrJ_credentials_iternextnameprintsysstderr Exceptionrrk)r6rTrSrjrKr]rverbosecatalog_credential_providervstorage_update_options_rFrJr7ers r8rezCatalog._init_credentials[sU       & ( (" -J0J'"##_,,),--4&? *%U' ' ' # +==??A%)!WW "% X"Wo&;"W@V"W    ,'_ %.7!77(0777c ++++F,F+ ' , , , ,'_ %.D!DD!DD(0DD8=DDDc ++++ ,6      10 J/+   s2C D:DDctjdd}t|tjdj}|d|ddddS)Nzdatabricks.sdkz?could not get Databricks token: databricks-sdk is not installedzdatabricks.sdk.coreDefaultCredentialsConfig Authorization) importlibutil find_spec ImportError import_module__dict__)clsr7ms r8r3zCatalog._get_databricks_tokens > # #$4 5 5 =SCc"" "  #$9 : : C7(&q%&((87799/J122NNr:)r"r#r r$r!r%r&r')r&r;)r>r#r&r?)r>r#rBr#r&rC)r>r#rBr#rFr#r&r )rJr#rKr%r&rL)r>r#rBr#rFr#rQrVrRrWrSrWrTrXrUrYr&r)rqrr>r#rBr#rFr#rnrrrorWrprWrSrsrTrXr&rt)r>r#r{r$r|r$r&r )r>r#rr%r&r') r>r#rBr#r{r$r|r$r&r )r>r#rBr#rr%r&r')r>r#rBr#rFr#rrrrrlrr{r$r|r$rrsr&r )r>r#rBr#rFr#r&r') rTrXrSrWrjr rKr%r]r#r&r)r&r#)__name__ __module__ __qualname____doc__r9r=rArErHrOrirzr~rrrrrre classmethodr3r:r8rr-sG$*" 7H7H7H7H7H7Hr,,,, : : : :AAAA"PPPP*RRRR<@5915 p p p p p p v 595915g+g+g+g+g+g+Z##'       > LLLLLL>##'       L       N7;"#',03 3 3 3 3 3 j! ! ! ! FIIIIVOOO[OOOr:rc*eZdZdZdd Zdd ZddZdS)rzGRetrieves credentials from the Unity catalog temporary credentials API.catalogrrJr#rKr%r&r'c0||_||_||_dSN)rrJrK)r6rrJrKs r8r9z"CatalogCredentialProvider.__init__s    r:rc>|\}\}}||fSr)r)r6rcredsexpirys r8__call__z"CatalogCredentialProvider.__call__s&!3355?E6f}r:Generator[Any]c#K|j|j|j\}}}|V|s|j}d|}t |||fVdS)NrzJdid not receive credentials from temporary credentials API for table_id = )rrOrJrKr)r6rrrrJr7s r8rz+CatalogCredentialProvider._credentials_iters15 0S0S M1T1 1 -%v%$$$ !}H!!! C.. Vmr:N)rrrJr#rKr%r&r')r&r)r&r)rrrrr9rrrr:r8rrsVQQ r:rrjr operationr#r&tuple[str, DataSourceFormat]c|jd|d}t||jd|d}t||j|jfS)Nzcannot z: no storage_location foundz: no data_source_format found)rkr-rl)rjrr7s r8rdrds_"*> >>>oo$,@ @@@oo  & (E EEr:)rjr rr#r&r), __future__r contextlibrr.rtypingrrrpolars._utils.unstablerpolars._utils.wraprpolars.catalog.unity.modelsr r r r collections.abcr r deltalakepolars._typingrrrpolars.dataframe.framerpolars.io.cloudrrrrpolars.lazyframersuppressr polars.polarsr init_classesrrrdrr:r8rs"""""" ..........999999''''''+))))))!!!!!!))))))GGGGGGGG000000WVVVVV******Z%%------ O $( " B OB OB OB OB OB OB OB OJ@ F F F F F FsB**B.1B.