diff --git a/README.md b/README.md index 77996fb..511095c 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,13 @@ A client CLI utility for [Datasette](https://datasette.io/) instances +## Things you can do with dclient + +- Run SQL queries against Datasette and returning the results as JSON +- Run queries against authenticated Datasette instances +- Create aliases and store authentication tokens for convenient access to Datasette +- Insert data into Datasette using the [insert API](https://docs.datasette.io/en/latest/json_api.html#the-json-write-api) (Datasette 1.0 alpha or higher) + ## Installation Install this tool using `pip`: diff --git a/dclient/cli.py b/dclient/cli.py index ac660f4..49ff338 100644 --- a/dclient/cli.py +++ b/dclient/cli.py @@ -1,7 +1,11 @@ import click import httpx +import io +import itertools import json import pathlib +from sqlite_utils.utils import rows_from_file, Format, TypeTracker, progressbar +import sys from .utils import token_for_url @@ -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,159 @@ 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( + "filepath", type=click.Path("rb", readable=True, allow_dash=True, dir_okay=False) +) +@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("--encoding", help="Character encoding for CSV/TSV") +@click.option( + "--no-detect-types", is_flag=True, help="Don't detect column types for CSV/TSV" +) +@click.option( + "--replace", is_flag=True, help="Replace rows with a matching primary key" +) +@click.option("--ignore", is_flag=True, help="Ignore rows with a matching primary key") +@click.option("--create", is_flag=True, help="Create table if it does not exist") +@click.option( + "pks", + "--pk", + multiple=True, + help="Columns to use as the primary key when creating the table", +) +@click.option( + "--batch-size", type=int, default=100, help="Send rows in batches of this size" +) +@click.option("--token", "-t", help="API token") +@click.option("--silent", is_flag=True, help="Don't output progress") +def insert( + url_or_alias, + table, + filepath, + format_csv, + format_tsv, + format_json, + format_nl, + encoding, + no_detect_types, + replace, + ignore, + create, + pks, + batch_size, + token, + silent, +): + """ + Insert data into a remote Datasette instance + + Example usage: + + \b + dclient insert \\ + https://private.datasette.cloud/data \\ + mytable data.csv --pk id --create + """ + 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")) + + format = None + if format_csv: + format = Format.CSV + elif format_tsv: + format = Format.TSV + elif format_json: + format = Format.JSON + elif format_nl: + format = Format.NL + if format is None and filepath == "-": + raise click.ClickException( + "An explicit format is required - e.g. --csv " + "- when reading from standard input" + ) + + if filepath != "-": + file_size = pathlib.Path(filepath).stat().st_size + fp = open(filepath, "rb") + else: + fp = sys.stdin.buffer + file_size = None + + try: + rows, format = rows_from_file(fp, format=format, encoding=encoding) + except Exception as ex: + raise click.ClickException(str(ex)) + + if format in (Format.JSON, Format.NL): + # Disable progress bar - it can't handle these formats + file_size = None + no_detect_types = True + + first = True + + with progressbar( + length=file_size, + label="Inserting rows", + silent=silent or (file_size is None), + show_percent=True, + ) as bar: + bytes_so_far = 0 + for batch in _batches(rows, batch_size): + if file_size is not None: + try: + bytes_consumed_so_far = fp.tell() + new_bytes = bytes_consumed_so_far - bytes_so_far + bar.update(new_bytes) + bytes_so_far += new_bytes + except ValueError: + # File has likely been closed, so fp.tell() fails + pass + 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 + _insert_batch( + url=url, + table=table, + batch=batch, + token=token, + create=create, + pks=pks, + replace=replace, + ignore=ignore, + ) + + @cli.group() def alias(): "Manage aliases for different instances" @@ -194,3 +353,56 @@ def _load_auths(auth_file): else: auths = {} return auths + + +def _batches(iterable, size): + iterable = iter(iterable) + while True: + batch = list(itertools.islice(iterable, size)) + if not batch: + return + yield batch + + +def _insert_batch(*, url, table, batch, token, create, pks, replace, ignore): + if create: + data = { + "table": table, + "rows": batch, + } + if replace: + data["replace"] = True + if ignore: + data["ignore"] = True + if pks: + if len(pks) == 1: + data["pk"] = pks[0] + else: + data["pks"] = pks + url = "{}/-/create".format(url) + else: + data = { + "rows": batch, + } + if replace: + data["replace"] = True + if ignore: + data["ignore"] = True + url = "{}/{}/-/insert".format(url, table) + response = httpx.post( + url, + headers={ + "Authorization": "Bearer {}".format(token), + "Content-Type": "application/json", + }, + json=data, + timeout=40.0, + ) + if str(response.status_code)[0] != "2": + # Is there an error we can show? + if "/json" in response.headers["content-type"]: + data = response.json() + if "errors" in data: + raise click.ClickException("\n".join(data["errors"])) + response.raise_for_status() + return response.json() diff --git a/docs/authentication.md b/docs/authentication.md index 5f0bd18..40e2dd7 100644 --- a/docs/authentication.md +++ b/docs/authentication.md @@ -1,3 +1,5 @@ +(authentication)= + # Authentication `dclient` can handle API tokens for Datasette instances that require authentication. diff --git a/docs/index.md b/docs/index.md index 90f9eba..b330d2e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -41,4 +41,5 @@ maxdepth: 3 queries aliases authentication +inserting ``` diff --git a/docs/inserting.md b/docs/inserting.md new file mode 100644 index 0000000..e664a29 --- /dev/null +++ b/docs/inserting.md @@ -0,0 +1,71 @@ +# Inserting data + +The `dclient insert` command can be used to insert data from a local file directly into a Datasette instance, via the [Write API](https://docs.datasette.io/en/latest/json_api.html#the-json-write-api) introduced in the Datasette 1.0 alphas. + +First you'll need to {ref}`authenticate ` with the instance. + +To insert data from a `data.csv` file into a table called `my_table`, creating that table if it does not exist: + +```bash +dclient insert https://my-private-space.datasette.cloud/data my_table data.csv --create +``` +You can also pipe data into standard input: +```bash +curl -s 'https://api.github.com/repos/simonw/dclient/issues' | \ + dclient insert \ + https://my-private-space.datasette.cloud/data \ + issues - --create +``` + +## Supported formats + +Data can be inserted from CSV, TSV, JSON or newline-delimited JSON files. + +The format of the file will be automatically detected. You can override this by using one of the following options: + +- `--csv` +- `--tsv` +- `--json` +- `--nl` for newline-delimited JSON + +Use `--encoding ` to specify the encoding of the file. The default is `utf-8`. + +### JSON + +JSON files should be formatted like this: +```json +[ + { + "id": 1 + "column1": "value1", + "column2": "value2" + }, + { + "id": 2 + "column1": "value1", + "column2": "value2" + } +] +``` +Newline-delimited files like this: +``` +{"id": 1, "column1": "value1", "column2": "value2"} +{"id": 2, "column1": "value1", "column2": "value2"} +``` + +### CSV and TSV + +CSV and TSV files should have a header row containing the names of the columns. + +By default, `dclient` will attempt to detect the types of the different columns in the CSV and TSV files - so if a column only ever contains numeric integers it will be stored as integers in the SQLite database. + +You can disable this and have every value treated as a string using `--no-detect-types`. + +### Other options + +- `--create` - create the table if it doesn't already exist +- `--replace` - replace any rows with a matching primary key +- `--ignore` - ignore any rows with a matching existing primary key +- `--pk id` - set a primary key (for if the table is being created) + +If you use `--create` a table will be created with rows to match the columns in your uploaded data - using the correctly detected types, unless you use `--no-detect-types` in which case every column will be of type `text`. 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..048546d 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", @@ -39,7 +39,7 @@ setup( "pytest-httpx", "cogapp", "pytest-mock", - "datasette", + "datasette>=1.0a2", ] }, python_requires=">=3.7", diff --git a/tests/test_insert.py b/tests/test_insert.py new file mode 100644 index 0000000..096ad32 --- /dev/null +++ b/tests/test_insert.py @@ -0,0 +1,318 @@ +import asyncio +from collections import namedtuple +from click.testing import CliRunner +from datasette.app import Datasette +from dclient.cli import cli +import httpx +import json +import pathlib +import pytest + + +@pytest.fixture +def assert_all_responses_were_requested() -> bool: + return False + + +@pytest.fixture +def non_mocked_hosts(): + # This ensures httpx-mock will not affect Datasette's own + # httpx calls made in the tests by datasette.client: + return ["localhost"] + + +def test_insert_mocked(httpx_mock, tmpdir): + httpx_mock.add_response( + json={ + "ok": True, + "database": "data", + "table": "table1", + "table_url": "http://datasette.example.com/data/table1", + "table_api_url": "http://datasette.example.com/data/table1.json", + "schema": "CREATE TABLE [table1] (...)", + "row_count": 100, + } + ) + path = pathlib.Path(tmpdir) / "data.csv" + path.write_text("a,b,c\n1,2,3\n") + runner = CliRunner() + result = runner.invoke( + cli, + [ + "insert", + "https://datasette.example.com/data", + "table1", + str(path), + "--csv", + "--token", + "x", + ], + catch_exceptions=False, + ) + assert result.exit_code == 0 + assert result.output == "Inserting rows\n" + request = httpx_mock.get_request() + assert request.headers["authorization"] == "Bearer x" + assert json.loads(request.read()) == {"rows": [{"a": 1, "b": 2, "c": 3}]} + + +SIMPLE_CSV = "a,b,c\n1,2,3\n" +SIMPLE_TSV = "a\tb\tc\n1\t2\t3\n" +SIMPLE_JSON = json.dumps( + [ + { + "a": 1, + "b": 2, + "c": 3, + } + ] +) +SIMPLE_JSON_NL = '{"a": 1, "b": 2, "c": 3}\n' +LATIN1_CSV = ( + b"date,name,latitude,longitude\n" + b"2020-01-01,Barra da Lagoa,-27.574,-48.422\n" + b"2020-03-04,S\xe3o Paulo,-23.561,-46.645\n" + b"2020-04-05,Salta,-24.793:-65.408" +) + + +InsertTest = namedtuple( + "InsertTest", + ( + "input_data", + "cmd_args", + "table_exists", + "expected_output", + "should_error", + "expected_table_json", + ), +) + + +def make_format_test(content, arg): + return InsertTest( + input_data=content, + # Using --silent to force no display of progress bar, since it won't + # be shown for the JSON formats anyway + cmd_args=["--silent", "--create"] + ([arg] if arg is not None else []), + table_exists=False, + expected_output="", + should_error=False, + expected_table_json=[{"rowid": 1, "a": 1, "b": 2, "c": 3}], + ) + + +@pytest.mark.parametrize( + "input_data,cmd_args,table_exists,expected_output,should_error,expected_table_json", + ( + # Auto-detect formats + make_format_test(SIMPLE_CSV, None), + make_format_test(SIMPLE_TSV, None), + make_format_test(SIMPLE_JSON, None), + make_format_test(SIMPLE_JSON_NL, None), + # Explicit formats + make_format_test(SIMPLE_CSV, "--csv"), + make_format_test(SIMPLE_TSV, "--tsv"), + make_format_test(SIMPLE_JSON, "--json"), + make_format_test(SIMPLE_JSON_NL, "--nl"), + # No --create option should error: + InsertTest( + input_data=SIMPLE_CSV, + cmd_args=[], + table_exists=False, + expected_output="Inserting rows\nError: Table not found: table1\n", + should_error=True, + expected_table_json=None, + ), + # --no-detect-types + InsertTest( + input_data=SIMPLE_CSV, + cmd_args=["--no-detect-types", "--create"], + table_exists=False, + expected_output="Inserting rows\n", + should_error=False, + expected_table_json=[{"rowid": 1, "a": "1", "b": "2", "c": "3"}], + ), + # --encoding - without it this should error: + InsertTest( + input_data=LATIN1_CSV, + cmd_args=["--no-detect-types", "--create", "--csv"], + table_exists=False, + expected_output="Inserting rows\n", + should_error=True, + expected_table_json=None, + ), + # --encoding - with it this should work: + InsertTest( + input_data=LATIN1_CSV, + cmd_args=[ + "--no-detect-types", + "--create", + "--encoding", + "latin-1", + "--csv", + ], + table_exists=False, + expected_output="Inserting rows\n", + should_error=False, + expected_table_json=[ + { + "rowid": 1, + "date": "2020-01-01", + "name": "Barra da Lagoa", + "latitude": "-27.574", + "longitude": "-48.422", + }, + { + "rowid": 2, + "date": "2020-03-04", + "name": "São Paulo", + "latitude": "-23.561", + "longitude": "-46.645", + }, + { + "rowid": 3, + "date": "2020-04-05", + "name": "Salta", + "latitude": "-24.793:-65.408", + "longitude": None, + }, + ], + ), + # Existing table, conflicting pk + InsertTest( + input_data=SIMPLE_CSV, + cmd_args=[], + table_exists=True, + expected_output="Inserting rows\nUNIQUE constraint failed: table1.a\nError: UNIQUE constraint failed: table1.a\n", + should_error=True, + expected_table_json=[{"a": 1, "b": 2, "c": 3}, {"a": 4, "b": 5, "c": 6}], + ), + # Existing table, --replace + InsertTest( + input_data="a,b,c\n1,2,4\n", + cmd_args=["--replace"], + table_exists=True, + expected_output="Inserting rows\n", + should_error=False, + expected_table_json=[{"a": 1, "b": 2, "c": 4}, {"a": 4, "b": 5, "c": 6}], + ), + # Existing table, --ignore + InsertTest( + input_data="a,b,c\n1,2,4\n", + cmd_args=["--ignore"], + table_exists=True, + expected_output="Inserting rows\n", + should_error=False, + expected_table_json=[{"a": 1, "b": 2, "c": 3}, {"a": 4, "b": 5, "c": 6}], + ), + ), +) +def test_insert_against_datasette( + httpx_mock, + tmpdir, + input_data, + cmd_args, + table_exists, + expected_output, + should_error, + expected_table_json, +): + ds = Datasette( + metadata={ + "permissions": { + "create-table": {"id": "*"}, + "insert-row": {"id": "*"}, + "update-row": {"id": "*"}, + } + } + ) + db = ds.add_memory_database("data") + loop = asyncio.get_event_loop() + + # Drop all tables in the database each time, because in-memory + # databases persist in between test runs + drop_all_tables(db, loop) + + if table_exists: + + async def run_table_exists(): + await db.execute_write( + "create table table1 (a integer primary key, b integer, c integer)" + ) + await db.execute_write( + "insert into table1 (a, b, c) values (1, 2, 3), (4, 5, 6)" + ) + + loop.run_until_complete(run_table_exists()) + + token = ds.create_token("actor") + + # These are useful with pytest --pdb to see what happened + datasette_requests = [] + datasette_responses = [] + + def custom_response(request: httpx.Request): + # Need to run this in async loop, because dclient itself uses + # sync HTTPX and not async HTTPX + async def run(): + datasette_requests.append(request) + response = await ds.client.request( + request.method, + request.url.path, + json=json.loads(request.read()), + headers=request.headers, + ) + # Create a fresh response to avoid an error where stream has been consumed + response = httpx.Response( + status_code=response.status_code, + headers=response.headers, + content=response.content, + ) + datasette_responses.append(response) + return response + + return loop.run_until_complete(run()) + + httpx_mock.add_callback(custom_response) + + path = pathlib.Path(tmpdir) / "data.txt" + if isinstance(input_data, str): + path.write_text(input_data) + else: + path.write_bytes(input_data) + runner = CliRunner() + result = runner.invoke( + cli, + [ + "insert", + "http://datasette.example.com/data", + "table1", + str(path), + "--token", + token, + ] + + cmd_args, + ) + if not should_error: + assert result.exit_code == 0 + else: + assert result.exit_code != 0 + assert result.output == expected_output + + if expected_table_json: + + async def fetch_table(): + response = await ds.client.get("/data/table1.json?_shape=array") + return response + + response = loop.run_until_complete(fetch_table()) + assert response.json() == expected_table_json + + +def drop_all_tables(db, loop): + async def run(): + for table in await db.table_names(): + await db.execute_write("drop table {}".format(table)) + + loop.run_until_complete(run())