+ try:
+ conn = psycopg2.connect(dsn, connection_factory=Connection)
+ ctxmgr = cast(ConnectionContext, contextlib.closing(conn))
+ ctxmgr.connection = conn
+ return ctxmgr
+ except psycopg2.OperationalError as err:
+ raise UsageError(f"Cannot connect to database: {err}") from err
+
+
+# Translation from PG connection string parameters to PG environment variables.
+# Derived from https://www.postgresql.org/docs/current/libpq-envars.html.
+_PG_CONNECTION_STRINGS = {
+ 'host': 'PGHOST',
+ 'hostaddr': 'PGHOSTADDR',
+ 'port': 'PGPORT',
+ 'dbname': 'PGDATABASE',
+ 'user': 'PGUSER',
+ 'password': 'PGPASSWORD',
+ 'passfile': 'PGPASSFILE',
+ 'channel_binding': 'PGCHANNELBINDING',
+ 'service': 'PGSERVICE',
+ 'options': 'PGOPTIONS',
+ 'application_name': 'PGAPPNAME',
+ 'sslmode': 'PGSSLMODE',
+ 'requiressl': 'PGREQUIRESSL',
+ 'sslcompression': 'PGSSLCOMPRESSION',
+ 'sslcert': 'PGSSLCERT',
+ 'sslkey': 'PGSSLKEY',
+ 'sslrootcert': 'PGSSLROOTCERT',
+ 'sslcrl': 'PGSSLCRL',
+ 'requirepeer': 'PGREQUIREPEER',
+ 'ssl_min_protocol_version': 'PGSSLMINPROTOCOLVERSION',
+ 'ssl_max_protocol_version': 'PGSSLMAXPROTOCOLVERSION',
+ 'gssencmode': 'PGGSSENCMODE',
+ 'krbsrvname': 'PGKRBSRVNAME',
+ 'gsslib': 'PGGSSLIB',
+ 'connect_timeout': 'PGCONNECT_TIMEOUT',
+ 'target_session_attrs': 'PGTARGETSESSIONATTRS',
+}
+
+
+def get_pg_env(dsn: str,
+ base_env: Optional[SysEnv] = None) -> Dict[str, str]:
+ """ Return a copy of `base_env` with the environment variables for
+ PostgresSQL set up from the given database connection string.
+ If `base_env` is None, then the OS environment is used as a base
+ environment.
+ """
+ env = dict(base_env if base_env is not None else os.environ)
+
+ for param, value in psycopg2.extensions.parse_dsn(dsn).items():
+ if param in _PG_CONNECTION_STRINGS:
+ env[_PG_CONNECTION_STRINGS[param]] = value
+ else:
+ LOG.error("Unknown connection parameter '%s' ignored.", param)
+
+ return env