Merge pull request #1 from panda-official/aliases

Add alias cmd

CHANGELOG.md

@@ -6,3 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+
+### Added
+
+- `drift-cli alias` command to manage aliases, [PR-1](https://github.com/panda-official/DriftCLI/pull/1)

README.md

@@ -8,7 +8,7 @@ Drift CLI is a command line client for [PANDA|Drift](https://driftpythonclient.r
 
 ## Features
 
-* ....
+*
 
 
 ## Requirements

docs/aliases.md

@@ -0,0 +1,36 @@
+# Aliases
+
+Aliases are used to simplify communication with a Drift instance.
+This way, users don't have to type the IP address or hostname and credentials for the Drift instance each time they want
+to use it.
+
+To create an alias, use the following `drift-cli` command:
+
+```shell
+drift-cli alias add local
+```
+
+Alternatively, you can provide the URL and API token for the storage engine as options:
+
+```shell
+drift-cli alias add -a 127.0.0.1 -p password -b bucket local
+```
+
+## Browsing aliases
+
+Once you've created an alias, you can use the `drift-cli` alias command to view it in a list or check its credentials:
+
+```shell
+
+```shell
+drift-cli alias ls
+drift-cli alias show local
+```
+
+## Removing an alias
+
+To remove an alias, use the rm subcommand of `drift-cli` alias:
+
+```shell
+drift-cli alias rm play
+```

drift_cli/alias.py

@@ -0,0 +1,85 @@
+"""Alias commands"""
+from typing import Optional
+
+import click
+from click import Abort
+
+from drift_cli.config import Config, read_config, write_config, Alias
+from drift_cli.utils.consoles import console, error_console
+from drift_cli.utils.error import error_handle
+from drift_cli.utils.helpers import get_alias
+
+
+@click.group()
+@click.pass_context
+def alias(ctx):
+ """Commands to manage aliases"""
+ ctx.obj["conf"] = read_config(ctx.obj["config_path"])
+
+
+@alias.command()
+@click.pass_context
+def ls(ctx):
+ """Print list of aliases"""
+ for name, _ in ctx.obj["conf"].aliases.items():
+ console.print(name)
+
+
+@alias.command()
+@click.argument("name")
+@click.pass_context
+def show(ctx, name: str):
+ """Show alias configuration"""
+ alias_: Alias = get_alias(ctx.obj["config_path"], name)
+ console.print(f"[bold]Address[/bold]: {alias_.address}")
+ console.print(f"[bold]Bucket[/bold] : {alias_.bucket}")
+
+
+@alias.command()
+@click.argument("name")
+@click.option(
+ "--address", "-A", help="Address of Drift instance, can be IP or hostname"
+)
+@click.option("--password", "-p", help="Password for Drift instance")
+@click.option("--bucket", "-b", help="Bucket to use, defaults to 'data'")
+@click.pass_context
+def add(
+ ctx,
+ name: str,
+ address: Optional[str],
+ password: Optional[str],
+ bucket: Optional[str],
+):
+ """Add a new alias with NAME"""
+ conf: Config = ctx.obj["conf"]
+ if name in conf.aliases:
+ error_console.print(f"Alias '{name}' already exists")
+ raise Abort()
+
+ if address is None or len(address) == 0:
+ address = click.prompt("IP or Hostname", type=str)
+ if password is None:
+ password = click.prompt("Password", type=str, hide_input=True)
+ if bucket is None:
+ bucket = click.prompt("Bucket", type=str, default="data")
+
+ with error_handle():
+ entry = Alias(address=address, password=password, bucket=bucket)
+
+ conf.aliases[name] = entry
+ write_config(ctx.obj["config_path"], conf)
+
+
+@alias.command()
+@click.argument("name")
+@click.pass_context
+def rm(ctx, name: str):
+ """
+ Remove alias with NAME
+ """
+ # Check if name exists
+ conf: Config = ctx.obj["conf"]
+ _ = get_alias(ctx.obj["config_path"], name)
+
+ conf.aliases.pop(name)
+ write_config(ctx.obj["config_path"], conf)

drift_cli/cli.py

@@ -5,7 +5,8 @@
 
 import click
 
-from drift_cli.config import write_config
+from drift_cli.alias import alias
+from drift_cli.config import write_config, Config
 
 
 @click.group()
@@ -46,8 +47,11 @@ def cli(
  parallel = 10
 
  if not Path.exists(config):
- write_config(config, {"aliases": {}})
+ write_config(config, Config(aliases={}))
 
  ctx.obj["config_path"] = config
  ctx.obj["timeout"] = timeout
  ctx.obj["parallel"] = parallel
+
+
+cli.add_command(alias, "alias")

drift_cli/config.py

@@ -26,10 +26,10 @@ def write_config(path: Path, config: Config):
  if not Path.exists(path):
  os.makedirs(path.parent, exist_ok=True)
  with open(path, "w", encoding="utf8") as config_file:
- config_file.write(toml.dumps(config))
+ toml.dump(config.dict(), config_file)
 
 
 def read_config(path: Path) -> Config:
  """Read config from TOML file"""
  with open(path, "r", encoding="utf8") as config_file:
- return toml.loads(config_file.read())
+ return Config.parse_obj(toml.load(config_file))

drift_cli/utils/consoles.py

@@ -0,0 +1,5 @@
+"""Rich consoles"""
+from rich.console import Console
+
+error_console = Console(stderr=True, style="bold red")
+console = Console()

drift_cli/utils/error.py

@@ -0,0 +1,16 @@
+"""Error handling"""
+from contextlib import contextmanager
+
+from click import Abort
+
+from drift_cli.utils.consoles import error_console
+
+
+@contextmanager
+def error_handle():
+ """Wrap try-catch block and print errorr"""
+ try:
+ yield
+ except Exception as err:
+ error_console.print(f"[{type(err).__name__}] {err}")
+ raise Abort() from err

drift_cli/utils/helpers.py

@@ -0,0 +1,153 @@
+"""Helper functions"""
+import asyncio
+import signal
+import time
+from asyncio import Semaphore, Queue
+from datetime import datetime
+from pathlib import Path
+from typing import Tuple, List
+
+from click import Abort
+from reduct import EntryInfo, Bucket
+from rich.progress import Progress
+
+from drift_cli.config import read_config, Alias
+from drift_cli.utils.consoles import error_console
+from drift_cli.utils.humanize import pretty_size
+
+signal_queue = Queue()
+
+
+def get_alias(config_path: Path, name: str) -> Alias:
+ """Helper method to parse alias from config"""
+ conf = read_config(config_path)
+
+ if name not in conf.aliases:
+ error_console.print(f"Alias '{name}' doesn't exist")
+ raise Abort()
+ alias_: Alias = conf.aliases[name]
+ return alias_
+
+
+def parse_path(path) -> Tuple[str, str]:
+ """Parse path ALIAS/RESOURCE"""
+ args = path.split("/")
+ if len(args) != 2:
+ raise RuntimeError(
+ f"Path {path} has wrong format. It must be 'ALIAS/BUCKET_NAME'"
+ )
+ return tuple(args)
+
+
+async def read_records_with_progress(
+ entry: EntryInfo,
+ bucket: Bucket,
+ progress: Progress,
+ sem: Semaphore,
+ **kwargs,
+): # pylint: disable=too-many-locals
+ """Read records from entry and show progress
+ Args:
+ entry (EntryInfo): Entry to read records from
+ bucket (Bucket): Bucket to read records from
+ progress (Progress): Progress bar to show progress
+ sem (Semaphore): Semaphore to limit parallelism
+ Keyword Args:
+ start (Optional[datetime]): Start time point
+ stop (Optional[datetime]): Stop time point
+ timeout (int): Timeout for read operation
+ Yields:
+ Record: Record from entry
+ """
+
+ def _to_timestamp(date: str) -> int:
+ return int(
+ datetime.fromisoformat(date.replace("Z", "+00:00")).timestamp() * 1000_000
+ )
+
+ start = _to_timestamp(kwargs["start"]) if kwargs["start"] else entry.oldest_record
+ stop = _to_timestamp(kwargs["stop"]) if kwargs["stop"] else entry.latest_record
+
+ include = {}
+ for item in kwargs["include"]:
+ if item:
+ key, value = item.split("=")
+ include[key] = value
+
+ exclude = {}
+ for item in kwargs["exclude"]:
+ if item:
+ key, value = item.split("=")
+ exclude[key] = value
+
+ last_time = start
+ task = progress.add_task(f"Entry '{entry.name}' waiting", total=stop - start)
+ async with sem:
+ exported_size = 0
+ count = 0
+ stats = []
+ speed = 0
+
+ def stop_signal():
+ signal_queue.put_nowait("stop")
+
+ asyncio.get_event_loop().add_signal_handler(signal.SIGINT, stop_signal)
+ asyncio.get_event_loop().add_signal_handler(signal.SIGTERM, stop_signal)
+
+ async for record in bucket.query(
+ entry.name,
+ start=start,
+ stop=stop,
+ include=include,
+ exclude=exclude,
+ ttl=kwargs["timeout"] * sem._value, # pylint: disable=protected-access
+ ):
+ if signal_queue.qsize() > 0:
+ # stop signal received
+ progress.update(
+ task,
+ description=f"Entry '{entry.name}' "
+ f"(copied {count} records ({pretty_size(exported_size)}), stopped",
+ refresh=True,
+ )
+ return
+
+ exported_size += record.size
+ stats.append((record.size, time.time()))
+ if len(stats) > 10:
+ stats.pop(0)
+
+ if len(stats) > 1:
+ speed = sum(s[0] for s in stats) / (stats[-1][1] - stats[0][1])
+
+ yield record
+
+ progress.update(
+ task,
+ description=f"Entry '{entry.name}' "
+ f"(copied {count} records ({pretty_size(exported_size)}), "
+ f"speed {pretty_size(speed)}/s)",
+ advance=record.timestamp - last_time,
+ refresh=True,
+ )
+ last_time = record.timestamp
+ count += 1
+
+ progress.update(task, total=1, completed=True)
+
+
+def filter_entries(entries: List[EntryInfo], names: List[str]) -> List[EntryInfo]:
+ """Filter entries by names"""
+ if not names or len(names) == 0:
+ return entries
+
+ if len(names) == 1 and names[0] == "":
+ return entries
+
+ def _filter(entry):
+ for name in names:
+ if name == entry.name:
+ return True
+ return False
+
+ return list(filter(_filter, entries))