Prevent race conditions by adding a random suffix to snowflake stage (#…

…166) * Prevent race conditions by adding a random suffix to snowflake stage * Bump version * Use UUID instead
instacart · May 5, 2020 · f4789b1 · f4789b1
1 parent 575b5ce
commit f4789b1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 4 deletions.
diff --git a/lore/__init__.py b/lore/__init__.py
@@ -17,7 +17,7 @@
 __copyright__ = 'Copyright © 2018, Instacart'
 __credits__ = ['Montana Low', 'Jeremy Stanley', 'Emmanuel Turlay', 'Shrikar Archak', 'Ganesh Krishnan']
 __license__ = 'MIT'
-__version__ = '0.8.2'
+__version__ = '0.8.3'
 __maintainer__ = 'Montana Low'
 __email__ = '[email protected]'
 __status__ = 'Development Status :: 4 - Beta'

diff --git a/lore/io/connection.py b/lore/io/connection.py
@@ -11,6 +11,7 @@
 import sys
 import tempfile
 import threading
+import uuid
 
 from datetime import datetime
 
@@ -188,15 +189,20 @@ def insert(self, table, dataframe, batch_size=10 ** 5):
                     tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.csv.gz')
                     tmp.close()
                     batch.to_csv(tmp.name, index=False, header=False, sep='|', na_rep='\\N', quoting=csv.QUOTE_NONE, compression='gzip')
-                    self._connection.connection.cursor().execute('PUT file://%(path)s @~/staged;' % {'path': tmp.name})
+                    suffix = str(uuid.uuid4())
+                    stage_name = 'staged_' + suffix
+                    self.execute('REMOVE @~/%(stage_name)s' % {'stage_name': stage_name})
+                    self._connection.connection.cursor().execute('PUT file://%(path)s @~/%(stage_name)s;' % {'path': tmp.name, 'stage_name': stage_name})
                     self._connection.connection.cursor().execute(
                         'COPY INTO %(table)s '
-                        'FROM @~/staged '
+                        'FROM @~/%(stage_name)s '
                         'FILE_FORMAT = (TYPE = CSV FIELD_DELIMITER = \'|\' SKIP_HEADER = 0 COMPRESSION = GZIP) '
                         'PURGE = TRUE' % {
-                            'table': table
+                            'table': table,
+                            'stage_name': stage_name
                         }
                     )
+                    self.execute('REMOVE @~/%(stage_name)s' % {'stage_name': stage_name})
                 finally:
                     os.remove(tmp.name)