Use toolforge envvars if available

If both of the relevant envvars are set (TOOL_REPLICA_USER +
TOOL_REPLICA_PASSWORD for the replicas, TOOL_TOOLSDB_USER +
TOOL_TOOLSDB_PASSWORD for toolsdb – they’re currently the same but might
diverge later), use them and skip reading the replica.my.cnf, which may
not be mounted. If only one envvar is set, don’t use it – this is mostly
“should never happen” territory but I think that makes slightly more
sense than trying to combine an incomplete set of envvars with the
replica.my.cnf file.

Bug: T339940

Author: lucaswerkmeister

.local-dictionary.txt

@@ -4,6 +4,7 @@ MERCHANTABILITY
 Mehta
 PrivateFileWorldReadableError
 dbname
+delenv
 fileno
 func
 intersphinx
@@ -15,6 +16,7 @@ pygments
 pymysql
 qualname
 rtd
+setenv
 sitematrix
 tmp
 toolforge

src/toolforge/__init__.py

@@ -74,6 +74,15 @@ def connect(
         host = f"{extension}.{host}"
     host = kwargs.pop("host", host)
 
+    try:
+        user = os.environ["TOOL_REPLICA_USER"]
+        password = os.environ["TOOL_REPLICA_PASSWORD"]
+    except KeyError:
+        kwargs.setdefault("read_default_file", os.path.expanduser("~/replica.my.cnf"))
+    else:
+        kwargs.setdefault("user", user)
+        kwargs.setdefault("password", password)
+
     return _connect(
         database=dbname + "_p",
         host=host,
@@ -84,7 +93,6 @@ def connect(
 def _connect(*args: str, **kwargs: str) -> _Connection:  # pragma: no cover
     """Wraper for pymysql.connect to make testing easier."""
     kw = {
-        "read_default_file": os.path.expanduser("~/replica.my.cnf"),
         "charset": "utf8mb4",
     }
     kw.update(kwargs)
@@ -98,6 +106,15 @@ def toolsdb(dbname: str, **kwargs: str) -> _Connection:
     :param `**kwargs`: For :meth:`pymysql.connect <pymysql.connections.Connection.__init__>`
     :return: :class:`pymysql.connections.Connection`
     """
+    try:
+        user = os.environ["TOOL_TOOLSDB_USER"]
+        password = os.environ["TOOL_TOOLSDB_PASSWORD"]
+    except KeyError:
+        kwargs.setdefault("read_default_file", os.path.expanduser("~/replica.my.cnf"))
+    else:
+        kwargs.setdefault("user", user)
+        kwargs.setdefault("password", password)
+
     return _connect(
         database=dbname,
         host="tools.db.svc.wikimedia.cloud",

tests/main_test.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 import io
 import json
+import os
 
 import pytest
 import requests
@@ -63,6 +64,7 @@ def test_set_user_agent(self, tool, url, email, expect):
                 {
                     "database": "enwiki_p",
                     "host": "enwiki.web.db.svc.wikimedia.cloud",
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
                 },
             ),
             (
@@ -71,6 +73,7 @@ def test_set_user_agent(self, tool, url, email, expect):
                 {
                     "database": "enwiki_p",
                     "host": "enwiki.web.db.svc.wikimedia.cloud",
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
                 },
             ),
             (
@@ -79,6 +82,7 @@ def test_set_user_agent(self, tool, url, email, expect):
                 {
                     "database": "enwiki_p",
                     "host": "enwiki.analytics.db.svc.wikimedia.cloud",
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
                 },
             ),
             (
@@ -87,6 +91,7 @@ def test_set_user_agent(self, tool, url, email, expect):
                 {
                     "database": "meta_p",
                     "host": "s7.analytics.db.svc.wikimedia.cloud",
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
                 },
             ),
             (
@@ -95,13 +100,77 @@ def test_set_user_agent(self, tool, url, email, expect):
                 {
                     "database": "wikidatawiki_p",
                     "host": "termstore.wikidatawiki.web.db.svc.wikimedia.cloud",
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
                 },
             ),
         ],
     )
     def test_connect(self, mocker, args, kwargs, expects):
         self._assert_connect(mocker, toolforge.connect, args, kwargs, expects)
 
+    @pytest.mark.parametrize(
+        ("env", "expects"),
+        [
+            (
+                {
+                    "TOOL_REPLICA_USER": "user name",
+                    "TOOL_REPLICA_PASSWORD": "password",
+                },
+                {
+                    "user": "user name",
+                    "password": "password",
+                },
+            ),
+            (
+                {
+                    "TOOL_REPLICA_USER": None,
+                    "TOOL_REPLICA_PASSWORD": None,
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+            (
+                {
+                    "TOOL_REPLICA_USER": "unused user name",
+                    "TOOL_REPLICA_PASSWORD": None,
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+            (
+                {
+                    "TOOL_REPLICA_USER": None,
+                    "TOOL_REPLICA_PASSWORD": "unused password",
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+            (
+                {
+                    "TOOL_TOOLSDB_USER": "unused",
+                    "TOOL_TOOLSDB_PASSWORD": "(toolsdb != replica)",
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+        ],
+    )
+    def test_connect_env(self, mocker, monkeypatch, env, expects):
+        for name, value in env.items():
+            if value is None:
+                monkeypatch.delenv(name, raising=False)
+            else:
+                monkeypatch.setenv(name, value)
+        args = ["enwiki_p"]
+        kwargs = {}
+        expects["database"] = "enwiki_p"
+        expects["host"] = "enwiki.web.db.svc.wikimedia.cloud"
+        self._assert_connect(mocker, toolforge.connect, args, kwargs, expects)
+
     def _assert_connect(self, mocker, func, args, kwargs, expect):
         """Mock toolforge._connect and assert it is called as expected.
 
@@ -134,13 +203,15 @@ def test_connect_rejects_unknown_cluster(self, mocker):
                 {
                     "database": "s12345__foo",
                     "host": "tools.db.svc.wikimedia.cloud",
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
                 },
             ),
             (
                 ["s12345__foo_p"],
                 {
                     "database": "s12345__foo_p",
                     "host": "tools.db.svc.wikimedia.cloud",
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
                 },
             ),
         ],
@@ -149,6 +220,69 @@ def test_toolsdb(self, mocker, args, expects):
         kwargs = {}
         self._assert_connect(mocker, toolforge.toolsdb, args, kwargs, expects)
 
+    @pytest.mark.parametrize(
+        ("env", "expects"),
+        [
+            (
+                {
+                    "TOOL_TOOLSDB_USER": "user name",
+                    "TOOL_TOOLSDB_PASSWORD": "password",
+                },
+                {
+                    "user": "user name",
+                    "password": "password",
+                },
+            ),
+            (
+                {
+                    "TOOL_TOOLSDB_USER": None,
+                    "TOOL_TOOLSDB_PASSWORD": None,
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+            (
+                {
+                    "TOOL_TOOLSDB_USER": "unused user name",
+                    "TOOL_TOOLSDB_PASSWORD": None,
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+            (
+                {
+                    "TOOL_TOOLSDB_USER": None,
+                    "TOOL_TOOLSDB_PASSWORD": "unused password",
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+            (
+                {
+                    "TOOL_REPLICA_USER": "unused",
+                    "TOOL_REPLICA_PASSWORD": "(toolsdb != replica)",
+                },
+                {
+                    "read_default_file": os.path.expanduser("~/replica.my.cnf"),
+                },
+            ),
+        ],
+    )
+    def test_toolsdb_env(self, mocker, monkeypatch, env, expects):
+        for name, value in env.items():
+            if value is None:
+                monkeypatch.delenv(name, raising=False)
+            else:
+                monkeypatch.setenv(name, value)
+        args = ["s12345__foo"]
+        kwargs = {}
+        expects["database"] = "s12345__foo"
+        expects["host"] = "tools.db.svc.wikimedia.cloud"
+        self._assert_connect(mocker, toolforge.toolsdb, args, kwargs, expects)
+
     def test_assert_private_file_no_args(self, mocker):
         load = mocker.Mock(
             return_value="data",