Merge branch 'sqlalchemy-dev' of https://github.com/overcoil/fork-databricks-sql-python into sqlalchemy-dev

Author: overcoil

.github/.github/pull_request_template.md

@@ -0,0 +1,23 @@
+<!-- We welcome contributions. All patches must include a sign-off. Please see CONTRIBUTING.md for details -->
+
+
+## What type of PR is this? 
+<!-- Check all that apply, delete what doesn't apply. -->
+
+- [ ] Refactor
+- [ ] Feature
+- [ ] Bug Fix
+- [ ] Other
+
+## Description
+
+## How is this tested?
+
+- [ ] Unit tests 
+- [ ] E2E Tests 
+- [ ] Manually
+- [ ] N/A
+
+<!-- If Manually, please describe. -->
+
+## Related Tickets & Documents

.github/workflows/code-quality-checks.yml

@@ -1,7 +1,7 @@
 name: Code Quality Checks
 on: [push]
 jobs:
-  run-tests:
+  run-unit-tests:
     runs-on: ubuntu-latest
     steps:
       #----------------------------------------------
@@ -48,7 +48,7 @@ jobs:
       #              run test suite
       #----------------------------------------------
       - name: Run tests
-        run: poetry run pytest tests/
+        run: poetry run python -m pytest tests/unit
   check-linting:
     runs-on: ubuntu-latest
     steps:

.github/workflows/dco-check.yml

@@ -0,0 +1,25 @@
+name: DCO Check
+
+on: [pull_request]
+
+jobs:
+  check:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Check for DCO
+        id: dco-check
+        uses: tisonkun/[email protected]
+      - name: Comment about DCO status
+        uses: actions/github-script@v6
+        if: ${{ failure() }}
+        with:
+          script: |
+            github.rest.issues.createComment({
+              issue_number: context.issue.number,
+              owner: context.repo.owner,
+              repo: context.repo.repo,
+              body: `Thanks for your contribution! To satisfy the DCO policy in our \
+              [contributing guide](https://github.com/databricks/databricks-sql-python/blob/main/CONTRIBUTING.md) \
+              every commit message must include a sign-off message. One or more of your commits is missing this message. \
+              You can reword previous commit messages with an interactive rebase (\`git rebase -i main\`).`
+            })

CONTRIBUTING.md

@@ -1,6 +1,73 @@
-# Contributing
+# Contributing Guide
 
-To contribute to this repository, fork it and send pull requests.
+We happily welcome contributions to the `databricks-sql-connector` package. We use [GitHub Issues](https://github.com/databricks/databricks-sql-python/issues) to track community reported issues and [GitHub Pull Requests](https://github.com/databricks/databricks-sql-python/pulls) for accepting changes.
+
+Contributions are licensed on a license-in/license-out basis.
+
+## Communication
+Before starting work on a major feature, please reach out to us via GitHub, Slack, email, etc. We will make sure no one else is already working on it and ask you to open a GitHub issue.
+A "major feature" is defined as any change that is > 100 LOC altered (not including tests), or changes any user-facing behavior.
+We will use the GitHub issue to discuss the feature and come to agreement.
+This is to prevent your time being wasted, as well as ours.
+The GitHub review process for major features is also important so that organizations with commit access can come to agreement on design.
+If it is appropriate to write a design document, the document must be hosted either in the GitHub tracking issue, or linked to from the issue and hosted in a world-readable location.
+Specifically, if the goal is to add a new extension, please read the extension policy.
+Small patches and bug fixes don't need prior communication.
+
+## Coding Style
+We follow [PEP 8](https://www.python.org/dev/peps/pep-0008/) with one exception: lines can be up to 100 characters in length, not 79.
+
+## Sign your work
+The sign-off is a simple line at the end of the explanation for the patch. Your signature certifies that you wrote the patch or otherwise have the right to pass it on as an open-source patch. The rules are pretty simple: if you can certify the below (from developercertificate.org):
+
+```
+Developer Certificate of Origin
+Version 1.1
+
+Copyright (C) 2004, 2006 The Linux Foundation and its contributors.
+1 Letterman Drive
+Suite D4700
+San Francisco, CA, 94129
+
+Everyone is permitted to copy and distribute verbatim copies of this
+license document, but changing it is not allowed.
+
+
+Developer's Certificate of Origin 1.1
+
+By making a contribution to this project, I certify that:
+
+(a) The contribution was created in whole or in part by me and I
+    have the right to submit it under the open source license
+    indicated in the file; or
+
+(b) The contribution is based upon previous work that, to the best
+    of my knowledge, is covered under an appropriate open source
+    license and I have the right under that license to submit that
+    work with modifications, whether created in whole or in part
+    by me, under the same open source license (unless I am
+    permitted to submit under a different license), as indicated
+    in the file; or
+
+(c) The contribution was provided directly to me by some other
+    person who certified (a), (b) or (c) and I have not modified
+    it.
+
+(d) I understand and agree that this project and the contribution
+    are public and that a record of the contribution (including all
+    personal information I submit with it, including my sign-off) is
+    maintained indefinitely and may be redistributed consistent with
+    this project or the open source license(s) involved.
+```
+
+Then you just add a line to every git commit message:
+
+```
+Signed-off-by: Joe Smith <[email protected]>
+Use your real name (sorry, no pseudonyms or anonymous contributions.)
+```
+
+If you set your `user.name` and `user.email` git configs, you can sign your commit automatically with `git commit -s`.
 
 ## Set up your environment
 
@@ -9,40 +76,61 @@ This project uses [Poetry](https://python-poetry.org/) for dependency management
 1. Clone this respository
 2. Run `poetry install` 
 
-### Unit Tests
+### Run tests
+
+We use [Pytest](https://docs.pytest.org/en/7.1.x/) as our test runner. Invoke it with `poetry run python -m pytest`, all other arguments are passed directly to `pytest`.
 
-We use [Pytest](https://docs.pytest.org/en/7.1.x/) as our test runner. Invoke it with `poetry run pytest`, all other arguments are passed directly to `pytest`.
+#### Unit tests
+
+Unit tests do not require a Databricks account.
 
-#### All tests
 ```bash
-poetry run pytest tests
+poetry run python -m pytest tests/unit
 ```
-
 #### Only a specific test file
 
 ```bash
-poetry run pytest tests/tests.py
+poetry run python -m pytest tests/unit/tests.py
 ```
 
 #### Only a specific method
 
 ```bash
-poetry run pytest tests/tests.py::ClientTestSuite::test_closing_connection_closes_commands
+poetry run python -m pytest tests/unit/tests.py::ClientTestSuite::test_closing_connection_closes_commands
+```
+
+#### e2e Tests
+
+End-to-end tests require a Databricks account. Before you can run them, you must set connection details for a Databricks SQL endpoint in your environment:
+
+```bash
+export host=""
+export http_path=""
+export access_token=""
+```
+
+There are several e2e test suites available:
+- `PySQLCoreTestSuite`
+- `PySQLLargeQueriesSuite`
+- `PySQLRetryTestSuite.HTTP503Suite` **[not documented]**
+- `PySQLRetryTestSuite.HTTP429Suite` **[not documented]**
+- `PySQLUnityCatalogTestSuite` **[not documented]**
+
+To execute the core test suite:
+
+```bash
+poetry run python -m pytest tests/e2e/driver_tests.py::PySQLCoreTestSuite
 ```
 
+The suites marked `[not documented]` require additional configuration which will be documented at a later time.
 ### Code formatting
 
 This project uses [Black](https://pypi.org/project/black/).
 
 ```
-poetry run black src
+poetry run python3 -m black src --check
 ```
-## Pull Request Process
 
-1. Update the [CHANGELOG.md](README.md) or similar documentation with details of changes you wish to make, if applicable.
-2. Add any appropriate tests.
-3. Make your code or other changes.
-4. Review guidelines such as
-   [How to write the perfect pull request][github-perfect-pr], thanks!
+Remove the `--check` flag to write reformatted files to disk.
 
-[github-perfect-pr]: https://blog.github.com/2015-01-21-how-to-write-the-perfect-pull-request/
+To simplify reviews you can format your changes in a separate commit.

README.md

@@ -11,7 +11,7 @@ You are welcome to file an issue here for general use cases. You can also contac
 
 ## Requirements
 
-Python 3.7 or above is required.
+A development machine running Python >=3.7, <3.10.
 
 ## Documentation
 

pyproject.toml

@@ -18,6 +18,9 @@ PyHive = "^0.6.5"
 [tool.poetry.plugins."sqlalchemy.dialects"]
 "databricks.thrift" = "databricks.sqlalchemy:DatabricksDialect"
 
+[tool.poetry.plugins."sqlalchemy.dialects"]
+"databricks.thrift" = "databricks.sqlalchemy:DatabricksDialect"
+
 [tool.poetry.dev-dependencies]
 pytest = "^7.1.2"
 mypy = "^0.950"

src/databricks/sql/thrift_backend.py

@@ -1,4 +1,5 @@
 from decimal import Decimal
+import errno
 import logging
 import math
 import time
@@ -15,6 +16,9 @@
 
 from databricks.sql.thrift_api.TCLIService import TCLIService, ttypes
 from databricks.sql import *
+from databricks.sql.thrift_api.TCLIService.TCLIService import (
+    Client as TCLIServiceClient,
+)
 from databricks.sql.utils import (
     ArrowQueue,
     ExecuteResponse,
@@ -39,6 +43,7 @@
     "_retry_delay_max": (float, 60, 5, 3600),
     "_retry_stop_after_attempts_count": (int, 30, 1, 60),
     "_retry_stop_after_attempts_duration": (float, 900, 1, 86400),
+    "_retry_delay_default": (float, 5, 1, 60),
 }
 
 
@@ -71,6 +76,8 @@ def __init__(
         # _retry_delay_min                      (default: 1)
         # _retry_delay_max                      (default: 60)
         #   {min,max} pre-retry delay bounds
+        # _retry_delay_default                   (default: 5)
+        #   Only used when GetOperationStatus fails due to a TCP/OS Error.
         # _retry_stop_after_attempts_count      (default: 30)
         #   total max attempts during retry sequence
         # _retry_stop_after_attempts_duration   (default: 900)
@@ -158,7 +165,7 @@ def _initialize_retry_args(self, kwargs):
                 "retry parameter: {} given_or_default {}".format(key, given_or_default)
             )
             if bound != given_or_default:
-                logger.warn(
+                logger.warning(
                     "Override out of policy retry parameter: "
                     + "{} given {}, restricted to {}".format(
                         key, given_or_default, bound
@@ -243,7 +250,9 @@ def _handle_request_error(self, error_info, attempt, elapsed):
     # FUTURE: Consider moving to https://github.com/litl/backoff or
     # https://github.com/jd/tenacity for retry logic.
     def make_request(self, method, request):
-        """Execute given request, attempting retries when receiving HTTP 429/503.
+        """Execute given request, attempting retries when
+            1. Receiving HTTP 429/503 from server
+            2. OSError is raised during a GetOperationStatus
 
         For delay between attempts, honor the given Retry-After header, but with bounds.
         Use lower bound of expontial-backoff based on _retry_delay_min,
@@ -260,17 +269,21 @@ def make_request(self, method, request):
         def get_elapsed():
             return time.time() - t0
 
+        def bound_retry_delay(attempt, proposed_delay):
+            """bound delay (seconds) by [min_delay*1.5^(attempt-1), max_delay]"""
+            delay = int(proposed_delay)
+            delay = max(delay, self._retry_delay_min * math.pow(1.5, attempt - 1))
+            delay = min(delay, self._retry_delay_max)
+            return delay
+
         def extract_retry_delay(attempt):
             # encapsulate retry checks, returns None || delay-in-secs
             # Retry IFF 429/503 code + Retry-After header set
             http_code = getattr(self._transport, "code", None)
             retry_after = getattr(self._transport, "headers", {}).get("Retry-After")
             if http_code in [429, 503] and retry_after:
                 # bound delay (seconds) by [min_delay*1.5^(attempt-1), max_delay]
-                delay = int(retry_after)
-                delay = max(delay, self._retry_delay_min * math.pow(1.5, attempt - 1))
-                delay = min(delay, self._retry_delay_max)
-                return delay
+                return bound_retry_delay(attempt, int(retry_after))
             return None
 
         def attempt_request(attempt):
@@ -279,24 +292,57 @@ def attempt_request(attempt):
             # - non-None method_return -> success, return and be done
             # - non-None retry_delay -> sleep delay before retry
             # - error, error_message always set when available
+
+            error, error_message, retry_delay = None, None, None
             try:
                 logger.debug("Sending request: {}".format(request))
                 response = method(request)
                 logger.debug("Received response: {}".format(response))
                 return response
-            except Exception as error:
+            except OSError as err:
+                error = err
+                error_message = str(err)
+
+                gos_name = TCLIServiceClient.GetOperationStatus.__name__
+                if method.__name__ == gos_name:
+                    retry_delay = bound_retry_delay(attempt, self._retry_delay_default)
+
+                    # fmt: off
+                    # The built-in errno package encapsulates OSError codes, which are OS-specific.
+                    # log.info for errors we believe are not unusual or unexpected. log.warn for
+                    # for others like EEXIST, EBADF, ERANGE which are not expected in this context.
+                    #
+                    # I manually tested this retry behaviour using mitmweb and confirmed that 
+                    # GetOperationStatus requests are retried when I forced network connection
+                    # interruptions / timeouts / reconnects. See #24 for more info.
+                                            # | Debian | Darwin |
+                    info_errs = [           # |--------|--------|         
+                        errno.ESHUTDOWN,    # |   32   |   32   |
+                        errno.EAFNOSUPPORT, # |   97   |   47   |
+                        errno.ECONNRESET,   # |   104  |   54   |
+                        errno.ETIMEDOUT,    # |   110  |   60   |
+                    ]
+
+                    # fmt: on
+                    log_string = f"{gos_name} failed with code {err.errno} and will attempt to retry"
+                    if err.errno in info_errs:
+                        logger.info(log_string)
+                    else:
+                        logger.warning(log_string)
+            except Exception as err:
+                error = err
                 retry_delay = extract_retry_delay(attempt)
                 error_message = ThriftBackend._extract_error_message_from_headers(
                     getattr(self._transport, "headers", {})
                 )
-                return RequestErrorInfo(
-                    error=error,
-                    error_message=error_message,
-                    retry_delay=retry_delay,
-                    http_code=getattr(self._transport, "code", None),
-                    method=method.__name__,
-                    request=request,
-                )
+            return RequestErrorInfo(
+                error=error,
+                error_message=error_message,
+                retry_delay=retry_delay,
+                http_code=getattr(self._transport, "code", None),
+                method=method.__name__,
+                request=request,
+            )
 
         # The real work:
         # - for each available attempt: