From 15b19a0a5713eec7cc740f5725df971f407d73e4 Mon Sep 17 00:00:00 2001
From: bojanserafimov <bojan.serafimov7@gmail.com>
Date: Wed, 9 Mar 2022 14:47:06 -0500
Subject: [PATCH] [proxy] Test connstr options (#1344)

* Add proxy test
* Fix typo
---
 test_runner/batch_others/test_proxy.py  | 13 +++++++++++
 test_runner/fixtures/zenith_fixtures.py | 29 ++++++++++++++++++-------
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/test_runner/batch_others/test_proxy.py b/test_runner/batch_others/test_proxy.py
index 9510e880b2..d2039f9758 100644
--- a/test_runner/batch_others/test_proxy.py
+++ b/test_runner/batch_others/test_proxy.py
@@ -1,2 +1,15 @@
+import pytest
+
+
 def test_proxy_select_1(static_proxy):
     static_proxy.safe_psql("select 1;")
+
+
+@pytest.mark.xfail  # Proxy eats the extra connection options
+def test_proxy_options(static_proxy):
+    schema_name = "tmp_schema_1"
+    with static_proxy.connect(schema=schema_name) as conn:
+        with conn.cursor() as cur:
+            cur.execute("SHOW search_path;")
+            search_path = cur.fetchall()[0][0]
+            assert schema_name == search_path
diff --git a/test_runner/fixtures/zenith_fixtures.py b/test_runner/fixtures/zenith_fixtures.py
index 252ca9b3c1..4d6e84048c 100644
--- a/test_runner/fixtures/zenith_fixtures.py
+++ b/test_runner/fixtures/zenith_fixtures.py
@@ -242,15 +242,20 @@ class PgProtocol:
                  host: str,
                  port: int,
                  username: Optional[str] = None,
-                 password: Optional[str] = None):
+                 password: Optional[str] = None,
+                 dbname: Optional[str] = None,
+                 schema: Optional[str] = None):
         self.host = host
         self.port = port
         self.username = username
         self.password = password
+        self.dbname = dbname
+        self.schema = schema
 
     def connstr(self,
                 *,
-                dbname: str = 'postgres',
+                dbname: Optional[str] = None,
+                schema: Optional[str] = None,
                 username: Optional[str] = None,
                 password: Optional[str] = None) -> str:
         """
@@ -259,6 +264,8 @@ class PgProtocol:
 
         username = username or self.username
         password = password or self.password
+        dbname = dbname or self.dbname or "postgres"
+        schema = schema or self.schema
         res = f'host={self.host} port={self.port} dbname={dbname}'
 
         if username:
@@ -267,13 +274,17 @@ class PgProtocol:
         if password:
             res = f'{res} password={password}'
 
+        if schema:
+            res = f"{res} options='-c search_path={schema}'"
+
         return res
 
     # autocommit=True here by default because that's what we need most of the time
     def connect(self,
                 *,
                 autocommit=True,
-                dbname: str = 'postgres',
+                dbname: Optional[str] = None,
+                schema: Optional[str] = None,
                 username: Optional[str] = None,
                 password: Optional[str] = None) -> PgConnection:
         """
@@ -282,11 +293,13 @@ class PgProtocol:
         This method passes all extra params to connstr.
         """
 
-        conn = psycopg2.connect(self.connstr(
-            dbname=dbname,
-            username=username,
-            password=password,
-        ))
+        conn = psycopg2.connect(
+            self.connstr(
+                dbname=dbname,
+                schema=schema,
+                username=username,
+                password=password,
+            ))
         # WARNING: this setting affects *all* tests!
         conn.autocommit = autocommit
         return conn