bump version to 3.1.1

Signed-off-by: Andre Furlan <andre.furlan@databricks.com>

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Release History
 
+# 3.1.1 (2024-02-21)
+
+- Fix: Cloud fetch file download errors (#356)
+
 # 3.1.0 (2024-02-16)
 
 - Revert retry-after behavior to be exponential backoff (#349)

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "databricks-sql-connector"
-version = "3.1.0"
+version = "3.1.1"
 description = "Databricks SQL Connector for Python"
 authors = ["Databricks <databricks-sql-connector-maintainers@databricks.com>"]
 license = "Apache-2.0"

src/databricks/sql/__init__.py

@@ -94,7 +94,7 @@ def __repr__(self):
 DATE = DBAPITypeObject("date")
 ROWID = DBAPITypeObject()
 
-__version__ = "3.1.0"
+__version__ = "3.1.1"
 USER_AGENT_NAME = "PyDatabricksSqlConnector"
 
 # These two functions are pyhive legacy

src/databricks/sql/cloudfetch/downloader.py

@@ -10,7 +10,7 @@
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_CLOUD_FILE_TIMEOUT = int(os.getenv("DATABRICKS_CLOUD_FILE_TIMEOUT", 60))
+DEFAULT_CLOUD_FILE_TIMEOUT = int(os.getenv("DATABRICKS_CLOUD_FILE_TIMEOUT", 180))
 
 
 @dataclass
@@ -221,7 +221,7 @@ def http_get_with_retry(url, max_retries=5, backoff_factor=2, download_timeout=6
                 return response
             else:
                 logger.error(response)
-        except requests.RequestException as e:
+        except Exception as e:
             # if this is not redacted, it will print the pre-signed URL
             logger.error(f"request failed with exception: {re.sub(pattern, mask, str(e))}")
         finally:

src/databricks/sql/utils.py

@@ -236,6 +236,9 @@ def _create_next_table(self) -> Union[pyarrow.Table, None]:
         # The server rarely prepares the exact number of rows requested by the client in cloud fetch.
         # Subsequently, we drop the extraneous rows in the last file if more rows are retrieved than requested
         if arrow_table.num_rows > downloaded_file.row_count:
+            logger.debug(
+                f"received {arrow_table.num_rows} rows, expected {downloaded_file.row_count} rows. Dropping extraneous rows."
+            )
             self.start_row_index += downloaded_file.row_count
             return arrow_table.slice(0, downloaded_file.row_count)