SQOOP-3105: Add cleanup logic for PostgreSQL related test cases

(Boglarka Egyed via Attila Szabo)

src/test/com/cloudera/sqoop/manager/PostgresqlExportTest.java

@@ -73,6 +73,10 @@ protected boolean useHsqldbTestServer() {
     return false;
   }
 
+  private String getDropTableStatement(String tableName, String schema) {
+    return "DROP TABLE IF EXISTS " + quoteTableOrSchemaName(schema) + "." + quoteTableOrSchemaName(tableName);
+  }
+
   @Before
   public void setUp() {
     super.setUp();
@@ -98,6 +102,16 @@ public void setUp() {
 
   @Override
   public void tearDown() {
+    try {
+      Statement stmt = connection.createStatement();
+      stmt.executeUpdate(getDropTableStatement(TABLE_NAME, SCHEMA_PUBLIC));
+      stmt.executeUpdate(getDropTableStatement(STAGING_TABLE_NAME, SCHEMA_PUBLIC));
+      stmt.executeUpdate(getDropTableStatement(TABLE_NAME, SCHEMA_SPECIAL));
+      stmt.executeUpdate(getDropTableStatement(STAGING_TABLE_NAME, SCHEMA_SPECIAL));
+    } catch(SQLException e) {
+      LOG.error("Can't clean up the database:", e);
+    }
+
     super.tearDown();
 
     try {
@@ -150,9 +164,9 @@ public void createIt(
           + "AS $$ "
           + "BEGIN "
           + "INSERT INTO "
-          + escapeTableOrSchemaName(SCHEMA_PUBLIC)
+          + quoteTableOrSchemaName(SCHEMA_PUBLIC)
           + "."
-          + escapeTableOrSchemaName(TABLE_NAME)
+          + quoteTableOrSchemaName(TABLE_NAME)
           + " ("
           + manager.escapeColName("id")
           +", "
@@ -199,7 +213,7 @@ private void create(
       // Create schema if not exists in dummy way (always create and ignore
       // errors.
       try {
-        st.executeUpdate("CREATE SCHEMA " + escapeTableOrSchemaName(schema));
+        st.executeUpdate("CREATE SCHEMA " + quoteTableOrSchemaName(schema));
         connection.commit();
       } catch (SQLException e) {
         LOG.info("Couldn't create schema " + schema + " (is o.k. as long as"
@@ -207,8 +221,8 @@ private void create(
         connection.rollback();
       }
 
-      String fullTableName = escapeTableOrSchemaName(schema)
+      String fullTableName = quoteTableOrSchemaName(schema)
-        + "." + escapeTableOrSchemaName(name);
+        + "." + quoteTableOrSchemaName(name);
 
       try {
         // Try to remove the table first. DROP TABLE IF EXISTS didn't
@@ -304,7 +318,7 @@ public void testExport() throws IOException, SQLException {
 
     runExport(getArgv(true));
 
-    assertRowCount(2, escapeTableOrSchemaName(TABLE_NAME), connection);
+    assertRowCount(2, quoteTableOrSchemaName(TABLE_NAME), connection);
   }
 
   @Test
@@ -316,7 +330,7 @@ public void testExportUsingProcedure() throws IOException, SQLException {
 
     runExport(getArgv(false));
 
-    assertRowCount(2, escapeTableOrSchemaName(TABLE_NAME), connection);
+    assertRowCount(2, quoteTableOrSchemaName(TABLE_NAME), connection);
   }
 
   @Test
@@ -330,7 +344,7 @@ public void testExportStaging() throws IOException, SQLException {
 
     runExport(getArgv(true, extra));
 
-    assertRowCount(2, escapeTableOrSchemaName(TABLE_NAME), connection);
+    assertRowCount(2, quoteTableOrSchemaName(TABLE_NAME), connection);
   }
 
   @Test
@@ -344,7 +358,7 @@ public void testExportDirect() throws IOException, SQLException {
 
     runExport(getArgv(true, extra));
 
-    assertRowCount(2, escapeTableOrSchemaName(TABLE_NAME), connection);
+    assertRowCount(2, quoteTableOrSchemaName(TABLE_NAME), connection);
   }
 
   @Test
@@ -362,8 +376,8 @@ public void testExportCustomSchema() throws IOException, SQLException {
     runExport(getArgv(true, extra));
 
     assertRowCount(2,
-      escapeTableOrSchemaName(SCHEMA_SPECIAL)
+      quoteTableOrSchemaName(SCHEMA_SPECIAL)
-        + "." + escapeTableOrSchemaName(TABLE_NAME),
+        + "." + quoteTableOrSchemaName(TABLE_NAME),
       connection);
   }
 
@@ -385,8 +399,8 @@ public void testExportCustomSchemaStaging() throws IOException, SQLException {
     runExport(getArgv(true, extra));
 
     assertRowCount(2,
-      escapeTableOrSchemaName(SCHEMA_SPECIAL)
+      quoteTableOrSchemaName(SCHEMA_SPECIAL)
-        + "." + escapeTableOrSchemaName(TABLE_NAME),
+        + "." + quoteTableOrSchemaName(TABLE_NAME),
       connection);
   }
 
@@ -410,8 +424,8 @@ public void testExportCustomSchemaStagingClear()
     runExport(getArgv(true, extra));
 
     assertRowCount(2,
-      escapeTableOrSchemaName(SCHEMA_SPECIAL)
+      quoteTableOrSchemaName(SCHEMA_SPECIAL)
-        + "." + escapeTableOrSchemaName(TABLE_NAME),
+        + "." + quoteTableOrSchemaName(TABLE_NAME),
       connection);
   }
 
@@ -432,8 +446,8 @@ public void testExportCustomSchemaDirect() throws IOException, SQLException {
     runExport(getArgv(true, extra));
 
     assertRowCount(2,
-      escapeTableOrSchemaName(SCHEMA_SPECIAL)
+      quoteTableOrSchemaName(SCHEMA_SPECIAL)
-        + "." + escapeTableOrSchemaName(TABLE_NAME),
+        + "." + quoteTableOrSchemaName(TABLE_NAME),
       connection);
   }
 
@@ -468,7 +482,7 @@ public static void assertRowCount(long expected,
     }
   }
 
-  public String escapeTableOrSchemaName(String tableName) {
+  public String quoteTableOrSchemaName(String tableName) {
     return "\"" + tableName + "\"";
   }
 }

src/test/com/cloudera/sqoop/manager/PostgresqlImportTest.java

@@ -31,6 +31,7 @@
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -109,11 +110,21 @@ public class PostgresqlImportTest extends ImportJobTestCase {
   static final String SCHEMA_SPECIAL = "special";
   static final String CONNECT_STRING = HOST_URL + DATABASE_NAME;
 
+  protected Connection connection;
+
   @Override
   protected boolean useHsqldbTestServer() {
     return false;
   }
 
+  public String quoteTableOrSchemaName(String tableName) {
+    return "\"" + tableName + "\"";
+  }
+
+  private String getDropTableStatement(String tableName, String schema) {
+    return "DROP TABLE IF EXISTS " + quoteTableOrSchemaName(schema) + "." + quoteTableOrSchemaName(tableName);
+  }
+
   @Before
   public void setUp() {
     super.setUp();
@@ -128,13 +139,35 @@ public void setUp() {
     LOG.debug("setUp complete.");
   }
 
+  @After
+  public void tearDown() {
+    try {
+      Statement stmt = connection.createStatement();
+      stmt.executeUpdate(getDropTableStatement(TABLE_NAME, SCHEMA_PUBLIC));
+      stmt.executeUpdate(getDropTableStatement(NULL_TABLE_NAME, SCHEMA_PUBLIC));
+      stmt.executeUpdate(getDropTableStatement(SPECIAL_TABLE_NAME, SCHEMA_PUBLIC));
+      stmt.executeUpdate(getDropTableStatement(DIFFERENT_TABLE_NAME, SCHEMA_SPECIAL));
+    } catch (SQLException e) {
+      LOG.error("Can't clean up the database:", e);
+    }
+
+    super.tearDown();
+
+    try {
+      connection.close();
+    } catch (SQLException e) {
+      LOG.error("Ignoring exception in tearDown", e);
+    }
+  }
+
+
+
   public void setUpData(String tableName, String schema, boolean nullEntry) {
     SqoopOptions options = new SqoopOptions(CONNECT_STRING, tableName);
     options.setUsername(DATABASE_USER);
     options.setPassword(PASSWORD);
 
     ConnManager manager = null;
-    Connection connection = null;
     Statement st = null;
 
     try {