Add support for extension databases

Via a new keyword-only extension parameter to the connect function.

Bug: T396115

Author: lucaswerkmeister

src/toolforge/__init__.py

@@ -25,6 +25,7 @@
 from typing import Callable
 from typing import IO
 from typing import NewType
+from typing import Optional
 from typing import cast
 
 import decorator
@@ -41,7 +42,13 @@
 )
 
 
-def connect(dbname: str, cluster: str = "web", **kwargs: str) -> _Connection:
+def connect(
+    dbname: str,
+    cluster: str = "web",
+    *,
+    extension: Optional[str] = None,
+    **kwargs: str,
+) -> _Connection:
     """
     Get a database connection for the specified wiki.
 
@@ -63,6 +70,8 @@ def connect(dbname: str, cluster: str = "web", **kwargs: str) -> _Connection:
         host = f"s7.{domain}"
     else:
         host = f"{dbname}.{domain}"
+    if extension is not None:
+        host = f"{extension}.{host}"
     host = kwargs.pop("host", host)
 
     return _connect(

tests/main_test.py

@@ -55,51 +55,64 @@ def test_set_user_agent(self, tool, url, email, expect):
         }
 
     @pytest.mark.parametrize(
-        ("args", "expects"),
+        ("args", "kwargs", "expects"),
         [
             (
                 ["enwiki_p"],
+                {},
                 {
                     "database": "enwiki_p",
                     "host": "enwiki.web.db.svc.wikimedia.cloud",
                 },
             ),
             (
                 ["enwiki"],
+                {},
                 {
                     "database": "enwiki_p",
                     "host": "enwiki.web.db.svc.wikimedia.cloud",
                 },
             ),
             (
                 ["enwiki", "analytics"],
+                {},
                 {
                     "database": "enwiki_p",
                     "host": "enwiki.analytics.db.svc.wikimedia.cloud",
                 },
             ),
             (
                 ["meta", "analytics"],
+                {},
                 {
                     "database": "meta_p",
                     "host": "s7.analytics.db.svc.wikimedia.cloud",
                 },
             ),
+            (
+                ["wikidatawiki"],
+                {"extension": "termstore"},
+                {
+                    "database": "wikidatawiki_p",
+                    "host": "termstore.wikidatawiki.web.db.svc.wikimedia.cloud",
+                },
+            ),
         ],
     )
-    def test_connect(self, mocker, args, expects):
-        self._assert_connect(mocker, toolforge.connect, args, expects)
+    def test_connect(self, mocker, args, kwargs, expects):
+        self._assert_connect(mocker, toolforge.connect, args, kwargs, expects)
 
-    def _assert_connect(self, mocker, func, args, expect):
+    def _assert_connect(self, mocker, func, args, kwargs, expect):
         """Mock toolforge._connect and assert it is called as expected.
 
         :param func: Function to call after mocking toolforge._connect
-        :param args: Arguments for calling func
+        :param args: Positional arguments for calling func
+        :param kwargs: Keyword arguments for calling func
         :param expect: Dict of expected arguments to toolforge._connect
         """
         mm = mocker.patch("toolforge._connect")
         mm.return_value = None
-        conn = func(*args)
+        conn = func(*args, **kwargs)
         assert conn is None
         mm.assert_called_once_with(**expect)
 
@@ -109,6 +122,7 @@ def test_connect_rejects_unknown_cluster(self, mocker):
                 mocker,
                 toolforge.connect,
                 ["ignored", "not web or analytics"],
+                {},
                 {"ignored": "ignored"},
             )
 
@@ -132,7 +146,8 @@ def test_connect_rejects_unknown_cluster(self, mocker):
         ],
     )
     def test_toolsdb(self, mocker, args, expects):
-        self._assert_connect(mocker, toolforge.toolsdb, args, expects)
+        kwargs = {}
+        self._assert_connect(mocker, toolforge.toolsdb, args, kwargs, expects)
 
     def test_assert_private_file_no_args(self, mocker):
         load = mocker.Mock(