dclient insert command (#13)

* Move making queries to own page, refs #7
* Documentation for the dclient insert feature, refs #8
* Implemented progress bar for insert, refs #8
* --pk option
* --ignore and --replace for insert, refs #8
* --batch-size option, refs #8
* First insert test, using a mock - refs #8
* insert test that exercises against an in-memory Datasette instance, refs #8
* Insert tests now cover an error case, refs #8
* Tests for --ignore and --replace, refs #8
* Tests for different formats, refs #8
* Test for --no-detect-types
* Support for --encoding, refs #8

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`:

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:
+    if url_or_alias in aliases:
-        url = aliases[url]
+        url = aliases[url_or_alias]
-    if not url.endswith(".json"):
+    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()

docs/authentication.md

@@ -1,3 +1,5 @@
+(authentication)=
+
 # Authentication
 
 `dclient` can handle API tokens for Datasette instances that require authentication.

docs/index.md

@@ -41,4 +41,5 @@ maxdepth: 3
 queries
 aliases
 authentication
+inserting
 ```

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 <authentication>` 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 <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`.

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
 

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",

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())