Merge pull request #466 from henrykironde/improve_bulk_insert_auto_co…

…mmit add execution for setting auto-comit off when bulk insert
weecology · Apr 25, 2016 · 840e92e · 840e92e
2 parents c3e161f + b810fbb
commit 840e92e
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/engines/mysql.py b/engines/mysql.py
@@ -45,6 +45,10 @@ def create_db_statement(self):
     def insert_data_from_file(self, filename):
         """Calls MySQL "LOAD DATA LOCAL INFILE" statement to perform a bulk
         insert."""
+
+        mysql_set_autocommit_off = """SET autocommit=0; SET UNIQUE_CHECKS=0; SET FOREIGN_KEY_CHECKS=0; SET sql_log_bin=0;"""
+        mysql_set_autocommit_on = """SET GLOBAL innodb_flush_log_at_trx_commit=1; COMMIT; SET autocommit=1; SET unique_checks=1; SET foreign_key_checks=1;"""
+
         self.get_cursor()
         ct = len([True for c in self.table.columns if c[1][0][:3] == "ct-"]) != 0
         if (self.table.cleanup.function == no_cleanup and
@@ -64,10 +68,13 @@ def insert_data_from_file(self, filename):
 IGNORE """ + str(self.table.header_rows) + """ LINES
 (""" + columns + ")"
             try:
+                self.cursor.execute(mysql_set_autocommit_off)
                 self.cursor.execute(statement)
+
+                self.cursor.execute(mysql_set_autocommit_on)
             except Exception as e:
-                print "Failed bulk insert (%s), inserting manually" % e
                 self.disconnect()  # If the execute fails the database connection can get hung up
+                self.cursor.execute(mysql_set_autocommit_on)
                 return Engine.insert_data_from_file(self, filename)
         else:
             return Engine.insert_data_from_file(self, filename)

diff --git a/engines/postgres.py b/engines/postgres.py
@@ -1,5 +1,5 @@
 import os
-import platform
+
 from retriever.lib.models import Engine, no_cleanup
 
 
@@ -72,11 +72,10 @@ def insert_data_from_file(self, filename):
         self.get_cursor()
         ct = len([True for c in self.table.columns if c[1][0][:3] == "ct-"]) != 0
         if (([self.table.cleanup.function, self.table.delimiter,
-             self.table.header_rows] == [no_cleanup, ",", 1])
+              self.table.header_rows] == [no_cleanup, ",", 1])
             and not self.table.fixed_width
             and not ct
-            and (not hasattr(self.table, "do_not_bulk_insert") or not self.table.do_not_bulk_insert)
-            ):
+            and (not hasattr(self.table, "do_not_bulk_insert") or not self.table.do_not_bulk_insert)):
             columns = self.table.get_insert_columns()
             filename = os.path.abspath(filename)
             statement = """
@@ -85,7 +84,9 @@ def insert_data_from_file(self, filename):
 WITH DELIMITER ','
 CSV HEADER"""
             try:
+                self.execute("BEGIN")
                 self.execute(statement)
+                self.execute("COMMIT")
             except:
                 self.connection.rollback()
                 return Engine.insert_data_from_file(self, filename)