diff --git a/dclient/cli.py b/dclient/cli.py index ac660f4..f65ef6b 100644 --- a/dclient/cli.py +++ b/dclient/cli.py @@ -1,9 +1,13 @@ import click import httpx +import itertools import json import pathlib +from sqlite_utils.utils import rows_from_file, Format, TypeTracker from .utils import token_for_url +INSERT_BATCH_SIZE = 50 + def get_config_dir(): return pathlib.Path(click.get_app_dir("io.datasette.dclient")) @@ -16,10 +20,10 @@ def cli(): @cli.command() -@click.argument("url") +@click.argument("url_or_alias") @click.argument("sql") @click.option("--token", "-t", help="API token") -def query(url, sql, token): +def query(url_or_alias, sql, token): """ Run a SQL query against a Datasette database URL @@ -27,9 +31,11 @@ def query(url, sql, token): """ aliases_file = get_config_dir() / "aliases.json" aliases = _load_aliases(aliases_file) - if url in aliases: - url = aliases[url] - if not url.endswith(".json"): + if url_or_alias in aliases: + url = aliases[url_or_alias] + else: + url = url_or_alias + if not url_or_alias.endswith(".json"): url += ".json" if token is None: # Maybe there's a token in auth.json? @@ -77,6 +83,97 @@ def query(url, sql, token): click.echo(json.dumps(response.json()["rows"], indent=2)) +@cli.command() +@click.argument("url_or_alias") +@click.argument("table") +@click.argument("file", type=click.File("rb")) +@click.option("format_csv", "--csv", is_flag=True, help="Input is CSV") +@click.option("format_tsv", "--tsv", is_flag=True, help="Input is TSV") +@click.option("format_json", "--json", is_flag=True, help="Input is JSON") +@click.option("format_nl", "--nl", is_flag=True, help="Input is newline-delimited JSON") +@click.option( + "--no-detect-types", is_flag=True, help="Don't detect column types for CSV/TSV" +) +@click.option("--create", is_flag=True, help="Create table if it does not exist") +@click.option("--token", "-t", help="API token") +def insert( + url_or_alias, + table, + file, + format_csv, + format_tsv, + format_json, + format_nl, + no_detect_types, + create, + token, +): + """ + Insert data into a remote Datasette instance + + Example usage: + + \b + dclient insert \\ + https://private.datasette.cloud/data \\ + mytable data.csv + """ + aliases_file = get_config_dir() / "aliases.json" + aliases = _load_aliases(aliases_file) + if url_or_alias in aliases: + url = aliases[url_or_alias] + else: + url = url_or_alias + + if token is None: + token = token_for_url(url, _load_auths(get_config_dir() / "auth.json")) + + print(url, table) + format = None + if format_csv: + format = Format.CSV + elif format_tsv: + format = Format.CSV + elif format_json: + format = Format.JSON + elif format_nl: + format = Format.NL + if format is None and file.name == "": + raise click.ClickException( + "An explicit format is required - e.g. --csv " + "- when reading from standard input" + ) + rows, format = rows_from_file(file, format=format) + + first = True + + for batch in _batches(rows, INSERT_BATCH_SIZE): + types = None + if first and not no_detect_types: + # Detect types on first batch + tracker = TypeTracker() + list(tracker.wrap(batch)) + types = tracker.types + # Convert types + for row in batch: + for key, value in row.items(): + if value is None: + continue + if types[key] == "integer": + if not value: + row[key] = None + else: + row[key] = int(value) + elif types[key] == "float": + if not value: + row[key] = None + else: + row[key] = float(value) + first = False + response = _insert_batch(url, table, batch, token=token, create=create) + print(response) + + @cli.group() def alias(): "Manage aliases for different instances" @@ -194,3 +291,36 @@ def _load_auths(auth_file): else: auths = {} return auths + + +def _batches(iterable, size): + while True: + batch = list(itertools.islice(iterable, size)) + if not batch: + return + yield batch + + +def _insert_batch(url, table, batch, token, create): + if create: + data = { + "table": table, + "rows": batch, + } + url = "{}/-/create".format(url) + else: + data = { + "rows": batch, + } + url = "{}/{}/-/insert".format(url, table) + response = httpx.post( + url, + headers={ + "Authorization": "Bearer {}".format(token), + "Content-Type": "application/json", + }, + json=data, + timeout=40.0, + ) + response.raise_for_status() + return response.json() diff --git a/docs/queries.md b/docs/queries.md index fa3123e..99d5ff3 100644 --- a/docs/queries.md +++ b/docs/queries.md @@ -37,7 +37,7 @@ cog.out( ) ]]] --> ``` -Usage: dclient query [OPTIONS] URL SQL +Usage: dclient query [OPTIONS] URL_OR_ALIAS SQL Run a SQL query against a Datasette database URL diff --git a/setup.py b/setup.py index 7cc7600..78557ea 100644 --- a/setup.py +++ b/setup.py @@ -31,7 +31,7 @@ setup( "datasette": ["client = dclient.plugin"], "console_scripts": ["dclient = dclient.cli:cli"], }, - install_requires=["click", "httpx"], + install_requires=["click", "httpx", "sqlite-utils"], extras_require={ "test": [ "pytest",