Fix query timeout and improve error messages (#18)

.github/workflows/ci.yml

@@ -11,7 +11,7 @@ jobs:
  format:
  runs-on: ubuntu-latest
  steps:
- - uses: actions/checkout@v2
+ - uses: actions/checkout@v3
  - name: Update pip
  run: python3 -m pip install --no-cache --upgrade pip setuptools wheel
 
@@ -25,7 +25,7 @@ jobs:
  needs: format
  runs-on: ubuntu-latest
  steps:
- - uses: actions/checkout@v2
+ - uses: actions/checkout@v3
  - name: Update pip
  run: python3 -m pip install --no-cache --upgrade pip
 
@@ -68,7 +68,7 @@ jobs:
  - uses: actions/setup-python@v4
  with:
  python-version: ${{matrix.python}}
- - uses: actions/checkout@v2
+ - uses: actions/checkout@v3
  - uses: actions/download-artifact@master
  with:
  name: package
@@ -88,7 +88,7 @@ jobs:
  needs: build
  runs-on: ubuntu-latest
  steps:
- - uses: actions/checkout@v2
+ - uses: actions/checkout@v3
  - uses: actions/download-artifact@master
  with:
  name: package

CHANGELOG.md

@@ -7,17 +7,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Query timeout and improve error messages, [PR-18](https://github.com/panda-official/DriftCLI/pull/18)
+
 ## 0.9.0 - 2023-09-14
 
 ### Added
 
-- support for TypedData for CSV export, [PR-17](https://github.com/panda-official/DriftCLI/pull/17)
+- Support for TypedData for CSV export, [PR-17](https://github.com/panda-official/DriftCLI/pull/17)
 
 ## 0.8.0 - 2023-07-11
 
 ### Added
 
-- Supports for labels in metadata, [PR-15](https://github.com/panda-official/DriftCLI/pull/15)
+- Support for labels in metadata, [PR-15](https://github.com/panda-official/DriftCLI/pull/15)
 
 ## 0.7.1 - 2023-06-29
 

drift_cli/export_impl/raw.py

@@ -165,8 +165,6 @@ def _next():
  return next(it)
  except StopIteration:
  return None
- except DriftClientError:
- return None
 
  pkg = await asyncio.get_running_loop().run_in_executor(pool, _next)
  if pkg is None or pkg.meta.type == MetaInfo.TIME_SERIES:
@@ -336,7 +334,7 @@ async def export_raw(client: DriftClient, dest: str, parallel: int, **kwargs):
  progress,
  sem,
  topics=topics,
- parallel=parallel // len(topics) + 1,
+ parallel=min(parallel, len(topics)),
  **kwargs,
  )
  for topic in topics

drift_cli/utils/helpers.py

@@ -99,20 +99,22 @@ def stop_signal():
  "Signals are not supported on this platform. No graceful shutdown possible."
  )
 
- it = client.walk(topic, start=start, stop=stop)
+ it = client.walk(topic, start=start, stop=stop, ttl=180 * parallel)
 
  def _next():
  try:
  return next(it)
  except StopIteration:
  return None
- except DriftClientError:
- return None
 
  while True:
- drift_pkg = await loop.run_in_executor(pool, _next)
- if drift_pkg is None:
- break
+ try:
+ drift_pkg = await loop.run_in_executor(pool, _next)
+ if drift_pkg is None:
+ break
+ except DriftClientError as err:
+ progress.update(task, description=f"[ERROR] {err}", refresh=True)
+ return
 
  if signal_queue.qsize() > 0:
  # stop signal received

pyproject.toml

@@ -29,7 +29,7 @@ classifiers = [
 ]
 
 dependencies = [
- "drift-python-client~=0.8.1",
+ "drift-python-client~=0.9.0",
  "click~=8.1",
  "tomlkit~=0.11",
  "rich~=12.6",

tests/export_test.py

@@ -192,7 +192,8 @@ def test__export_raw_data(runner, client, conf, export_path, topics, timeseries)
  """Test export raw data"""
  client.walk.side_effect = [Iterator(timeseries), Iterator(timeseries)]
  result = runner(
- f"-c {conf} -p 2 export raw test {export_path} --start 2022-01-01 --stop 2022-01-02"
+ f"-c {conf} -p 2 export raw test {export_path} "
+ f"--start 2022-01-01T00:00:00Z --stop 2022-01-02T00:00:00Z"
  )
  assert f"Topic '{topics[0]}' (copied 2 packages (943 B)" in result.output
  assert f"Topic '{topics[1]}' (copied 2 packages (943 B)" in result.output
@@ -203,6 +204,12 @@ def test__export_raw_data(runner, client, conf, export_path, topics, timeseries)
  assert (export_path / topics[1] / "1.dp").exists()
  assert (export_path / topics[1] / "2.dp").exists()
 
+ assert client.walk.call_count == 2
+ assert client.walk.call_args_list[0][0][0] == topics[0]
+ assert client.walk.call_args_list[0][1]["start"] == 1640995200.0
+ assert client.walk.call_args_list[0][1]["stop"] == 1641081600.0
+ assert client.walk.call_args_list[0][1]["ttl"] == 360
+
 
 @pytest.mark.usefixtures("set_alias")
 def test__export_raw_data_with_metadata(