Enhance SEA HTTP Client (databricks#618)

* allow empty schema bytes for alignment with SEA

Signed-off-by: varun-edachali-dbx <[email protected]>

* pass is_vl_op to Sea backend ExecuteResponse

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove catalog requirement in get_tables

Signed-off-by: varun-edachali-dbx <[email protected]>

* move filters.py to SEA utils

Signed-off-by: varun-edachali-dbx <[email protected]>

* ensure SeaResultSet

Signed-off-by: varun-edachali-dbx <[email protected]>

* prevent circular imports

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove unused imports

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove cast, throw error if not SeaResultSet

Signed-off-by: varun-edachali-dbx <[email protected]>

* pass param as TSparkParameterValue

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove failing test (temp)

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove SeaResultSet type assertion

Signed-off-by: varun-edachali-dbx <[email protected]>

* change errors to align with spec, instead of arbitrary ValueError

Signed-off-by: varun-edachali-dbx <[email protected]>

* make SEA backend methods return SeaResultSet

Signed-off-by: varun-edachali-dbx <[email protected]>

* use spec-aligned Exceptions in SEA backend

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove defensive row type check

Signed-off-by: varun-edachali-dbx <[email protected]>

* raise ProgrammingError for invalid id

Signed-off-by: varun-edachali-dbx <[email protected]>

* make is_volume_operation strict bool

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove complex types code

Signed-off-by: varun-edachali-dbx <[email protected]>

* Revert "remove complex types code"

This reverts commit 138359d.

* introduce type conversion for primitive types for JSON + INLINE

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove SEA running on metadata queries (known failures

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary docstrings

Signed-off-by: varun-edachali-dbx <[email protected]>

* align expected types with databricks sdk

Signed-off-by: varun-edachali-dbx <[email protected]>

* link rest api reference to validate types

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove test_catalogs_returns_arrow_table test

metadata commands not expected to pass

Signed-off-by: varun-edachali-dbx <[email protected]>

* fix fetchall_arrow and fetchmany_arrow

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove thrift aligned test_cancel_during_execute from SEA tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary changes in example scripts

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary chagnes in example scripts

Signed-off-by: varun-edachali-dbx <[email protected]>

* _convert_json_table -> _create_json_table

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove accidentally removed test

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove new unit tests (to be re-added based on new arch)

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove changes in sea_result_set functionality (to be re-added)

Signed-off-by: varun-edachali-dbx <[email protected]>

* introduce more integration tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove SEA tests in parameterized queries

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove partial parameter fix changes

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary timestamp tests

(pass with minor disparity)

Signed-off-by: varun-edachali-dbx <[email protected]>

* slightly stronger typing of _convert_json_types

Signed-off-by: varun-edachali-dbx <[email protected]>

* stronger typing of json utility func s

Signed-off-by: varun-edachali-dbx <[email protected]>

* stronger typing of fetch*_json

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove unused helper methods in SqlType

Signed-off-by: varun-edachali-dbx <[email protected]>

* line breaks after multi line pydocs, remove excess logs

Signed-off-by: varun-edachali-dbx <[email protected]>

* line breaks after multi line pydocs, reduce diff of redundant changes

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff of redundant changes

Signed-off-by: varun-edachali-dbx <[email protected]>

* mandate ResultData in SeaResultSet constructor

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove complex type conversion

Signed-off-by: varun-edachali-dbx <[email protected]>

* correct fetch*_arrow

Signed-off-by: varun-edachali-dbx <[email protected]>

* recover old sea tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* move queue and result set into SEA specific dir

Signed-off-by: varun-edachali-dbx <[email protected]>

* pass ssl_options into CloudFetchQueue

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove redundant conversion.py

Signed-off-by: varun-edachali-dbx <[email protected]>

* fix type issues

Signed-off-by: varun-edachali-dbx <[email protected]>

* ValueError not ProgrammingError

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff

Signed-off-by: varun-edachali-dbx <[email protected]>

* introduce SEA cloudfetch e2e tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* allow empty cloudfetch result

Signed-off-by: varun-edachali-dbx <[email protected]>

* add unit tests for CloudFetchQueue and SeaResultSet

Signed-off-by: varun-edachali-dbx <[email protected]>

* skip pyarrow dependent tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* simplify download process: no pre-fetching

Signed-off-by: varun-edachali-dbx <[email protected]>

* correct class name in logs

Signed-off-by: varun-edachali-dbx <[email protected]>

* align with old impl

Signed-off-by: varun-edachali-dbx <[email protected]>

* align next_n_rows with prev imple

Signed-off-by: varun-edachali-dbx <[email protected]>

* align remaining_rows with prev impl

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary Optional params

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary changes in thrift field if tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove unused imports

Signed-off-by: varun-edachali-dbx <[email protected]>

* init hybrid

* run large queries

Signed-off-by: varun-edachali-dbx <[email protected]>

* hybrid disposition

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-ncessary log

Signed-off-by: varun-edachali-dbx <[email protected]>

* formatting (black)

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove redundant tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* multi frame decompression of lz4

Signed-off-by: varun-edachali-dbx <[email protected]>

* ensure no compression (temp)

Signed-off-by: varun-edachali-dbx <[email protected]>

* introduce separate link fetcher

Signed-off-by: varun-edachali-dbx <[email protected]>

* log time to create table

Signed-off-by: varun-edachali-dbx <[email protected]>

* add chunk index to table creation time log

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove custom multi-frame decompressor for lz4

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove excess logs

* remove redundant tests (temp)

Signed-off-by: varun-edachali-dbx <[email protected]>

* add link to download manager before notifying consumer

Signed-off-by: varun-edachali-dbx <[email protected]>

* move link fetching immediately before table creation so link expiry is not an issue

Signed-off-by: varun-edachali-dbx <[email protected]>

* resolve merge artifacts

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove redundant methods

Signed-off-by: varun-edachali-dbx <[email protected]>

* formatting (black)

Signed-off-by: varun-edachali-dbx <[email protected]>

* introduce callback to handle link expiry

Signed-off-by: varun-edachali-dbx <[email protected]>

* fix types

Signed-off-by: varun-edachali-dbx <[email protected]>

* fix param type in unit tests

Signed-off-by: varun-edachali-dbx <[email protected]>

* formatting + minor type fixes

Signed-off-by: varun-edachali-dbx <[email protected]>

* Revert "introduce callback to handle link expiry"

This reverts commit bd51b1c.

* remove unused callback (to be introduced later)

Signed-off-by: varun-edachali-dbx <[email protected]>

* correct param extraction

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove common constructor for databricks client abc

Signed-off-by: varun-edachali-dbx <[email protected]>

* make SEA Http Client instance a private member

Signed-off-by: varun-edachali-dbx <[email protected]>

* make GetChunksResponse model more robust

Signed-off-by: varun-edachali-dbx <[email protected]>

* add link to doc of GetChunk response model

Signed-off-by: varun-edachali-dbx <[email protected]>

* pass result_data instead of "initial links" into SeaCloudFetchQueue

Signed-off-by: varun-edachali-dbx <[email protected]>

* move download_manager init into parent CloudFetchQueue

Signed-off-by: varun-edachali-dbx <[email protected]>

* raise ServerOperationError for no 0th chunk

Signed-off-by: varun-edachali-dbx <[email protected]>

* unused iports

Signed-off-by: varun-edachali-dbx <[email protected]>

* return None in case of empty respose

Signed-off-by: varun-edachali-dbx <[email protected]>

* ensure table is empty on no initial link s

Signed-off-by: varun-edachali-dbx <[email protected]>

* account for total chunk count

Signed-off-by: varun-edachali-dbx <[email protected]>

* iterate by chunk index instead of link

Signed-off-by: varun-edachali-dbx <[email protected]>

* make LinkFetcher convert link static

Signed-off-by: varun-edachali-dbx <[email protected]>

* add helper for link addition, check for edge case to prevent inf wait

Signed-off-by: varun-edachali-dbx <[email protected]>

* add unit tests for LinkFetcher

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary download manager check

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove un-necessary string literals around param type

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove duplicate download_manager init

Signed-off-by: varun-edachali-dbx <[email protected]>

* account for empty response in LinkFetcher init

Signed-off-by: varun-edachali-dbx <[email protected]>

* make get_chunk_link return mandatory ExternalLink

Signed-off-by: varun-edachali-dbx <[email protected]>

* set shutdown_event instead of breaking on completion so get_chunk_link is informed

Signed-off-by: varun-edachali-dbx <[email protected]>

* docstrings, logging, pydoc

Signed-off-by: varun-edachali-dbx <[email protected]>

* use total_chunk_cound > 0

Signed-off-by: varun-edachali-dbx <[email protected]>

* clarify that link has already been submitted on getting row_offset

Signed-off-by: varun-edachali-dbx <[email protected]>

* return None for out of range

Signed-off-by: varun-edachali-dbx <[email protected]>

* default link_fetcher to None

Signed-off-by: varun-edachali-dbx <[email protected]>

---------

Signed-off-by: varun-edachali-dbx <[email protected]>

* Chunk download latency (databricks#634)

* chunk download latency

Signed-off-by: Sai Shree Pradhan <[email protected]>

* formatting

Signed-off-by: Sai Shree Pradhan <[email protected]>

* test fixes

Signed-off-by: Sai Shree Pradhan <[email protected]>

* sea-migration static type checking fixes

Signed-off-by: Sai Shree Pradhan <[email protected]>

* check types fix

Signed-off-by: Sai Shree Pradhan <[email protected]>

* fix type issues

Signed-off-by: varun-edachali-dbx <[email protected]>

* type fix revert

Signed-off-by: Sai Shree Pradhan <[email protected]>

* -

Signed-off-by: Sai Shree Pradhan <[email protected]>

* statement id in get metadata functions

Signed-off-by: Sai Shree Pradhan <[email protected]>

* removed result set extractor

Signed-off-by: Sai Shree Pradhan <[email protected]>

* databricks client type

Signed-off-by: Sai Shree Pradhan <[email protected]>

* formatting

Signed-off-by: Sai Shree Pradhan <[email protected]>

* remove defaults, fix chunk id

Signed-off-by: Sai Shree Pradhan <[email protected]>

* added statement type to command id

Signed-off-by: Sai Shree Pradhan <[email protected]>

* check types fix

Signed-off-by: Sai Shree Pradhan <[email protected]>

* renamed chunk_id to num_downloaded_chunks

Signed-off-by: Sai Shree Pradhan <[email protected]>

* set statement type to query for chunk download

Signed-off-by: Sai Shree Pradhan <[email protected]>

* comment fix

Signed-off-by: Sai Shree Pradhan <[email protected]>

* removed dup check for trowset

Signed-off-by: Sai Shree Pradhan <[email protected]>

---------

Signed-off-by: Sai Shree Pradhan <[email protected]>

* acquire lock before notif + formatting (black)

Signed-off-by: varun-edachali-dbx <[email protected]>

* fix imports

Signed-off-by: varun-edachali-dbx <[email protected]>

* add get_chunk_link s

Signed-off-by: varun-edachali-dbx <[email protected]>

* simplify description extraction

Signed-off-by: varun-edachali-dbx <[email protected]>

* pass session_id_hex to ThriftResultSet

Signed-off-by: varun-edachali-dbx <[email protected]>

* revert to main's extract description

Signed-off-by: varun-edachali-dbx <[email protected]>

* validate row count for sync query tests as well

Signed-off-by: varun-edachali-dbx <[email protected]>

* guid_hex -> hex_guid

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff

Signed-off-by: varun-edachali-dbx <[email protected]>

* set .value in compression

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff

Signed-off-by: varun-edachali-dbx <[email protected]>

* formatting (black)

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove redundant test

Signed-off-by: varun-edachali-dbx <[email protected]>

* move extra_params to the back

Signed-off-by: varun-edachali-dbx <[email protected]>

* is_direct_results -> has_more_rows

Signed-off-by: varun-edachali-dbx <[email protected]>

* Revert "is_direct_results -> has_more_rows"

This reverts commit 0e87374.

* stop passing session_id_hex

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove redundant comment

Signed-off-by: varun-edachali-dbx <[email protected]>

* add extra_params param

Signed-off-by: varun-edachali-dbx <[email protected]>

* pass extra_params into test_...unset...

Signed-off-by: varun-edachali-dbx <[email protected]>

* remove excess session_id_he

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce changes in DatabricksRetryPolicy

Signed-off-by: varun-edachali-dbx <[email protected]>

* reduce diff in DatabricksRetryPolicy

Signed-off-by: varun-edachali-dbx <[email protected]>

* simple comments on proxy setting

Signed-off-by: varun-edachali-dbx <[email protected]>

* link docs for getproxies)(

Signed-off-by: varun-edachali-dbx <[email protected]>

* rename proxy specific attrs with proxy prefix

Signed-off-by: varun-edachali-dbx <[email protected]>

---------

Signed-off-by: varun-edachali-dbx <[email protected]>

Author: varun-edachali-dbx

src/databricks/sql/auth/retry.py

@@ -127,7 +127,7 @@ def __init__(
             total=_attempts_remaining,
             respect_retry_after_header=True,
             backoff_factor=self.delay_min,
-            allowed_methods=["POST"],
+            allowed_methods=["POST", "GET", "DELETE"],
             status_forcelist=[429, 503, *self.force_dangerous_codes],
         )