Initial oauth flow for use with datasette-oauth

This commit is contained in:
Simon Willison 2026-02-26 15:00:15 -08:00
commit 4df076d1f0
7 changed files with 503 additions and 102 deletions

View file

@ -2,6 +2,7 @@ from click.testing import CliRunner
from dclient.cli import cli
import pathlib
import json
import pytest
def test_auth(mocker, tmpdir):
@ -39,3 +40,139 @@ def test_auth(mocker, tmpdir):
# Check the tokens file
auth_file = pathlib.Path(tmpdir) / "auth.json"
assert json.loads(auth_file.read_text()) == {}
# -- login command (OAuth device flow) --
DEVICE_RESPONSE = {
"device_code": "devcode123",
"user_code": "ABCD-EFGH",
"verification_uri": "https://example.com/-/oauth/device/verify",
"expires_in": 900,
"interval": 0,
}
TOKEN_SUCCESS = {
"access_token": "dstok_abc123",
"token_type": "bearer",
"expires_in": 3600,
}
def test_login_with_url(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
mocker.patch("dclient.cli.time.sleep")
httpx_mock.add_response(json=DEVICE_RESPONSE, status_code=200)
httpx_mock.add_response(json=TOKEN_SUCCESS, status_code=200)
runner = CliRunner()
result = runner.invoke(cli, ["login", "https://example.com/"])
assert result.exit_code == 0
assert "ABCD-EFGH" in result.output
assert "Login successful" in result.output
# Token should be saved
auth_file = pathlib.Path(tmpdir) / "auth.json"
auths = json.loads(auth_file.read_text())
assert auths["https://example.com/"] == "dstok_abc123"
def test_login_adds_trailing_slash(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
mocker.patch("dclient.cli.time.sleep")
httpx_mock.add_response(json=DEVICE_RESPONSE, status_code=200)
httpx_mock.add_response(json=TOKEN_SUCCESS, status_code=200)
runner = CliRunner()
result = runner.invoke(cli, ["login", "https://example.com"])
assert result.exit_code == 0
# Check that the device request went to the right URL
requests = httpx_mock.get_requests()
assert str(requests[0].url) == "https://example.com/-/oauth/device"
def test_login_with_alias(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
mocker.patch("dclient.cli.time.sleep")
# Set up an alias first
config_file = pathlib.Path(tmpdir) / "config.json"
config_file.write_text(
json.dumps(
{
"default_instance": None,
"instances": {
"prod": {
"url": "https://prod.example.com",
"default_database": None,
}
},
}
)
)
httpx_mock.add_response(json=DEVICE_RESPONSE, status_code=200)
httpx_mock.add_response(json=TOKEN_SUCCESS, status_code=200)
runner = CliRunner()
result = runner.invoke(cli, ["login", "prod"])
assert result.exit_code == 0
assert "Login successful" in result.output
# Token should be saved by alias name
auth_file = pathlib.Path(tmpdir) / "auth.json"
auths = json.loads(auth_file.read_text())
assert auths["prod"] == "dstok_abc123"
def test_login_interactive_prompt(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
mocker.patch("dclient.cli.time.sleep")
httpx_mock.add_response(json=DEVICE_RESPONSE, status_code=200)
httpx_mock.add_response(json=TOKEN_SUCCESS, status_code=200)
runner = CliRunner()
result = runner.invoke(cli, ["login"], input="https://example.com/\n")
assert result.exit_code == 0
assert "Instance URL or alias" in result.output
assert "Login successful" in result.output
def test_login_access_denied(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
mocker.patch("dclient.cli.time.sleep")
httpx_mock.add_response(json=DEVICE_RESPONSE, status_code=200)
httpx_mock.add_response(json={"error": "access_denied"}, status_code=400)
runner = CliRunner()
result = runner.invoke(cli, ["login", "https://example.com/"])
assert result.exit_code == 1
assert "Authorization denied" in result.output
def test_login_expired_token(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
mocker.patch("dclient.cli.time.sleep")
httpx_mock.add_response(json=DEVICE_RESPONSE, status_code=200)
httpx_mock.add_response(json={"error": "expired_token"}, status_code=400)
runner = CliRunner()
result = runner.invoke(cli, ["login", "https://example.com/"])
assert result.exit_code == 1
assert "expired" in result.output
def test_login_pending_then_success(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
mocker.patch("dclient.cli.time.sleep")
httpx_mock.add_response(json=DEVICE_RESPONSE, status_code=200)
# First poll: pending
httpx_mock.add_response(json={"error": "authorization_pending"}, status_code=400)
# Second poll: success
httpx_mock.add_response(json=TOKEN_SUCCESS, status_code=200)
runner = CliRunner()
result = runner.invoke(cli, ["login", "https://example.com/"])
assert result.exit_code == 0
assert "Login successful" in result.output
auth_file = pathlib.Path(tmpdir) / "auth.json"
auths = json.loads(auth_file.read_text())
assert auths["https://example.com/"] == "dstok_abc123"
def test_login_device_endpoint_error(httpx_mock, mocker, tmpdir):
mocker.patch("dclient.cli.get_config_dir", return_value=pathlib.Path(tmpdir))
httpx_mock.add_response(status_code=403)
runner = CliRunner()
result = runner.invoke(cli, ["login", "https://example.com/"])
assert result.exit_code == 1
assert "Failed to start login flow" in result.output

View file

@ -6,7 +6,6 @@ import json
import pathlib
import pytest
# -- databases command --
@ -22,9 +21,7 @@ def test_databases_json(httpx_mock, mocker, tmpdir):
status_code=200,
)
runner = CliRunner()
result = runner.invoke(
cli, ["databases", "-i", "https://example.com", "--json"]
)
result = runner.invoke(cli, ["databases", "-i", "https://example.com", "--json"])
assert result.exit_code == 0
data = json.loads(result.output)
assert len(data) == 2
@ -175,9 +172,7 @@ def test_tables_hidden(httpx_mock, mocker, tmpdir):
)
runner = CliRunner()
# Without --hidden
result = runner.invoke(
cli, ["tables", "-i", "https://example.com", "-d", "db"]
)
result = runner.invoke(cli, ["tables", "-i", "https://example.com", "-d", "db"])
assert "visible" in result.output
assert "hidden_t" not in result.output
@ -213,9 +208,7 @@ def test_plugins_json(httpx_mock, mocker, tmpdir):
status_code=200,
)
runner = CliRunner()
result = runner.invoke(
cli, ["plugins", "-i", "https://example.com", "--json"]
)
result = runner.invoke(cli, ["plugins", "-i", "https://example.com", "--json"])
assert result.exit_code == 0
data = json.loads(result.output)
assert len(data) == 2
@ -262,9 +255,7 @@ def test_schema_all_tables(httpx_mock, mocker, tmpdir):
status_code=200,
)
runner = CliRunner()
result = runner.invoke(
cli, ["schema", "-i", "https://example.com", "-d", "main"]
)
result = runner.invoke(cli, ["schema", "-i", "https://example.com", "-d", "main"])
assert result.exit_code == 0
assert "CREATE TABLE users" in result.output
assert "CREATE VIEW user_count" in result.output
@ -372,9 +363,7 @@ def test_default_query_with_database_override(httpx_mock, mocker, tmpdir):
status_code=200,
)
runner = CliRunner()
result = runner.invoke(
cli, ["select count(*) from events", "-d", "analytics"]
)
result = runner.invoke(cli, ["select count(*) from events", "-d", "analytics"])
assert result.exit_code == 0
request = httpx_mock.get_request()
assert request.url.path == "/analytics.json"
@ -410,9 +399,7 @@ def test_default_query_with_instance_override(httpx_mock, mocker, tmpdir):
status_code=200,
)
runner = CliRunner()
result = runner.invoke(
cli, ["select count(*) from users", "-i", "staging"]
)
result = runner.invoke(cli, ["select count(*) from users", "-i", "staging"])
assert result.exit_code == 0
request = httpx_mock.get_request()
assert request.url.host == "staging.example.com"
@ -528,9 +515,7 @@ def test_get_command(httpx_mock, mocker, tmpdir):
status_code=200,
)
runner = CliRunner()
result = runner.invoke(
cli, ["get", "/-/plugins.json", "-i", "https://example.com"]
)
result = runner.invoke(cli, ["get", "/-/plugins.json", "-i", "https://example.com"])
assert result.exit_code == 0
data = json.loads(result.output)
assert data == {"hello": "world"}
@ -549,8 +534,14 @@ def test_instances_plain(mocker, tmpdir):
{
"default_instance": "prod",
"instances": {
"prod": {"url": "https://prod.example.com", "default_database": "main"},
"staging": {"url": "https://staging.example.com", "default_database": None},
"prod": {
"url": "https://prod.example.com",
"default_database": "main",
},
"staging": {
"url": "https://staging.example.com",
"default_database": None,
},
},
}
)
@ -570,7 +561,10 @@ def test_instances_json(mocker, tmpdir):
{
"default_instance": "prod",
"instances": {
"prod": {"url": "https://prod.example.com", "default_database": "main"},
"prod": {
"url": "https://prod.example.com",
"default_database": "main",
},
},
}
)

View file

@ -12,7 +12,6 @@ import json
import pathlib
import pytest
# -- Config loading/saving --
@ -196,7 +195,9 @@ def test_resolve_token_from_flag(tmpdir):
"""An explicit --token flag is used directly."""
auth_file = pathlib.Path(tmpdir) / "auth.json"
config_file = pathlib.Path(tmpdir) / "config.json"
token = _resolve_token("explicit-token", "https://example.com", auth_file, config_file)
token = _resolve_token(
"explicit-token", "https://example.com", auth_file, config_file
)
assert token == "explicit-token"
@ -218,7 +219,9 @@ def test_resolve_token_from_auth_by_alias(tmpdir):
}
)
)
token = _resolve_token(None, "https://myapp.datasette.cloud", auth_file, config_file)
token = _resolve_token(
None, "https://myapp.datasette.cloud", auth_file, config_file
)
assert token == "tok123"

View file

@ -3,7 +3,6 @@ from dclient.cli import cli
import json
import pathlib
QUERY_RESPONSE = {
"ok": True,
"database": "data",
@ -42,9 +41,13 @@ def test_token_flag_overrides_datasette_token(httpx_mock, mocker, tmpdir):
result = runner.invoke(
cli,
[
"query", "data", "select 1",
"-i", "https://example.com",
"--token", "flag-token",
"query",
"data",
"select 1",
"-i",
"https://example.com",
"--token",
"flag-token",
],
)
assert result.exit_code == 0

View file

@ -9,9 +9,7 @@ def test_migration_simple(tmpdir):
"""Migrate a simple aliases.json with a database-in-URL alias."""
config_dir = pathlib.Path(tmpdir)
aliases_file = config_dir / "aliases.json"
aliases_file.write_text(
json.dumps({"content": "https://datasette.io/content"})
)
aliases_file.write_text(json.dumps({"content": "https://datasette.io/content"}))
_migrate_v1_to_v2(config_dir)
@ -29,9 +27,7 @@ def test_migration_no_path_segment(tmpdir):
"""Migrate an alias that has no database in the URL."""
config_dir = pathlib.Path(tmpdir)
aliases_file = config_dir / "aliases.json"
aliases_file.write_text(
json.dumps({"local": "http://localhost:8001"})
)
aliases_file.write_text(json.dumps({"local": "http://localhost:8001"}))
_migrate_v1_to_v2(config_dir)
@ -44,13 +40,9 @@ def test_migration_with_auth(tmpdir):
"""Auth keys are migrated from URLs to alias names."""
config_dir = pathlib.Path(tmpdir)
aliases_file = config_dir / "aliases.json"
aliases_file.write_text(
json.dumps({"content": "https://datasette.io/content"})
)
aliases_file.write_text(json.dumps({"content": "https://datasette.io/content"}))
auth_file = config_dir / "auth.json"
auth_file.write_text(
json.dumps({"https://datasette.io/content": "tok123"})
)
auth_file.write_text(json.dumps({"https://datasette.io/content": "tok123"}))
_migrate_v1_to_v2(config_dir)
@ -68,9 +60,7 @@ def test_migration_auth_url_fallback(tmpdir):
aliases_file = config_dir / "aliases.json"
aliases_file.write_text(json.dumps({}))
auth_file = config_dir / "auth.json"
auth_file.write_text(
json.dumps({"https://other.example.com": "tok456"})
)
auth_file.write_text(json.dumps({"https://other.example.com": "tok456"}))
_migrate_v1_to_v2(config_dir)
@ -105,9 +95,7 @@ def test_migration_multi_path_segments(tmpdir):
"""URL with multiple path segments stores URL as-is."""
config_dir = pathlib.Path(tmpdir)
aliases_file = config_dir / "aliases.json"
aliases_file.write_text(
json.dumps({"deep": "https://example.com/a/b/c"})
)
aliases_file.write_text(json.dumps({"deep": "https://example.com/a/b/c"}))
_migrate_v1_to_v2(config_dir)

View file

@ -53,7 +53,9 @@ def test_query(httpx_mock, with_token):
# Check the request
request = httpx_mock.get_request()
assert str(request.url) == "https://example.com/content.json?sql=hello&_shape=objects"
assert (
str(request.url) == "https://example.com/content.json?sql=hello&_shape=objects"
)
if with_token:
assert request.headers["authorization"] == "Bearer xyz"
else:
@ -67,9 +69,7 @@ def test_aliases(mocker, tmpdir, httpx_mock):
assert result.exit_code == 0
assert result.output == ""
result = runner.invoke(
cli, ["alias", "add", "foo", "https://example.com"]
)
result = runner.invoke(cli, ["alias", "add", "foo", "https://example.com"])
assert result.exit_code == 0
assert result.output == ""