Issue 473: PgServer missing -key option, patch from Andrew Franklin

h2/src/docsrc/html/changelog.html

@@ -36,6 +36,7 @@ Change Log
 </li><li>Improve error message when we have a unique constraint violation, displays the offending key in the error message
 </li><li>Issue 478: Support for "SHOW TRANSACTION ISOLATION LEVEL", patch from Andrew Franklin
 </li><li>Issue 475: PgServer: add support for CancelRequest, patch from Andrew Franklin
+</li><li>Issue 473: PgServer missing -key option, patch from Andrew Franklin
 </li></ul>
 
 <h2>Version 1.3.172 (2013-05-25)</h2>

h2/src/main/org/h2/server/pg/PgServer.java

@@ -20,6 +20,7 @@ import java.util.Collections;
 import java.util.HashSet;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
+import org.h2.constant.ErrorCode;
 import org.h2.engine.Constants;
 import org.h2.message.DbException;
 import org.h2.server.Service;
@@ -84,6 +85,7 @@ public class PgServer implements Service {
     private boolean allowOthers;
     private boolean isDaemon;
     private boolean ifExists;
+    private String key, keyDatabase;
 
     @Override
     public void init(String... args) {
@@ -103,6 +105,9 @@ public class PgServer implements Service {
                 isDaemon = true;
             } else if (Tool.isOption(a, "-ifExists")) {
                 ifExists = true;
+            } else if (Tool.isOption(a, "-key")) {
+                key = args[++i];
+                keyDatabase = args[++i];
             }
         }
         org.h2.Driver.load();
@@ -517,6 +522,25 @@ public class PgServer implements Service {
         }
     }
 
+    /**
+     * If no key is set, return the original database name. If a key is set,
+     * check if the key matches. If yes, return the correct database name. If
+     * not, throw an exception.
+     *
+     * @param db the key to test (or database name if no key is used)
+     * @return the database name
+     * @throws DbException if a key is set but doesn't match
+     */
+    public String checkKeyAndGetDatabaseName(String db) {
+        if (key == null) {
+            return db;
+        }
+        if (key.equals(db)) {
+            return keyDatabase;
+        }
+        throw DbException.get(ErrorCode.WRONG_USER_OR_PASSWORD);
+    }
+
     @Override
     public boolean isDaemon() {
         return isDaemon;

h2/src/main/org/h2/server/pg/PgServerThread.java

@@ -171,7 +171,7 @@ public class PgServerThread implements Runnable {
                     if ("user".equals(param)) {
                         this.userName = value;
                     } else if ("database".equals(param)) {
-                        this.databaseName = value;
+                        this.databaseName = server.checkKeyAndGetDatabaseName(value);
                     } else if ("client_encoding".equals(param)) {
                         // UTF8
                         clientEncoding = value;

h2/src/test/org/h2/test/unit/TestPgServer.java

@@ -37,6 +37,11 @@ public class TestPgServer extends TestBase {
 
     @Override
     public void test() throws SQLException {
+        testPGAdapter();
+        testKeyAlias();
+    }
+
+    private void testPGAdapter() throws SQLException {
         deleteDb("test");
         Server server = Server.createPgServer("-baseDir", getBaseDir(), "-pgPort", "5535", "-pgDaemon");
         assertEquals(5535, server.getPort());
@@ -259,4 +264,26 @@ public class TestPgServer extends TestBase {
         conn.close();
     }
 
+    private void testKeyAlias() throws SQLException {
+        Server server = Server.createPgServer("-pgPort", "5535", "-pgDaemon", "-key", "test", "mem:test");
+        server.start();
+        try {
+            Class.forName("org.postgresql.Driver");
+
+            Connection conn = DriverManager.getConnection("jdbc:postgresql://localhost:5535/test", "sa", "sa");
+            Statement stat = conn.createStatement();
+
+            // confirm that we've got the in memory implementation by creating a table and checking flags
+            stat.execute("create table test(id int primary key, name varchar)");
+            ResultSet rs = stat.executeQuery("select storage_type from information_schema.tables where table_name = 'TEST'");
+            assertTrue(rs.next());
+            assertEquals("MEMORY", rs.getString(1));
+
+            conn.close();
+        } catch (ClassNotFoundException e) {
+            println("PostgreSQL JDBC driver not found - PgServer not tested");
+        } finally {
+            server.stop();
+        }
+    }
 }