Refactor format_dsn() method (#3069)

so that it doesn't take any decisions about which keywords should appear in the connection string and just uses a provided dict.

patroni/postgresql/bootstrap.py

@@ -216,15 +216,13 @@ class Bootstrap(object):
         cmd = config.get('post_bootstrap') or config.get('post_init')
         if cmd:
             r = self._postgresql.connection_pool.conn_kwargs
-            connstring = self._postgresql.config.format_dsn(r, True)
-            if 'host' not in r:
-                # https://www.postgresql.org/docs/current/static/libpq-pgpass.html
-                # A host name of localhost matches both TCP (host name localhost) and Unix domain socket
-                # (pghost empty or the default socket directory) connections coming from the local machine.
-                r['host'] = 'localhost'  # set it to localhost to write into pgpass
 
-            env = self._postgresql.config.write_pgpass(r)
+            # https://www.postgresql.org/docs/current/static/libpq-pgpass.html
+            # A host name of localhost matches both TCP (host name localhost) and Unix domain socket
+            # (pghost empty or the default socket directory) connections coming from the local machine.
+            env = self._postgresql.config.write_pgpass({'host': 'localhost', **r})
             env['PGOPTIONS'] = '-c synchronous_commit=local -c statement_timeout=0'
+            connstring = self._postgresql.config.format_dsn({**r, 'password': None})
 
             try:
                 ret = self._postgresql.cancellable.call(shlex.split(cmd) + [connstring], env=env)
@@ -258,7 +256,7 @@ class Bootstrap(object):
             r = clone_member.conn_kwargs(self._postgresql.config.replication)
             # add the credentials to connect to the replica origin to pgpass.
             env = self._postgresql.config.write_pgpass(r)
-            connstring = self._postgresql.config.format_dsn(r, True)
+            connstring = self._postgresql.config.format_dsn({**r, 'password': None})
         else:
             connstring = ''
             env = os.environ.copy()

patroni/postgresql/config.py

@@ -574,27 +574,26 @@ class ConfigHandler(object):
             del ret['dbname']
         return ret
 
-    def format_dsn(self, params: Dict[str, Any], include_dbname: bool = False) -> str:
+    def format_dsn(self, params: Dict[str, Any]) -> str:
+        """Format connection string from connection parameters.
+
+        .. note::
+            only parameters from the below list are considered and values are escaped.
+
+        :param params: :class:`dict` object with connection parameters.
+
+        :returns: a connection string in a format "key1=value2 key2=value2"
+        """
         # A list of keywords that can be found in a conninfo string. Follows what is acceptable by libpq
         keywords = ('dbname', 'user', 'passfile' if params.get('passfile') else 'password', 'host', 'port',
                     'sslmode', 'sslcompression', 'sslcert', 'sslkey', 'sslpassword', 'sslrootcert', 'sslcrl',
                     'sslcrldir', 'application_name', 'krbsrvname', 'gssencmode', 'channel_binding',
                     'target_session_attrs')
-        if include_dbname:
-            params = params.copy()
-            if 'dbname' not in params:
-                params['dbname'] = self._postgresql.database
-            # we are abusing information about the necessity of dbname
-            # dsn should contain passfile or password only if there is no dbname in it (it is used in recovery.conf)
-            skip = {'passfile', 'password'}
-        else:
-            skip = {'dbname'}
 
         def escape(value: Any) -> str:
             return re.sub(r'([\'\\ ])', r'\\\1', str(value))
 
-        return ' '.join('{0}={1}'.format(kw, escape(params[kw])) for kw in keywords
-                        if kw not in skip and params.get(kw) is not None)
+        return ' '.join('{0}={1}'.format(kw, escape(params[kw])) for kw in keywords if params.get(kw) is not None)
 
     def _write_recovery_params(self, fd: ConfigWriter, recovery_params: CaseInsensitiveDict) -> None:
         if self._postgresql.major_version >= 90500:
@@ -606,8 +605,7 @@ class ConfigHandler(object):
                 recovery_params.setdefault('pause_at_recovery_target', 'false')
         for name, value in sorted(recovery_params.items()):
             if name == 'primary_conninfo':
-                if 'password' in value and self._postgresql.major_version >= 100000:
-                    self.write_pgpass(value)
+                if self._postgresql.major_version >= 100000 and 'PGPASSFILE' in self.write_pgpass(value):
                     value['passfile'] = self._passfile = self._pgpass
                     self._passfile_mtime = mtime(self._pgpass)
                 value = self.format_dsn(value)
@@ -754,7 +752,7 @@ class ConfigHandler(object):
         if passfile_mtime:
             try:
                 with open(passfile) as f:
-                    wanted_lines = (self._pgpass_line(wanted_primary_conninfo) or '').splitlines()
+                    wanted_lines = (self._pgpass_content(wanted_primary_conninfo) or '').splitlines()
                     file_lines = f.read().splitlines()
                     if set(wanted_lines) == set(file_lines):
                         self._passfile = passfile
@@ -873,27 +871,38 @@ class ConfigHandler(object):
             os.unlink(name)
 
     @staticmethod
-    def _pgpass_line(record: Dict[str, Any]) -> Optional[str]:
+    def _pgpass_content(record: Dict[str, Any]) -> Optional[str]:
+        """Generate content of `pgpassfile` based on connection parameters.
+
+        .. note::
+            In case if ``host`` is a comma separated string we generate one line per host.
+
+        :param record: :class:`dict` object with connection parameters.
+        :returns: a string with generated content of pgpassfile or ``None`` if there is no ``password``.
+        """
         if 'password' in record:
             def escape(value: Any) -> str:
                 return re.sub(r'([:\\])', r'\\\1', str(value))
 
-            record = {n: escape(record.get(n) or '*') for n in ('host', 'port', 'user', 'password')}
-            # 'host' could be several comma-separated hostnames, in this case
-            # we need to write on pgpass line per host
-            line = ''
-            for hostname in record['host'].split(','):
-                line += hostname + ':{port}:*:{user}:{password}'.format(**record) + '\n'
-            return line.rstrip()
+            # 'host' could be several comma-separated hostnames, in this case we need to write on pgpass line per host
+            hosts = map(escape, filter(None, map(str.strip, (record.get('host') or '*').split(','))))
+            record = {n: escape(record.get(n) or '*') for n in ('port', 'user', 'password')}
+            return '\n'.join('{host}:{port}:*:{user}:{password}'.format(**record, host=host) for host in hosts)
 
     def write_pgpass(self, record: Dict[str, Any]) -> Dict[str, str]:
-        line = self._pgpass_line(record)
-        if not line:
+        """Maybe creates :attr:`_passfile` based on connection parameters.
+
+        :param record: :class:`dict` object with connection parameters.
+
+        :returns: a copy of environment variables, that will include ``PGPASSFILE`` in case if the file was written.
+        """
+        content = self._pgpass_content(record)
+        if not content:
             return os.environ.copy()
 
         with open(self._pgpass, 'w') as f:
             os.chmod(self._pgpass, stat.S_IWRITE | stat.S_IREAD)
-            f.write(line)
+            f.write(content)
 
         return {**os.environ, 'PGPASSFILE': self._pgpass}
 

patroni/postgresql/rewind.py

@@ -412,11 +412,22 @@ class Rewind(object):
             except Exception as e:
                 logger.warning('Unable to clean %s: %r', replslot_dir, e)
 
-    def pg_rewind(self, r: Dict[str, Any]) -> bool:
-        # prepare pg_rewind connection
-        env = self._postgresql.config.write_pgpass(r)
+    def pg_rewind(self, conn_kwargs: Dict[str, Any]) -> bool:
+        """Do pg_rewind.
+
+        .. note::
+            If ``pg_rewind`` doesn't support ``--restore-target-wal`` parameter and exited with non zero code,
+            Patroni will parse stderr/stdout to figure out if it failed due to a missing WAL file and will
+            repeat an attempt after downloading the missing file using ``restore_command``.
+
+        :param conn_kwargs: :class:`dict` object with connection parameters.
+
+        :returns: ``True`` if ``pg_rewind`` finished successfully, ``False`` otherwise.
+        """
+        # prepare pg_rewind connection string
+        env = self._postgresql.config.write_pgpass(conn_kwargs)
         env.update(LANG='C', LC_ALL='C', PGOPTIONS='-c statement_timeout=0')
-        dsn = self._postgresql.config.format_dsn(r, True)
+        dsn = self._postgresql.config.format_dsn({**conn_kwargs, 'password': None})
         logger.info('running pg_rewind from %s', dsn)
 
         restore_command = (self._postgresql.config.get('recovery_conf') or EMPTY_DICT).get('restore_command') \

tests/test_rewind.py

@@ -103,7 +103,10 @@ class TestRewind(BaseTestPostgresql):
     @patch.object(Postgresql, 'stop', Mock(return_value=False))
     @patch.object(Postgresql, 'start', Mock())
     def test_execute(self, mock_checkpoint):
-        self.r.execute(self.leader)
+        with patch('patroni.postgresql.rewind.logger.info') as mock_logger:
+            self.r.execute(self.leader)
+            self.assertEqual(mock_logger.call_args_list[0][0],
+                             ('running pg_rewind from %s', 'dbname=postgres user=foo host=127.0.0.1 port=5435'))
         with patch.object(Postgresql, 'major_version', PropertyMock(return_value=130000)):
             self.r.execute(self.leader)
             with patch.object(MockCursor, 'fetchone', Mock(side_effect=Exception)):