add support to run scripts on usage DB

kishan · kishan · commit 5401ee84ba9d · 2011-04-05T16:11:02.000+05:30
diff --git a/utils/src/com/cloud/utils/db/Transaction.java b/utils/src/com/cloud/utils/db/Transaction.java
@@ -161,6 +161,15 @@ public static Connection getStandaloneConnection() {
 		}
     }
     
+    public static Connection getStandaloneUsageConnection() {
+        try {
+            return s_usageDS.getConnection();
+        } catch (SQLException e) {
+            s_logger.warn("Unexpected exception: ", e);
+            return null;
+        }
+    }
+    
     protected static boolean checkAnnotation(int stack, Transaction txn) {
         final StackTraceElement[] stacks = Thread.currentThread().getStackTrace();
         StackElement se = txn.peekInStack(CURRENT_TXN);
diff --git a/utils/test/com/cloud/utils/db/DbTestUtils.java b/utils/test/com/cloud/utils/db/DbTestUtils.java
@@ -58,4 +58,34 @@ public static void executeScript(String file, boolean autoCommit, boolean stopOn
         }
     }
 
+    public static void executeUsageScript(String file, boolean autoCommit, boolean stopOnError) {
+        File cleanScript = PropertiesUtil.findConfigFile(file);
+        if (cleanScript == null) {
+            throw new RuntimeException("Unable to clean the database because I can't find " + file);
+        }
+        
+        Connection conn = Transaction.getStandaloneUsageConnection();
+        
+        ScriptRunner runner = new ScriptRunner(conn, autoCommit, stopOnError);
+        FileReader reader;
+        try {
+            reader = new FileReader(cleanScript);
+        } catch (FileNotFoundException e) {
+            throw new RuntimeException("Unable to read " + file, e);
+        } 
+        try {
+            runner.runScript(reader);
+        } catch (IOException e) {
+            throw new RuntimeException("Unable to read " + file, e);
+        } catch (SQLException e) {
+            throw new RuntimeException("Unable to execute " + file, e);
+        }
+        
+        try {
+            conn.close();
+        } catch (SQLException e) {
+            throw new RuntimeException("Unable to close DB connection", e);
+        }
+    }
+
 }
