yugabyte
diff --git a/‎java/yb-ysql-conn-mgr/src/test/java/org/yb/ysqlconnmgr/TestSessionParameters.java‎
Lines changed: 160 additions & 30 deletions b/‎java/yb-ysql-conn-mgr/src/test/java/org/yb/ysqlconnmgr/TestSessionParameters.java‎
Lines changed: 160 additions & 30 deletions
@@ -15,8 +15,13 @@
 
 import static org.yb.AssertionWrappers.assertTrue;
 import static org.yb.AssertionWrappers.assertFalse;
+import static org.yb.AssertionWrappers.assertEquals;
 import static org.yb.AssertionWrappers.fail;
+import static org.yb.AssertionWrappers.assertThrows;
 
+import java.io.File;
+import java.io.IOException;
+import java.net.InetSocketAddress;
 import java.sql.Connection;
 import java.sql.ResultSet;
 import java.sql.Statement;
@@ -25,9 +30,11 @@
 import java.util.Set;
 import java.util.HashSet;
 import java.util.Arrays;
-
+import java.util.List;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.yb.client.TestUtils;
+import org.yb.util.ProcessUtil;
 import org.yb.minicluster.MiniYBClusterBuilder;
 import org.yb.pgsql.ConnectionEndpoint;
 import com.yugabyte.PGConnection;
@@ -118,8 +125,9 @@ private static enum ExceptionType {
       new ExceptionType[] {}),
     new SessionParameter("search_path", "\"$user\", public", "test",
       new ExceptionType[] {}),
+    // JDBC sends TimeZone in the startup packet, and it can't be overridden using startup options
     new SessionParameter("TimeZone", "UTC", "EST",
-      new ExceptionType[] {}),
+      new ExceptionType[] { ExceptionType.INVALID_STARTUP }),
     new SessionParameter("default_transaction_isolation", "read committed", "serializable",
       new ExceptionType[] { ExceptionType.YB_CACHE_MAPPING }),
     new SessionParameter("transaction_isolation", "read committed",
@@ -546,14 +554,14 @@ public void testStartupParameters() throws Exception {
       String parameterName = sp.parameterName;
       String expectedValue = sp.expectedValue;
 
-      // (GH#22248) (rbarigidad) Startup parameters with spaces are not parsed
-      // correctly, they are ignored for now as well.
+      String valueWithEscapedSpaces = expectedValue.replace(" ", "\\ ");
+
       if (sp.exceptionSet.contains(ExceptionType.INVALID_STARTUP))
         continue;
 
       try (Connection conn = getConnectionBuilder()
           .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
-          .withOptions(String.format("-c %s=%s", parameterName, expectedValue))
+          .withOptions(String.format("-c %s=%s", parameterName, valueWithEscapedSpaces))
           .connect();
           Statement stmt = conn.createStatement()) {
         String fetchedValue = fetchParameterValue(stmt, sp);
@@ -566,56 +574,178 @@ public void testStartupParameters() throws Exception {
 
   @Test
   public void testTimeZoneSetting() throws Exception {
-    try (Connection conn =
-        getConnectionBuilder().withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
-            .withOptions("-c TimeZone=Asia/Kolkata").connect()) {
+    try (Connection conn = getConnectionBuilder()
+                               .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                               .connect();
+         Statement stmt = conn.createStatement()) {
+      stmt.execute("SET TimeZone = 'Asia/Kolkata'");
       PGConnection pgConn = (PGConnection) conn;
       Map<String, String> params = new HashMap<>(pgConn.getParameterStatuses());
 
       // Get the timezone directly
       String backendTimeZone = params.get("TimeZone");
-      assertTrue("Backend timezone should be set correctly from startup packet",
+      assertTrue("Backend timezone should be set correctly from SET command",
           backendTimeZone.equals("Asia/Kolkata"));
     }
 
-    try (Connection conn =
-        getConnectionBuilder().withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
-            .withOptions("-c timezone=Asia/Kolkata").connect()) {
+    try (Connection conn = getConnectionBuilder()
+                               .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                               .connect();
+         Statement stmt = conn.createStatement()) {
+      stmt.execute("SET timezone = 'Asia/Kolkata'");
       PGConnection pgConn = (PGConnection) conn;
       Map<String, String> params = new HashMap<>(pgConn.getParameterStatuses());
 
       // Get the timezone directly
       String backendTimeZone = params.get("TimeZone");
-      assertTrue("Backend timezone should be set correctly from startup packet",
+      assertTrue("Backend timezone should be set correctly from SET command",
           backendTimeZone.equals("Asia/Kolkata"));
     }
+  }
 
+  public void testStartupParameterPrecedence() throws Exception {
+    // Test the precedence between guc setting specified through "options" key and one specified
+    // directly in the startup packet. pgJDBC driver only allows us to use the former method.
+    // See getParametersForStartup() function in pgJDBC driver code.
+    //
+    // To still test for the precedence, we try to set "client_encoding" through "options". The
+    // driver internally adds a (client_encoding, UTF8) pair to the startup packet and we test that
+    // that setting has a higher precedence
+    //
+    // Note: We don't test that setting an option through "options" field works, that is
+    // accomplished in testStartupParameters() above
+
+    SessionParameter sp =
+        new SessionParameter("client_encoding", "SQL_ASCII", "UTF8", new ExceptionType[] {});
+    String parameterName = sp.parameterName;
+    String expectedValue = sp.expectedValue;
+
+    try (Connection conn = getConnectionBuilder()
+                               .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                               .withOptions("-c client_encoding=SQL_ASCII")
+                               .connect();
+         Statement stmt = conn.createStatement()) {
+      String fetchedValue = fetchParameterValue(stmt, sp);
+      assertTrue(String.format("expected value %s for %s, but fetched %s instead", expectedValue,
+                     parameterName, fetchedValue),
+          expectedValue.equals(fetchedValue));
+    }
+  }
+
+  @Test
+  public void testStartupParameterCaseInsensitivity() throws Exception {
+    // Test that sending a startup parameter in a different case than what Postgres sends in
+    // ParameterStatus works correctly.
+    // This mainly tests that whatever state the connection manager stores is case-insensitive:
+    // either by lowercasing all GUC names in case of auth passthrough or forwarding the startup
+    // packet to auth backend and using the name in returned ParameterStatus packet.
+    SessionParameter sp =
+        new SessionParameter("default_transaction_read_only", "off", "on", new ExceptionType[] {});
+    String parameterName = sp.parameterName;
+    String expectedValue = sp.expectedValue;
+
+    try (Connection conn = getConnectionBuilder()
+                               .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                               .withOptions("-c DEFAULT_TRANSACTION_READ_ONLY=on")
+                               .connect();
+         Statement stmt = conn.createStatement()) {
+      String fetchedValue = fetchParameterValue(stmt, sp);
+
+      assertTrue(String.format("expected value %s for %s, but fetched %s instead", expectedValue,
+                     parameterName, fetchedValue),
+          expectedValue.equals(fetchedValue));
+    }
+  }
+
+  @Test
+  public void testInvalidStartupParameter() throws Exception {
+    // Test that setting an invalid parameter in startup packet fails at the time of database
+    // connection. We use ysqlsh here since we don't want to execute ANY query after connecting,
+    // whereas JDBC tries to execute some SET statements after connecting. With JDBC, we try to
+    // deploy the state to a transactional backend which fails and we're not able to distinguish
+    // it from connection failure.
+
+    // Sleep for some time to wait for the cluster to come up. We have to do this since this
+    // driver doesn't have any retry mechanism
+    Thread.sleep(5000);
+
+    File pgBinDir = new File(TestUtils.getBuildRootDir(), "postgres/bin");
+    File ysqlshExec = new File(pgBinDir, "ysqlsh");
+    final InetSocketAddress postgresAddress = miniCluster.getYsqlConnMgrContactPoints().get(0);
+
+    List<String> args = Arrays.asList(
+      ysqlshExec.toString(),
+      "-h", postgresAddress.getHostName(),
+      "-p", Integer.toString(postgresAddress.getPort()),
+      "-U", "yugabyte",
+      "-v", "ON_ERROR_STOP=1",
+      "-c", "\\q"
+    );
+
+
+    Map<String, String> invalidEnv = new HashMap<String, String>();
+    invalidEnv.put("PGOPTIONS", "-c geqo=abc");
+    assertThrows(
+        IOException.class, () -> { ProcessUtil.runProcess(args, Integer.MAX_VALUE, invalidEnv); });
+
+    Map<String, String> validEnv = new HashMap<String, String>();
+    invalidEnv.put("PGOPTIONS", "-c geqo=off");
+    ProcessUtil.runProcess(args, Integer.MAX_VALUE, validEnv);
+  }
+
+  @Test
+  public void testSettingOverrideVariableInStartupPacket() throws Exception {
+    // Startup parameters are set with the source PGC_S_CLIENT while SET statements are set
+    // with the source PGC_S_SESSION. There can be settings with PGC_S_OVERRIDE which would
+    // not be overridden with startup options but would get overriden with SET statements.
+    // We test this fact with the GUC session_authorization.
+
+    final String roleName = "test_user_session_auth";
+
+    // Connect as yugabyte user to yugabyte db
     try (
-        Connection conn = getConnectionBuilder()
-            .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR).connect();
+        Connection conn =
+            getConnectionBuilder().withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                .withUser("yugabyte").withDatabase("yugabyte").connect();
         Statement stmt = conn.createStatement()) {
-      stmt.execute("SET TimeZone = 'Asia/Kolkata'");
-      PGConnection pgConn = (PGConnection) conn;
-      Map<String, String> params = new HashMap<>(pgConn.getParameterStatuses());
+      stmt.execute(String.format("CREATE ROLE %s WITH login", roleName));
+    }
 
-      // Get the timezone directly
-      String backendTimeZone = params.get("TimeZone");
-      assertTrue("Backend timezone should be set correctly from SET command",
-          backendTimeZone.equals("Asia/Kolkata"));
+    // Connect as test user to yugabyte db
+    try (
+        Connection connTestUser =
+            getConnectionBuilder().withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                .withUser(roleName).withDatabase("yugabyte").connect();
+        Statement stmtTestUser = connTestUser.createStatement()) {
+      try (ResultSet rs = stmtTestUser.executeQuery("SELECT 1")) {
+        assertTrue("SELECT 1 should return a row", rs.next());
+      }
     }
 
+    // Execute a global DDL as a workaround for role not applying when startup packet is processed
+    try (
+        Connection conn =
+            getConnectionBuilder().withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR)
+                .withUser("yugabyte").withDatabase("yugabyte").connect();
+        Statement stmt = conn.createStatement()) {
+      stmt.execute(
+          "create role tmp1; grant tmp1 to yugabyte; revoke tmp1 from yugabyte; drop role tmp1;");
+    }
+
+    // Try to set session authorization in startup packet
+    String sessionAuthOption = String.format("-c session_authorization=%s", roleName);
+
     try (
         Connection conn = getConnectionBuilder()
-            .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR).connect();
+            .withConnectionEndpoint(ConnectionEndpoint.YSQL_CONN_MGR).withUser("yugabyte")
+            .withDatabase("yugabyte").withOptions(sessionAuthOption).connect();
         Statement stmt = conn.createStatement()) {
-      stmt.execute("SET timezone = 'Asia/Kolkata'");
-      PGConnection pgConn = (PGConnection) conn;
-      Map<String, String> params = new HashMap<>(pgConn.getParameterStatuses());
 
-      // Get the timezone directly
-      String backendTimeZone = params.get("TimeZone");
-      assertTrue("Backend timezone should be set correctly from SET command",
-          backendTimeZone.equals("Asia/Kolkata"));
+      try (ResultSet rs = stmt.executeQuery("SHOW SESSION AUTHORIZATION")) {
+        assertTrue("SHOW SESSION AUTHORIZATION should return a row", rs.next());
+        String actualSessionAuth = rs.getString(1);
+        assertEquals("yugabyte", actualSessionAuth);
+      }
     }
   }
 }