From 98d714c5e1c3bf2c9e7210b7dd93f7d347fa234f Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Mon, 17 Jul 2023 16:03:05 -0700
Subject: [PATCH] Use longest possible matching URL for auth, refs #3

---
 dclient/cli.py      | 11 ++---------
 dclient/utils.py    | 12 ++++++++++++
 tests/test_utils.py | 28 +++++++++++++++++++++++++++-
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/dclient/cli.py b/dclient/cli.py
index 6ba3faa..7a51263 100644
--- a/dclient/cli.py
+++ b/dclient/cli.py
@@ -2,7 +2,7 @@ import click
 import httpx
 import json
 import pathlib
-from .utils import url_matches_prefix
+from .utils import token_for_url
 
 
 def get_config_dir():
@@ -33,7 +33,7 @@ def query(url, sql, token):
         url += ".json"
     if token is None:
         # Maybe there's a token in auth.json?
-        token = _token_for_url_from_auth(url, get_config_dir() / "auth.json")
+        token = token_for_url(url, _load_auths(get_config_dir() / "auth.json"))
     headers = {}
     if token:
         headers["Authorization"] = f"Bearer {token}"
@@ -77,13 +77,6 @@ def query(url, sql, token):
     click.echo(json.dumps(response.json()["rows"], indent=2))
 
 
-def _token_for_url_from_auth(url, auth_file):
-    auths = _load_auths(auth_file)
-    for auth_url, token in auths.items():
-        if url_matches_prefix(url, auth_url):
-            return token
-
-
 @cli.group()
 def alias():
     "Manage aliases for different instances"
diff --git a/dclient/utils.py b/dclient/utils.py
index b03df35..0a6576d 100644
--- a/dclient/utils.py
+++ b/dclient/utils.py
@@ -28,3 +28,15 @@ def url_matches_prefix(url: str, prefix: str) -> bool:
         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
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 2219a0b..81a63e8 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,4 +1,4 @@
-from dclient.utils import url_matches_prefix
+from dclient.utils import token_for_url, url_matches_prefix
 import pytest
 
 
@@ -39,3 +39,29 @@ import pytest
 )
 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