dclient auth commands

Closes #3
This commit is contained in:
Simon Willison 2023-07-17 16:23:53 -07:00 committed by GitHub
commit 8a205d81a3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 304 additions and 12 deletions

View file

@ -81,6 +81,30 @@ Options:
```
<!-- [[[end]]] -->
## Authentication
`dclient` can handle API tokens for Datasette instances that require authentication.
You can pass an API token to `query` using `-t/--token` like this:
```bash
dclient query https://latest.datasette.io/fixtures "select * from facetable" -t dstok_mytoken
```
You can also store tokens for a specific URL prefix. To always use `dstok_mytoken` for any URL on the `https://latest.datasette.io/` instance you can run this:
```bash
dclient auth add https://latest.datasette.io/
```
Then paste in the token and hit enter when prompted to do so.
To list which URLs you have set tokens for, run the `auth list` command:
```bash
dclient auth list
```
To delete the token for a specific URL, run `auth remove`:
```bash
dclient auth remove https://latest.datasette.io/
```
## Aliases
You can assign an alias to a Datasette database using the `dclient alias` command:

View file

@ -2,6 +2,7 @@ import click
import httpx
import json
import pathlib
from .utils import token_for_url
def get_config_dir():
@ -17,7 +18,8 @@ def cli():
@cli.command()
@click.argument("url")
@click.argument("sql")
def query(url, sql):
@click.option("--token", "-t", help="API token")
def query(url, sql, token):
"""
Run a SQL query against a Datasette database URL
@ -29,17 +31,50 @@ def query(url, sql):
url = aliases[url]
if not url.endswith(".json"):
url += ".json"
response = httpx.get(url, params={"sql": sql, "_shape": "objects"})
if response.status_code != 200 or not response.json()["ok"]:
data = response.json()
if token is None:
# Maybe there's a token in auth.json?
token = token_for_url(url, _load_auths(get_config_dir() / "auth.json"))
headers = {}
if token:
headers["Authorization"] = f"Bearer {token}"
response = httpx.get(url, params={"sql": sql, "_shape": "objects"}, headers=headers)
if response.status_code != 200:
# Is it valid JSON?
try:
data = response.json()
except json.JSONDecodeError:
raise click.ClickException(
"{} status code. Response was not valid JSON".format(
response.status_code
)
)
bits = []
if data.get("title"):
bits.append(data["title"])
if data.get("error"):
bits.append(data["error"])
raise click.ClickException(
"{} status code. {}".format(response.status_code, ": ".join(bits))
)
# We should have JSON now
try:
data = response.json()
except json.JSONDecodeError:
raise click.ClickException("Response was not valid JSON")
# ... but it may have a {"ok": false} error
if not data.get("ok"):
bits = []
if data.get("title"):
bits.append(data["title"])
if data.get("error"):
bits.append(data["error"])
if not bits:
bits = [json.dumps(data)]
raise click.ClickException(": ".join(bits))
else:
click.echo(json.dumps(response.json()["rows"], indent=2))
# Output results
click.echo(json.dumps(response.json()["rows"], indent=2))
@cli.group()
@ -60,10 +95,10 @@ def list_(_json):
click.echo(f"{alias} = {url}")
@alias.command()
@alias.command(name="add")
@click.argument("name")
@click.argument("url")
def add(name, url):
def alias_add(name, url):
"Add an alias"
config_dir = get_config_dir()
config_dir.mkdir(parents=True, exist_ok=True)
@ -73,9 +108,9 @@ def add(name, url):
aliases_file.write_text(json.dumps(aliases, indent=4))
@alias.command()
@alias.command(name="remove")
@click.argument("name")
def remove(name):
def alias_remove(name):
"Remove an alias"
config_dir = get_config_dir()
aliases_file = config_dir / "aliases.json"
@ -87,9 +122,75 @@ def remove(name):
raise click.ClickException("No such alias")
@cli.group()
def auth():
"Manage authentication for different instances"
@auth.command(name="add")
@click.argument("alias_or_url")
@click.option("--token", prompt=True, hide_input=True)
def auth_add(alias_or_url, token):
"""
Add an authentication token for an alias or URL
Example usage:
\b
dclient auth add https://datasette.io/content
Paste in the token when prompted.
"""
aliases_file = get_config_dir() / "aliases.json"
aliases = _load_aliases(aliases_file)
url = alias_or_url
if alias_or_url in aliases:
url = aliases[alias_or_url]
config_dir = get_config_dir()
config_dir.mkdir(parents=True, exist_ok=True)
auth_file = config_dir / "auth.json"
auths = _load_auths(auth_file)
auths[url] = token
auth_file.write_text(json.dumps(auths, indent=4))
@auth.command(name="list")
def auth_list():
"List stored API tokens"
auths_file = get_config_dir() / "auth.json"
click.echo("Tokens file: {}".format(auths_file))
auths = _load_auths(auths_file)
if auths:
click.echo()
for url, token in auths.items():
click.echo("{}:\t{}..".format(url, token[:1]))
@auth.command(name="remove")
@click.argument("alias_or_url")
def auth_remove(alias_or_url):
"Remove the API token for an alias or URL"
config_dir = get_config_dir()
auth_file = config_dir / "auth.json"
auths = _load_auths(auth_file)
try:
del auths[alias_or_url]
auth_file.write_text(json.dumps(auths, indent=4))
except KeyError:
raise click.ClickException("No such URL or alias")
def _load_aliases(aliases_file):
if aliases_file.exists():
aliases = json.loads(aliases_file.read_text())
else:
aliases = {}
return aliases
def _load_auths(auth_file):
if auth_file.exists():
auths = json.loads(auth_file.read_text())
else:
auths = {}
return auths

42
dclient/utils.py Normal file
View file

@ -0,0 +1,42 @@
from urllib.parse import urlparse
def url_matches_prefix(url: str, prefix: str) -> bool:
url_parsed = urlparse(url)
prefix_parsed = urlparse(prefix)
if (
url_parsed.scheme != prefix_parsed.scheme
or url_parsed.netloc != prefix_parsed.netloc
):
return False
if url_parsed.path == prefix_parsed.path:
return True
url_path = url_parsed.path
# Add '/' to the end of the paths if they don't already end with '/'
prefix_path = prefix_parsed.path
# Special treatment for /foo.json
if url_path.endswith(".json"):
if prefix_path + ".json" == url_path:
return True
if not prefix_path.endswith("/"):
prefix_path += "/"
return url_path.startswith(prefix_path)
def token_for_url(url: str, tokens: dict) -> str:
matches = []
for prefix_url, token in tokens.items():
if url_matches_prefix(url, prefix_url):
matches.append((prefix_url, token))
# Sort by length of prefix_url, so that the longest match is first
matches.sort(key=lambda x: len(x[0]), reverse=True)
if matches:
return matches[0][1]
return None

41
tests/test_cli_auth.py Normal file
View file

@ -0,0 +1,41 @@
from click.testing import CliRunner
from dclient.cli import cli
import pathlib
import json
def test_auth(mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
runner = CliRunner()
result = runner.invoke(cli, ["auth", "list"])
assert result.exit_code == 0
assert result.output.startswith("Tokens file:")
# Should only have one line
assert len([line for line in result.output.split("\n") if line.strip()]) == 1
# Now add a token
result2 = runner.invoke(cli, ["auth", "add", "https://example.com"], input="xyz\n")
assert result2.exit_code == 0
# Check the tokens file
auth_file = pathlib.Path(tmpdir) / "auth.json"
assert json.loads(auth_file.read_text()) == {"https://example.com": "xyz"}
# auth list should show that now
result3 = runner.invoke(cli, ["auth", "list"])
assert result3.output.startswith("Tokens file:")
assert "https://example.com" in result3.output
# Remove should fail with an incorrect URL
result4 = runner.invoke(cli, ["auth", "remove", "https://example.com/foo"])
assert result4.exit_code == 1
assert result4.output == "Error: No such URL or alias\n"
# Remove should work with the correct URL
result5 = runner.invoke(cli, ["auth", "remove", "https://example.com"])
assert result5.exit_code == 0
assert result5.output == ""
# Check the tokens file
auth_file = pathlib.Path(tmpdir) / "auth.json"
assert json.loads(auth_file.read_text()) == {}

View file

@ -2,6 +2,7 @@ from click.testing import CliRunner
from dclient.cli import cli
import json
import pathlib
import pytest
def test_query_error(httpx_mock):
@ -17,10 +18,14 @@ def test_query_error(httpx_mock):
runner = CliRunner()
result = runner.invoke(cli, ["query", "https://example.com", "hello"])
assert result.exit_code == 1
assert result.output == "Error: Invalid SQL: Statement must be a SELECT\n"
assert (
result.output
== "Error: 400 status code. Invalid SQL: Statement must be a SELECT\n"
)
def test_query(httpx_mock):
@pytest.mark.parametrize("with_token", (False, True))
def test_query(httpx_mock, with_token):
httpx_mock.add_response(
json={
"ok": True,
@ -37,10 +42,22 @@ def test_query(httpx_mock):
status_code=200,
)
runner = CliRunner()
result = runner.invoke(cli, ["query", "https://example.com", "hello"])
args = ["query", "https://example.com", "hello"]
if with_token:
args.append("--token")
args.append("xyz")
result = runner.invoke(cli, args)
assert result.exit_code == 0
assert json.loads(result.output) == [{"5 * 2": 10}]
# Check the request
request = httpx_mock.get_request()
assert str(request.url) == "https://example.com.json?sql=hello&_shape=objects"
if with_token:
assert request.headers["authorization"] == "Bearer xyz"
else:
assert "authorization" not in request.headers
def test_aliases(mocker, tmpdir, httpx_mock):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))

67
tests/test_utils.py Normal file
View file

@ -0,0 +1,67 @@
from dclient.utils import token_for_url, url_matches_prefix
import pytest
@pytest.mark.parametrize(
"url,prefix_url,expected",
(
("https://example.com/foo/bar", "https://example.com/foo", True),
("https://example.com/foo/bar2", "https://example.com/foo/bar", False),
("https://example.com/foo/bar/baz", "https://example.com/foo/bar", True),
("https://example.com/foo/bar/baz", "https://example.com/foo", True),
("https://example.com/foo.json", "https://example.com/foo", True),
# different scheme
(
"http://example.com/foo/bar",
"https://example.com/foo",
False,
),
# different netloc
(
"https://example.org/foo/bar",
"https://example.com/foo",
False,
),
# exactly the same
("https://example.com/foo", "https://example.com/foo", True),
# trailing '/'
(
"https://example.com/foo/bar",
"https://example.com/foo/bar/",
False,
),
(
"https://example.com/foo/bar/baz",
"https://example.com/foo/bar/",
True,
),
),
)
def test_url_matches_prefix(url, prefix_url, expected):
assert url_matches_prefix(url, prefix_url) == expected
@pytest.mark.parametrize(
"url,tokens,expected",
(
("https://foo.com/bar", {"https://foo.com": "foo"}, "foo"),
("https://foo.com/bar", {"https://foo.com/baz": "baz"}, None),
(
"https://foo.com/bar",
{"https://foo.com": "foo", "https://foo.com/bar": "bar"},
"bar",
),
(
"https://foo.com/bar/baz",
{
"https://foo.com": "foo",
"https://foo.com/bar/baz": "baz",
"https://foo.com/bar": "bar",
},
"baz",
),
),
)
def test_token_for_url(url, tokens, expected):
# Should always return longest matching of the available options
assert token_for_url(url, tokens) == expected