From a92acb2427ce35308fea6c2f1ef68f5b82d48770 Mon Sep 17 00:00:00 2001 From: matt-fleming Date: Wed, 31 Jan 2024 12:39:36 +0000 Subject: [PATCH] Rebase to upstream v3.0.2 --- .github/workflows/code-quality-checks.yml | 163 ---------------------- src/databricks/sql/thrift_backend.py | 16 ++- 2 files changed, 14 insertions(+), 165 deletions(-) delete mode 100644 .github/workflows/code-quality-checks.yml diff --git a/.github/workflows/code-quality-checks.yml b/.github/workflows/code-quality-checks.yml deleted file mode 100644 index bfb8ca94..00000000 --- a/.github/workflows/code-quality-checks.yml +++ /dev/null @@ -1,163 +0,0 @@ -name: Code Quality Checks -on: - push: - branches: - - main - pull_request: - branches: - - main -jobs: - run-unit-tests: - runs-on: ubuntu-latest - strategy: - matrix: - python-version: [3.8, 3.9, "3.10", "3.11"] - steps: - #---------------------------------------------- - # check-out repo and set-up python - #---------------------------------------------- - - name: Check out repository - uses: actions/checkout@v2 - - name: Set up python ${{ matrix.python-version }} - id: setup-python - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - #---------------------------------------------- - # ----- install & configure poetry ----- - #---------------------------------------------- - - name: Install Poetry - uses: snok/install-poetry@v1 - with: - virtualenvs-create: true - virtualenvs-in-project: true - installer-parallel: true - - #---------------------------------------------- - # load cached venv if cache exists - #---------------------------------------------- - - name: Load cached venv - id: cached-poetry-dependencies - uses: actions/cache@v2 - with: - path: .venv - key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ github.event.repository.name }}-${{ hashFiles('**/poetry.lock') }} - #---------------------------------------------- - # install dependencies if cache does not exist - #---------------------------------------------- - - name: Install dependencies - if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' - run: poetry install --no-interaction --no-root - #---------------------------------------------- - # install your root project, if required - #---------------------------------------------- - - name: Install library - run: poetry install --no-interaction - #---------------------------------------------- - # run test suite - #---------------------------------------------- - - name: Run tests - run: poetry run python -m pytest tests/unit - check-linting: - runs-on: ubuntu-latest - strategy: - matrix: - python-version: [3.8, 3.9, "3.10"] - steps: - #---------------------------------------------- - # check-out repo and set-up python - #---------------------------------------------- - - name: Check out repository - uses: actions/checkout@v2 - - name: Set up python ${{ matrix.python-version }} - id: setup-python - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - #---------------------------------------------- - # ----- install & configure poetry ----- - #---------------------------------------------- - - name: Install Poetry - uses: snok/install-poetry@v1 - with: - virtualenvs-create: true - virtualenvs-in-project: true - installer-parallel: true - - #---------------------------------------------- - # load cached venv if cache exists - #---------------------------------------------- - - name: Load cached venv - id: cached-poetry-dependencies - uses: actions/cache@v2 - with: - path: .venv - key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ github.event.repository.name }}-${{ hashFiles('**/poetry.lock') }} - #---------------------------------------------- - # install dependencies if cache does not exist - #---------------------------------------------- - - name: Install dependencies - if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' - run: poetry install --no-interaction --no-root - #---------------------------------------------- - # install your root project, if required - #---------------------------------------------- - - name: Install library - run: poetry install --no-interaction - #---------------------------------------------- - # black the code - #---------------------------------------------- - - name: Black - run: poetry run black --check src - - check-types: - runs-on: ubuntu-latest - strategy: - matrix: - python-version: [3.8, 3.9, "3.10"] - steps: - #---------------------------------------------- - # check-out repo and set-up python - #---------------------------------------------- - - name: Check out repository - uses: actions/checkout@v2 - - name: Set up python ${{ matrix.python-version }} - id: setup-python - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - #---------------------------------------------- - # ----- install & configure poetry ----- - #---------------------------------------------- - - name: Install Poetry - uses: snok/install-poetry@v1 - with: - virtualenvs-create: true - virtualenvs-in-project: true - installer-parallel: true - - #---------------------------------------------- - # load cached venv if cache exists - #---------------------------------------------- - - name: Load cached venv - id: cached-poetry-dependencies - uses: actions/cache@v2 - with: - path: .venv - key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ github.event.repository.name }}-${{ hashFiles('**/poetry.lock') }} - #---------------------------------------------- - # install dependencies if cache does not exist - #---------------------------------------------- - - name: Install dependencies - if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' - run: poetry install --no-interaction --no-root - #---------------------------------------------- - # install your root project, if required - #---------------------------------------------- - - name: Install library - run: poetry install --no-interaction - #---------------------------------------------- - # black the code - #---------------------------------------------- - - name: Mypy - run: poetry run mypy --install-types --non-interactive src diff --git a/src/databricks/sql/thrift_backend.py b/src/databricks/sql/thrift_backend.py index 69ac760a..e0247e6a 100644 --- a/src/databricks/sql/thrift_backend.py +++ b/src/databricks/sql/thrift_backend.py @@ -54,6 +54,10 @@ DATABRICKS_REASON_HEADER = "x-databricks-reason-phrase" TIMESTAMP_AS_STRING_CONFIG = "spark.thriftserver.arrowBasedRowSet.timestampAsString" + +# HACK! +THRIFT_SOCKET_TIMEOUT = os.getenv("THRIFT_SOCKET_TIMEOUT", None) + DEFAULT_SOCKET_TIMEOUT = float(900) # see Connection.__init__ for parameter descriptions. @@ -227,7 +231,12 @@ def __init__( **additional_transport_args, # type: ignore ) - timeout = kwargs.get("_socket_timeout", DEFAULT_SOCKET_TIMEOUT) + timeout = THRIFT_SOCKET_TIMEOUT or kwargs.get( + "_socket_timeout", DEFAULT_SOCKET_TIMEOUT + ) + # HACK! + logger.info(f"Setting timeout HACK! to {timeout}") + # setTimeout defaults to 15 minutes and is expected in ms self._transport.setTimeout(timeout and (float(timeout) * 1000.0)) @@ -637,7 +646,10 @@ def _create_arrow_table(self, t_row_set, lz4_compressed, schema_bytes, descripti num_rows, ) = convert_column_based_set_to_arrow_table(t_row_set.columns, description) elif t_row_set.arrowBatches is not None: - (arrow_table, num_rows,) = convert_arrow_based_set_to_arrow_table( + ( + arrow_table, + num_rows, + ) = convert_arrow_based_set_to_arrow_table( t_row_set.arrowBatches, lz4_compressed, schema_bytes ) else: