Merge ~cjwatson/launchpad:db-roles into launchpad:master

[Colin Watson]

Add ability to switch DB role after connecting

Our traditional strategy for deploying database credentials has been to have a password for each database user we want to use, and to deploy those to `.pgpass` files on each machine that needs to use each user. This works, but it's quite manual and cumbersome. With charmed deployments it gets harder: each `db` relation can only communicate a single set of credentials, which means that a script server would be stuck having to have perhaps tens of relations representing the wide variety of database users that the scripts it runs want to use: as well as being ugly, I expect this would also make Juju very slow.

The natural approach with the PostgreSQL charm (or when emulating it using `external-services`) is to set the `roles` attribute on the `db` relation, which causes the relevant user to be made a member of that set of roles. The `SET ROLE` SQL command can then switch to that role.

I don't see a straightforward way to make these two strategies entirely compatible, but it's simple enough to add a configuration flag that selects which strategy we're using, so I've taken that approach. `database.set_role_after_connecting` defaults to false which indicates our traditional strategy; setting it to True causes us to use the `SET ROLE` command after connecting instead, expecting the session user to have been made a member of the target role.

[Andrey Fedoseev (andrey-fedoseev)]

review: Approve

diff --git a/lib/lp/services/config/__init__.py b/lib/lp/services/config/__init__.py
index 859a38a..24ac875 100644
--- a/lib/lp/services/config/__init__.py
+++ b/lib/lp/services/config/__init__.py
@@ -289,6 +289,7 @@ class DatabaseConfig:
             "db_statement_timeout",
             "db_statement_timeout_precision",
             "isolation_level",
+            "set_role_after_connecting",
             "soft_request_timeout",
             "storm_cache",
             "storm_cache_size",
diff --git a/lib/lp/services/config/schema-lazr.conf b/lib/lp/services/config/schema-lazr.conf
index cf9c276..480d818 100644
--- a/lib/lp/services/config/schema-lazr.conf
+++ b/lib/lp/services/config/schema-lazr.conf
@@ -599,6 +599,13 @@ storm_cache_size: 10000
 # datatype: existing_directory
 replication_logdir: database/replication
 
+# If True, the connection string must contain a username, and the database
+# adapter will switch to the target `dbuser` using `SET ROLE` after
+# connecting (the session user must be a member of the target role).  If
+# False, the connection string must not contain a username, and the database
+# adapter will connect directly as the target `dbuser`.
+set_role_after_connecting: False
+
 
 [diff]
 # The maximum size in bytes to read from the librarian to make available in
diff --git a/lib/lp/services/webapp/adapter.py b/lib/lp/services/webapp/adapter.py
index a2585dd..a1b67ed 100644
--- a/lib/lp/services/webapp/adapter.py
+++ b/lib/lp/services/webapp/adapter.py
@@ -463,16 +463,24 @@ class LaunchpadDatabase(Postgres):
         # is reconnected it pays attention to any config changes.
         config_entry = "%s_%s" % (realm, flavor)
         connection_string = getattr(dbconfig, config_entry)
-        assert "user" not in parse_dsn(connection_string), (
-            "Database username should not be specified in "
-            "connection string (%s)." % connection_string
-        )
 
         # Try to lookup dbuser using the $realm_dbuser key. If this fails,
         # fallback to the dbuser key.
         dbuser = getattr(dbconfig, "%s_dbuser" % realm, dbconfig.dbuser)
+        parsed_dsn = parse_dsn(connection_string)
 
-        self._dsn = make_dsn(connection_string, user=dbuser)
+        if dbconfig.set_role_after_connecting:
+            assert "user" in parsed_dsn, (
+                "With set_role_after_connecting, database username must be "
+                "specified in connection string (%s)." % connection_string
+            )
+            self._dsn = make_dsn(connection_string)
+        else:
+            assert "user" not in parsed_dsn, (
+                "Database username must not be specified in "
+                "connection string (%s)." % connection_string
+            )
+            self._dsn = make_dsn(connection_string, user=dbuser)
 
         flags = _get_dirty_commit_flags()
 
@@ -483,6 +491,9 @@ class LaunchpadDatabase(Postgres):
 
         raw_connection = super().raw_connect()
 
+        if dbconfig.set_role_after_connecting and dbuser != parsed_dsn["user"]:
+            raw_connection.cursor().execute("SET ROLE %s", (dbuser,))
+
         # Set read only mode for the session.
         # An alternative would be to use the _ro users generated by
         # security.py, but this would needlessly double the number
diff --git a/lib/lp/services/webapp/tests/test_adapter.py b/lib/lp/services/webapp/tests/test_adapter.py
index ce03209..f409934 100644
--- a/lib/lp/services/webapp/tests/test_adapter.py
+++ b/lib/lp/services/webapp/tests/test_adapter.py
@@ -1,6 +1,7 @@
 # Copyright 2022 Canonical Ltd.  This software is licensed under the
 # GNU Affero General Public License version 3 (see the file LICENSE).
 
+from psycopg2.errors import InsufficientPrivilege
 from psycopg2.extensions import parse_dsn
 from zope.component import getUtility
 
@@ -19,36 +20,42 @@ class TestLaunchpadDatabase(TestCase):
 
     layer = DatabaseFunctionalLayer
 
-    def test_refuses_connection_string_with_user(self):
+    def setUp(self):
+        super().setUp()
         self.addCleanup(dbconfig.reset)
+
+    def assertCurrentUser(self, store, user):
+        self.assertEqual(
+            user, store.execute("SELECT current_user").get_one()[0]
+        )
+
+    def test_refuses_connstring_with_user(self):
         connstr = "dbname=%s user=foo" % DatabaseLayer._db_fixture.dbname
         dbconfig.override(rw_main_primary=connstr)
         disconnect_stores()
         self.assertRaisesWithContent(
             AssertionError,
-            "Database username should not be specified in connection string "
+            "Database username must not be specified in connection string "
             "(%s)." % connstr,
             getUtility(IStoreSelector).get,
             MAIN_STORE,
             PRIMARY_FLAVOR,
         )
 
-    def test_refuses_connection_string_uri_with_user(self):
-        self.addCleanup(dbconfig.reset)
+    def test_refuses_connstring_uri_with_user(self):
         connstr = "postgresql://foo@/%s" % DatabaseLayer._db_fixture.dbname
         dbconfig.override(rw_main_primary=connstr)
         disconnect_stores()
         self.assertRaisesWithContent(
             AssertionError,
-            "Database username should not be specified in connection string "
+            "Database username must not be specified in connection string "
             "(%s)." % connstr,
             getUtility(IStoreSelector).get,
             MAIN_STORE,
             PRIMARY_FLAVOR,
         )
 
-    def test_accepts_connection_string_without_user(self):
-        self.addCleanup(dbconfig.reset)
+    def test_accepts_connstring_without_user(self):
         connstr = "dbname=%s" % DatabaseLayer._db_fixture.dbname
         dbconfig.override(rw_main_primary=connstr, dbuser="ro")
         disconnect_stores()
@@ -57,9 +64,9 @@ class TestLaunchpadDatabase(TestCase):
             {"dbname": DatabaseLayer._db_fixture.dbname, "user": "ro"},
             parse_dsn(store.get_database()._dsn),
         )
+        self.assertCurrentUser(store, "ro")
 
-    def test_accepts_connection_string_uri_without_user(self):
-        self.addCleanup(dbconfig.reset)
+    def test_accepts_connstring_uri_without_user(self):
         connstr = "postgresql:///%s" % DatabaseLayer._db_fixture.dbname
         dbconfig.override(rw_main_primary=connstr, dbuser="ro")
         disconnect_stores()
@@ -68,3 +75,85 @@ class TestLaunchpadDatabase(TestCase):
             {"dbname": DatabaseLayer._db_fixture.dbname, "user": "ro"},
             parse_dsn(store.get_database()._dsn),
         )
+        self.assertCurrentUser(store, "ro")
+
+    def test_set_role_after_connecting_refuses_connstring_without_user(self):
+        connstr = "dbname=%s" % DatabaseLayer._db_fixture.dbname
+        dbconfig.override(
+            set_role_after_connecting=True,
+            rw_main_primary=connstr,
+            dbuser="read",
+        )
+        disconnect_stores()
+        self.assertRaisesWithContent(
+            AssertionError,
+            "With set_role_after_connecting, database username must be "
+            "specified in connection string (%s)." % connstr,
+            getUtility(IStoreSelector).get,
+            MAIN_STORE,
+            PRIMARY_FLAVOR,
+        )
+
+    def test_set_role_after_connecting_refuses_connstring_uri_without_user(
+        self,
+    ):
+        connstr = "postgresql:///%s" % DatabaseLayer._db_fixture.dbname
+        dbconfig.override(
+            set_role_after_connecting=True,
+            rw_main_primary=connstr,
+            dbuser="read",
+        )
+        disconnect_stores()
+        self.assertRaisesWithContent(
+            AssertionError,
+            "With set_role_after_connecting, database username must be "
+            "specified in connection string (%s)." % connstr,
+            getUtility(IStoreSelector).get,
+            MAIN_STORE,
+            PRIMARY_FLAVOR,
+        )
+
+    def test_set_role_after_connecting_accepts_connstring_with_user(self):
+        connstr = "dbname=%s user=ro" % DatabaseLayer._db_fixture.dbname
+        dbconfig.override(
+            set_role_after_connecting=True,
+            rw_main_primary=connstr,
+            dbuser="read",
+        )
+        disconnect_stores()
+        store = getUtility(IStoreSelector).get(MAIN_STORE, PRIMARY_FLAVOR)
+        self.assertEqual(
+            {"dbname": DatabaseLayer._db_fixture.dbname, "user": "ro"},
+            parse_dsn(store.get_database()._dsn),
+        )
+        self.assertCurrentUser(store, "read")
+
+    def test_set_role_after_connecting_accepts_connstring_uri_with_user(self):
+        connstr = "postgresql://ro@/%s" % DatabaseLayer._db_fixture.dbname
+        dbconfig.override(
+            set_role_after_connecting=True,
+            rw_main_primary=connstr,
+            dbuser="read",
+        )
+        disconnect_stores()
+        store = getUtility(IStoreSelector).get(MAIN_STORE, PRIMARY_FLAVOR)
+        self.assertEqual(
+            {"dbname": DatabaseLayer._db_fixture.dbname, "user": "ro"},
+            parse_dsn(store.get_database()._dsn),
+        )
+        self.assertCurrentUser(store, "read")
+
+    def test_set_role_after_connecting_not_member(self):
+        connstr = "dbname=%s user=ro" % DatabaseLayer._db_fixture.dbname
+        dbconfig.override(
+            set_role_after_connecting=True,
+            rw_main_primary=connstr,
+            dbuser="launchpad_main",
+        )
+        disconnect_stores()
+        self.assertRaises(
+            InsufficientPrivilege,
+            getUtility(IStoreSelector).get,
+            MAIN_STORE,
+            PRIMARY_FLAVOR,
+        )