suuport host/port arguments for postgres

abcdabcd3899 · abcdabcd3899 · commit b7501c7a8606 · 2021-03-09T12:15:04.000+08:00
diff --git a/src/sqlancer/MainOptions.java b/src/sqlancer/MainOptions.java
@@ -49,6 +49,13 @@ public class MainOptions {
     @Parameter(names = "--password", description = "The password used to log into the DBMS")
     private String password = "sqlancer"; // NOPMD
 
+    @Parameter(names = "--host", description = "The host used to log into the DBMS")
+    private String host = "sqlancer"; // NOPMD
+
+    @Parameter(names = "--port", description = "The port used to log into the DBMS")
+    private String port = "sqlancer"; // NOPMD
+
+
     @Parameter(names = "--print-progress-information", description = "Whether to print progress information such as the number of databases generated or queries issued", arity = 1)
     private boolean printProgressInformation = true; // NOPMD
 
@@ -151,6 +158,14 @@ public String getPassword() {
         return password;
     }
 
+    public String getHost(){
+        return host;
+    }
+
+    public String getPort(){
+        return port;
+    }
+
     public boolean printProgressInformation() {
         return printProgressInformation;
     }
diff --git a/src/sqlancer/postgres/PostgresProvider.java b/src/sqlancer/postgres/PostgresProvider.java
@@ -54,6 +54,8 @@ public class PostgresProvider extends SQLProviderAdapter<PostgresGlobalState, Po
     protected String password;
     protected String entryPath;
     protected String host;
+    protected String port;
+    protected String Scheme;
     protected String testURL;
     protected String databaseName;
     protected String createDatabaseCommand;
@@ -200,12 +202,17 @@ public SQLConnection createDatabase(PostgresGlobalState globalState) throws SQLE
 
         username = globalState.getOptions().getUserName();
         password = globalState.getOptions().getPassword();
+        host = globalState.getOptions().getHost();
+	port = globalState.getOptions().getPort();
         entryPath = "/test";
         entryURL = globalState.getDmbsSpecificOptions().connectionURL;
         // trim URL to exclude "jdbc:"
         if (entryURL.startsWith("jdbc:")) {
             entryURL = entryURL.substring(5);
         }
+	String entryDatabaseName = entryPath.substring(1);
+        databaseName = globalState.getDatabaseName();
+
         try {
             URI uri = new URI(entryURL);
             String userInfoURI = uri.getUserInfo();
@@ -228,12 +235,16 @@ public SQLConnection createDatabase(PostgresGlobalState globalState) throws SQLE
             if (pathURI != null) {
                 entryPath = pathURI;
             }
-            host = uri.getHost();
+	    if ("sqlancer".equals(host)){
+	        host = uri.getHost();
+	    }
+	    if("sqlancer".equals(port)){
+	        port = Integer.toString(uri.getPort());
+	    }
+            entryURL = uri.getScheme() + "://" + host + ":" + port + "/" + entryDatabaseName;
         } catch (URISyntaxException e) {
             throw new AssertionError(e);
         }
-        String entryDatabaseName = entryPath.substring(1);
-        databaseName = globalState.getDatabaseName();
         Connection con = DriverManager.getConnection("jdbc:" + entryURL, username, password);
         globalState.getState().logStatement(String.format("\\c %s;", entryDatabaseName));
         globalState.getState().logStatement("DROP DATABASE IF EXISTS " + databaseName);
@@ -251,6 +262,7 @@ public SQLConnection createDatabase(PostgresGlobalState globalState) throws SQLE
         String postDatabaseName = entryURL.substring(databaseIndex + entryDatabaseName.length());
         testURL = preDatabaseName + databaseName + postDatabaseName;
         globalState.getState().logStatement(String.format("\\c %s;", databaseName));
+	
         con = DriverManager.getConnection("jdbc:" + testURL, username, password);
         return new SQLConnection(con);
     }
