Pass schema in credentials to Postgresql

plugins/postgres/dbt/adapters/postgres/connections.py

@@ -33,6 +33,9 @@
  'schema': {
  'type': 'string',
  },
+ 'search_path': {
+ 'type': 'string',
+ },
  'keepalives_idle': {
  'type': 'integer',
  },
@@ -53,7 +56,7 @@ def type(self):
  return 'postgres'
 
  def _connection_keys(self):
- return ('host', 'port', 'user', 'database', 'schema')
+ return ('host', 'port', 'user', 'database', 'schema', 'search_path')
 
 
 class PostgresConnectionManager(SQLConnectionManager):
@@ -105,6 +108,14 @@ def open(cls, connection):
  if keepalives_idle:
  kwargs['keepalives_idle'] = keepalives_idle
 
+ # psycopg2 doesn't support search_path officially,
+ # see https:/psycopg/psycopg2/issues/465
+ search_path = credentials.get('search_path')
+ if search_path is not None and search_path != '':
+ # see https://postgresql.org/docs/9.5/libpq-connect.html
+ kwargs['options'] = '-c search_path={}'.format(
+ search_path.replace(' ', '\\ '))
+
  try:
  handle = psycopg2.connect(
  dbname=credentials.database,

plugins/redshift/dbt/adapters/redshift/connections.py

@@ -10,7 +10,6 @@
 
 drop_lock = multiprocessing.Lock()
 
-
 REDSHIFT_CREDENTIALS_CONTRACT = {
  'type': 'object',
  'additionalProperties': False,
@@ -55,6 +54,9 @@
  'If using IAM auth, the ttl for the temporary credentials'
  )
  },
+ 'search_path': {
+ 'type': 'string',
+ },
  'keepalives_idle': {
  'type': 'integer',
  },
@@ -75,7 +77,9 @@ def type(self):
  return 'redshift'
 
  def _connection_keys(self):
- return ('host', 'port', 'user', 'database', 'schema', 'method')
+ return (
+ 'host', 'port', 'user', 'database', 'schema', 'method',
+ 'search_path')
 
 
 class RedshiftConnectionManager(PostgresConnectionManager):
@@ -94,7 +98,6 @@ def fresh_transaction(self, name=None):
  to use the default.
  """
  with drop_lock:
-
  connection = self.get_thread_connection()
 
  if connection.transaction_open:

test/unit/test_postgres_adapter.py

@@ -125,6 +125,38 @@ def test_changed_keepalive(self, psycopg2):
  connect_timeout=10,
  keepalives_idle=256)
 
+ @mock.patch('dbt.adapters.postgres.connections.psycopg2')
+ def test_search_path(self, psycopg2):
+ self.config.credentials = self.config.credentials.incorporate(
+ search_path="test"
+ )
+ connection = self.adapter.acquire_connection('dummy')
+
+ psycopg2.connect.assert_called_once_with(
+ dbname='postgres',
+ user='root',
+ host='thishostshouldnotexist',
+ password='password',
+ port=5432,
+ connect_timeout=10,
+ options="-c search_path=test")
+
+ @mock.patch('dbt.adapters.postgres.connections.psycopg2')
+ def test_schema_with_space(self, psycopg2):
+ self.config.credentials = self.config.credentials.incorporate(
+ search_path="test test"
+ )
+ connection = self.adapter.acquire_connection('dummy')
+
+ psycopg2.connect.assert_called_once_with(
+ dbname='postgres',
+ user='root',
+ host='thishostshouldnotexist',
+ password='password',
+ port=5432,
+ connect_timeout=10,
+ options="-c search_path=test\ test")
+
  @mock.patch('dbt.adapters.postgres.connections.psycopg2')
  def test_set_zero_keepalive(self, psycopg2):
  self.config.credentials = self.config.credentials.incorporate(

test/unit/test_redshift_adapter.py

@@ -160,6 +160,40 @@ def test_changed_keepalive(self, psycopg2):
  connect_timeout=10,
  keepalives_idle=256)
 
+ @mock.patch('dbt.adapters.postgres.connections.psycopg2')
+ def test_search_path(self, psycopg2):
+ self.config.credentials = self.config.credentials.incorporate(
+ search_path="test"
+ )
+ connection = self.adapter.acquire_connection('dummy')
+
+ psycopg2.connect.assert_called_once_with(
+ dbname='redshift',
+ user='root',
+ host='thishostshouldnotexist',
+ password='password',
+ port=5439,
+ connect_timeout=10,
+ options="-c search_path=test",
+ keepalives_idle=RedshiftAdapter.ConnectionManager.DEFAULT_TCP_KEEPALIVE)
+
+ @mock.patch('dbt.adapters.postgres.connections.psycopg2')
+ def test_search_path_with_space(self, psycopg2):
+ self.config.credentials = self.config.credentials.incorporate(
+ search_path="test test"
+ )
+ connection = self.adapter.acquire_connection('dummy')
+
+ psycopg2.connect.assert_called_once_with(
+ dbname='redshift',
+ user='root',
+ host='thishostshouldnotexist',
+ password='password',
+ port=5439,
+ connect_timeout=10,
+ options="-c search_path=test\ test",
+ keepalives_idle=RedshiftAdapter.ConnectionManager.DEFAULT_TCP_KEEPALIVE)
+
  @mock.patch('dbt.adapters.postgres.connections.psycopg2')
  def test_set_zero_keepalive(self, psycopg2):
  self.config.credentials = self.config.credentials.incorporate(