@@ -267,7 +267,7 @@ def _redshift_types_from_path(
267
267
return redshift_types
268
268
269
269
270
- def _create_table ( # pylint: disable=too-many-locals,too-many-arguments
270
+ def _create_table ( # pylint: disable=too-many-locals,too-many-arguments,too-many-branches,too-many-statements
271
271
df : Optional [pd .DataFrame ],
272
272
path : Optional [Union [str , List [str ]]],
273
273
con : redshift_connector .Connection ,
@@ -292,6 +292,7 @@ def _create_table( # pylint: disable=too-many-locals,too-many-arguments
292
292
use_threads : Union [bool , int ] = True ,
293
293
boto3_session : Optional [boto3 .Session ] = None ,
294
294
s3_additional_kwargs : Optional [Dict [str , str ]] = None ,
295
+ lock : bool = False ,
295
296
) -> Tuple [str , Optional [str ]]:
296
297
if mode == "overwrite" :
297
298
if overwrite_method == "truncate" :
@@ -306,14 +307,21 @@ def _create_table( # pylint: disable=too-many-locals,too-many-arguments
306
307
_logger .debug (str (e ))
307
308
con .rollback ()
308
309
_begin_transaction (cursor = cursor )
310
+ if lock :
311
+ _lock (cursor , [table ], schema = schema )
309
312
elif overwrite_method == "delete" :
310
313
if _does_table_exist (cursor = cursor , schema = schema , table = table ):
314
+ if lock :
315
+ _lock (cursor , [table ], schema = schema )
311
316
# Atomic, but slow.
312
317
_delete_all (cursor = cursor , schema = schema , table = table )
313
318
else :
314
319
# Fast, atomic, but either fails if there are any dependent views or, in cascade mode, deletes them.
315
320
_drop_table (cursor = cursor , schema = schema , table = table , cascade = bool (overwrite_method == "cascade" ))
321
+ # No point in locking here, the oid will change.
316
322
elif _does_table_exist (cursor = cursor , schema = schema , table = table ) is True :
323
+ if lock :
324
+ _lock (cursor , [table ], schema = schema )
317
325
if mode == "upsert" :
318
326
guid : str = uuid .uuid4 ().hex
319
327
temp_table : str = f"temp_redshift_{ guid } "
@@ -378,6 +386,8 @@ def _create_table( # pylint: disable=too-many-locals,too-many-arguments
378
386
)
379
387
_logger .debug ("Create table query:\n %s" , sql )
380
388
cursor .execute (sql )
389
+ if lock :
390
+ _lock (cursor , [table ], schema = schema )
381
391
return table , schema
382
392
383
393
@@ -889,6 +899,7 @@ def to_sql( # pylint: disable=too-many-locals
889
899
primary_keys = primary_keys ,
890
900
varchar_lengths_default = varchar_lengths_default ,
891
901
varchar_lengths = varchar_lengths ,
902
+ lock = lock ,
892
903
)
893
904
if index :
894
905
df .reset_index (level = df .index .names , inplace = True )
@@ -905,8 +916,6 @@ def to_sql( # pylint: disable=too-many-locals
905
916
_logger .debug ("sql: %s" , sql )
906
917
cursor .executemany (sql , (parameters ,))
907
918
if table != created_table : # upsert
908
- if lock :
909
- _lock (cursor , [table ], schema = schema )
910
919
_upsert (
911
920
cursor = cursor ,
912
921
schema = schema ,
@@ -1385,10 +1394,8 @@ def copy_from_files( # pylint: disable=too-many-locals,too-many-arguments
1385
1394
use_threads = use_threads ,
1386
1395
boto3_session = boto3_session ,
1387
1396
s3_additional_kwargs = s3_additional_kwargs ,
1397
+ lock = lock ,
1388
1398
)
1389
- if lock and table == created_table :
1390
- # Lock before copy if copying into target (not temp) table
1391
- _lock (cursor , [table ], schema = schema )
1392
1399
_copy (
1393
1400
cursor = cursor ,
1394
1401
path = path ,
@@ -1404,8 +1411,6 @@ def copy_from_files( # pylint: disable=too-many-locals,too-many-arguments
1404
1411
manifest = manifest ,
1405
1412
)
1406
1413
if table != created_table : # upsert
1407
- if lock :
1408
- _lock (cursor , [table ], schema = schema )
1409
1414
_upsert (
1410
1415
cursor = cursor ,
1411
1416
schema = schema ,
0 commit comments