DRIFT-660: Implement export time series data as .dp and csv files (#2)

CHANGELOG.md

@@ -10,3 +10,5 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - `drift-cli alias` command to manage aliases, [PR-1](https://github.com/panda-official/DriftCLI/pull/1)
+- `drift-cli export raw` command to export data from blob
+ storage, [PR-2](https://github.com/panda-official/DriftCLI/pull/2)

README.md

@@ -8,8 +8,7 @@ Drift CLI is a command line client for [PANDA|Drift](https://driftpythonclient.r
 
 ## Features
 
-*
-
+* Export data from Drift Blob Storage
 
 ## Requirements
 
@@ -23,3 +22,13 @@ To install the Drift CLI, simply use pip:
 ```
 pip install drift-cli
 ```
+##
+
+```
+drift-cli --help
+drift-cli alias add drift-device --address 127.0.0.1 --password SOME_PASSWORD
+drift-cli export raw drift-device ./tmp --start 2021-01-01 --end 2021-01-02[export.md](..%2F..%2Freduct%2Freduct-cli%2Fdocs%2Fexport.md)[export.md](..%2F..%2Freduct%2Freduct-cli%2Fdocs%2Fexport.md)
+```
+## Links
+
+* [Documentation](https://driftcli.readthedocs.io/en/latest/)

docs/export.md

@@ -0,0 +1,53 @@
+# Export Data
+
+In this section, we will show how to export data from a Drift instance by using the `drift-cli export` command.
+
+## Export Raw Data
+
+
+The `drift-cli export raw` command allows you to export data from a Drift instance to a local folder
+on your computer. This can be useful if you want to make a copy of your data for backup or process them locally.
+
+The `drift-cli export raw` command has the following syntax:
+
+```
+drift-cli export raw [OPTIONS] SRC DEST
+```
+
+`SRC` should be an alias of a Drift instance you want to export data from.
+
+`DEST` should be the destination folder on your computer where you want to save the exported data.
+
+Here is an example of how you might use the `drift-cli export raw` command:
+
+```
+drift-cli export raw drift-device ./exported-data --start 2021-01-23 --end 2021-01-24
+```
+
+This will export all the raw data from the `drift-device` Drift instance to the `./exported-data` folder on your computer.
+For each topic the CLI will create a separate folder. Each package will be saved as a separate file with the name
+`<timestamp>.dp`.
+
+## Available options
+
+Here is a list of the options that you can use with the `drift-cli export` commands:
+
+* `--start`: This option allows you to specify a starting time point for the data that you want to export. Data with
+ timestamps newer than this time point will be included in the export. The time point should be in ISO format (e.g.,
+ 2022-01-01T00:00:00Z).
+
+* `--stop`: This option allows you to specify an ending time point for the data that you want to export. Data with
+ timestamps older than this time point will be included in the export. The time point should be in ISO format (e.g.,
+ 2022-01-01T00:00:00Z).
+
+* `--csv`: This option allows you to export data in csv format. It creates a separate csv file for each topic in the
+ exported data and save time series data in a single column with meta information in first row. The meta information
+ has the following format: `topic,package count, first timestamp, last timestamp`. The timestamp format is Unix time
+ in milliseconds.
+
+
+You also can use the global `--parallel` option to specify the number of entries that you want to export in parallel:
+
+```
+drift-cli --parallel 10 export raw drift-device ./exported-data --start 2021-01-01T00:00:00Z --end 2021-01-02T00:00:00Z
+```

drift_cli/cli.py

@@ -6,6 +6,8 @@
 import click
 
 from drift_cli.alias import alias
+from drift_cli.export import export
+
 from drift_cli.config import write_config, Config
 
 
@@ -17,12 +19,6 @@
  type=Path,
  help="Path to config file. Default ${HOME}/drift-cli/config.toml",
 )
-@click.option(
- "--timeout",
- "-t",
- type=int,
- help="Timeout for requests in seconds. Default 5",
-)
 @click.option(
  "--parallel",
  "-p",
@@ -33,25 +29,21 @@
 def cli(
  ctx,
  config: Optional[Path] = None,
- timeout: Optional[int] = None,
  parallel: Optional[int] = None,
 ):
  """CLI client for PANDA | Drift Platform"""
  if config is None:
  config = Path.home() / ".drift-cli" / "config.toml"
 
- if timeout is None:
- timeout = 5
-
  if parallel is None:
  parallel = 10
 
  if not Path.exists(config):
  write_config(config, Config(aliases={}))
 
  ctx.obj["config_path"] = config
- ctx.obj["timeout"] = timeout
  ctx.obj["parallel"] = parallel
 
 
 cli.add_command(alias, "alias")
+cli.add_command(export, "export")

drift_cli/export.py

@@ -0,0 +1,83 @@
+"""Export Command"""
+import asyncio
+
+import click
+from click import Abort
+from drift_client import DriftClient
+
+from drift_cli.config import Alias
+from drift_cli.config import read_config
+from drift_cli.export_impl.raw import export_raw
+from drift_cli.utils.consoles import error_console
+from drift_cli.utils.error import error_handle
+from drift_cli.utils.helpers import (
+ parse_path,
+)
+
+start_option = click.option(
+ "--start",
+ help="Export records with timestamps newer than this time point in ISO format e.g. 2023-01-01T00:00:00.000Z",
+)
+
+stop_option = click.option(
+ "--stop",
+ help="Export records with timestamps older than this time point in ISO format e.g 2023-01-01T00:00:00.000Z",
+)
+
+
+@click.group()
+def export():
+ """Export data from a bucket somewhere else"""
+
+
+@export.command()
+@click.argument("src")
+@click.argument("dest")
+@stop_option
+@start_option
+@click.option(
+ "--csv",
+ help="Export data as CSV instead of raw data (only for timeseries)",
+ default=False,
+ is_flag=True,
+)
+@click.pass_context
+def raw(
+ ctx,
+ src: str,
+ dest: str,
+ start: str,
+ stop: str,
+ csv: bool,
+): # pylint: disable=too-many-arguments
+ """Export data from SRC bucket to DST folder
+
+ SRC should be in the format of ALIAS/BUCKET_NAME.
+ DST should be a path to a folder.
+
+ As result, the folder will contain a folder for each entry in the bucket.
+ Each entry folder will contain a file for each record
+ in the entry with the timestamp as the name.
+ """
+ if start is None or stop is None:
+ error_console.print("Error: --start and --stop are required")
+ raise Abort()
+
+ alias_name, _ = parse_path(src)
+ alias: Alias = read_config(ctx.obj["config_path"]).aliases[alias_name]
+
+ loop = asyncio.get_event_loop()
+ run = loop.run_until_complete
+ client = DriftClient(alias.address, alias.password, loop=loop)
+
+ with error_handle():
+ run(
+ export_raw(
+ client,
+ dest,
+ parallel=ctx.obj["parallel"],
+ start=start,
+ stop=stop,
+ csv=csv,
+ )
+ )

drift_cli/export_impl/raw.py

@@ -0,0 +1,122 @@
+"""Export data"""
+import asyncio
+from concurrent.futures import ThreadPoolExecutor, Executor
+from pathlib import Path
+
+import numpy as np
+from drift_client import DriftClient
+from drift_protocol.meta import MetaInfo
+from rich.progress import Progress
+
+from drift_cli.utils.helpers import read_topic
+
+
+async def _export_topic(
+ pool: Executor,
+ client: DriftClient,
+ topic: str,
+ dest: str,
+ progress: Progress,
+ sem,
+ **kwargs,
+):
+ async for package, task in read_topic(pool, client, topic, progress, sem, **kwargs):
+ Path.mkdir(Path(dest) / topic, exist_ok=True, parents=True)
+ with open(Path(dest) / topic / f"{package.package_id}.dp", "wb") as file:
+ file.write(package.blob)
+
+
+async def _export_csv(
+ pool: Executor,
+ client: DriftClient,
+ curren_topic: str,
+ dest: str,
+ progress: Progress,
+ sem,
+ **kwargs,
+):
+ Path.mkdir(Path(dest), exist_ok=True, parents=True)
+
+ filename = Path(dest) / f"{curren_topic}.csv"
+ started = False
+ first_timestamp = 0
+ last_timestamp = 0
+ count = 0
+ async for package, task in read_topic(
+ pool, client, curren_topic, progress, sem, **kwargs
+ ):
+ meta = package.meta
+ if meta.type != MetaInfo.TIME_SERIES:
+ progress.update(
+ task,
+ description=f"[SKIPPED] Topic {curren_topic} is not a time series",
+ completed=True,
+ )
+ break
+
+ if not started:
+ with open(Path(dest) / f"{curren_topic}.csv", "w") as file:
+ file.write(" " * 256 + "\n")
+ started = True
+ first_timestamp = meta.time_series_info.start_timestamp.ToMilliseconds()
+ else:
+ if last_timestamp != meta.time_series_info.start_timestamp.ToMilliseconds():
+ progress.update(
+ task,
+ description=f"[ERROR] Topic {curren_topic} has gaps",
+ completed=True,
+ )
+ break
+
+ if package.status_code != 0:
+ progress.update(
+ task,
+ description=f"[ERROR] Topic {curren_topic} has a bad package",
+ completed=True,
+ )
+ break
+
+ last_timestamp = package.meta.time_series_info.stop_timestamp.ToMilliseconds()
+ with open(filename, "a") as file:
+ np.savetxt(file, package.as_np(), delimiter=",", fmt="%.5f")
+
+ count += 1
+
+ if started:
+ with open(filename, "r+") as file:
+ file.seek(0)
+ file.write(
+ ",".join(
+ [
+ curren_topic,
+ str(count),
+ str(first_timestamp),
+ str(last_timestamp),
+ ]
+ )
+ )
+
+
+async def export_raw(client: DriftClient, dest: str, parallel: int, **kwargs):
+ """Export data from Drift instance to DST folder
+ Args:
+ client: Drift client
+ dest: Path to a folder
+ parallel: Number of parallel tasks
+ KArgs:
+ start: Export records with timestamps newer than this time point in ISO format
+ stop: Export records with timestamps older than this time point in ISO format
+ csv: Export data as CSV instead of raw data
+ """
+ sem = asyncio.Semaphore(parallel)
+ with Progress() as progress:
+ with ThreadPoolExecutor() as pool:
+ topics = client.get_topics()
+ csv = kwargs.pop("csv", False)
+ task = _export_csv if csv else _export_topic
+
+ tasks = [
+ task(pool, client, topic, dest, progress, sem, topics=topics, **kwargs)
+ for topic in topics
+ ]
+ await asyncio.gather(*tasks)