PostgreSQL: allow connection parameters to be specified

redash/query_runner/pg.py

@@ -138,6 +138,15 @@ def _get_ssl_config(configuration):
     return ssl_config
 
 
+def _parse_dsn(configuration):
+    standard_params = {"user", "password", "host", "port", "dbname"}
+    params = psycopg2.extensions.parse_dsn(configuration.get("dsn", ""))
+    overlap = standard_params.intersection(params.keys())
+    if overlap:
+        raise ValueError("Extra parameters may not contain {}".format(overlap))
+    return params
+
+
 class PostgreSQL(BaseSQLQueryRunner):
     noop_query = "SELECT 1"
 
@@ -151,6 +160,7 @@ def configuration_schema(cls):
                 "host": {"type": "string", "default": "127.0.0.1"},
                 "port": {"type": "number", "default": 5432},
                 "dbname": {"type": "string", "title": "Database Name"},
+                "dsn": {"type": "string", "default": "application_name=redash", "title": "Parameters"},
                 "sslmode": {
                     "type": "string",
                     "title": "SSL Mode",
@@ -244,6 +254,7 @@ def _get_tables(self, schema):
 
     def _get_connection(self):
         self.ssl_config = _get_ssl_config(self.configuration)
+        self.dsn = _parse_dsn(self.configuration)
         connection = psycopg2.connect(
             user=self.configuration.get("user"),
             password=self.configuration.get("password"),
@@ -252,6 +263,7 @@ def _get_connection(self):
             dbname=self.configuration.get("dbname"),
             async_=True,
             **self.ssl_config,
+            **self.dsn,
         )
 
         return connection

tests/query_runner/test_pg.py

@@ -1,6 +1,16 @@
 from unittest import TestCase
 
-from redash.query_runner.pg import build_schema
+from redash.query_runner.pg import _parse_dsn, build_schema
+
+
+class TestParameters(TestCase):
+    def test_parse_dsn(self):
+        configuration = {"dsn": "application_name=redash connect_timeout=5"}
+        self.assertDictEqual(_parse_dsn(configuration), {"application_name": "redash", "connect_timeout": "5"})
+
+    def test_parse_dsn_not_permitted(self):
+        configuration = {"dsn": "password=xyz"}
+        self.assertRaises(ValueError, _parse_dsn, configuration)
 
 
 class TestBuildSchema(TestCase):