Fix Redshift Locking Behavior

awswrangler/redshift.py

@@ -267,7 +267,7 @@ def _redshift_types_from_path(
     return redshift_types
 
 
-def _create_table(  # pylint: disable=too-many-locals,too-many-arguments
+def _create_table(  # pylint: disable=too-many-locals,too-many-arguments,too-many-branches,too-many-statements
     df: Optional[pd.DataFrame],
     path: Optional[Union[str, List[str]]],
     con: redshift_connector.Connection,
@@ -292,6 +292,7 @@ def _create_table(  # pylint: disable=too-many-locals,too-many-arguments
     use_threads: Union[bool, int] = True,
     boto3_session: Optional[boto3.Session] = None,
     s3_additional_kwargs: Optional[Dict[str, str]] = None,
+    lock: bool = False,
 ) -> Tuple[str, Optional[str]]:
     if mode == "overwrite":
         if overwrite_method == "truncate":
@@ -306,14 +307,21 @@ def _create_table(  # pylint: disable=too-many-locals,too-many-arguments
                 _logger.debug(str(e))
                 con.rollback()
             _begin_transaction(cursor=cursor)
+            if lock:
+                _lock(cursor, [table], schema=schema)
         elif overwrite_method == "delete":
             if _does_table_exist(cursor=cursor, schema=schema, table=table):
+                if lock:
+                    _lock(cursor, [table], schema=schema)
                 # Atomic, but slow.
                 _delete_all(cursor=cursor, schema=schema, table=table)
         else:
             # Fast, atomic, but either fails if there are any dependent views or, in cascade mode, deletes them.
             _drop_table(cursor=cursor, schema=schema, table=table, cascade=bool(overwrite_method == "cascade"))
+            # No point in locking here, the oid will change.
     elif _does_table_exist(cursor=cursor, schema=schema, table=table) is True:
+        if lock:
+            _lock(cursor, [table], schema=schema)
         if mode == "upsert":
             guid: str = uuid.uuid4().hex
             temp_table: str = f"temp_redshift_{guid}"
@@ -378,6 +386,8 @@ def _create_table(  # pylint: disable=too-many-locals,too-many-arguments
     )
     _logger.debug("Create table query:\n%s", sql)
     cursor.execute(sql)
+    if lock:
+        _lock(cursor, [table], schema=schema)
     return table, schema
 
 
@@ -889,6 +899,7 @@ def to_sql(  # pylint: disable=too-many-locals
                 primary_keys=primary_keys,
                 varchar_lengths_default=varchar_lengths_default,
                 varchar_lengths=varchar_lengths,
+                lock=lock,
             )
             if index:
                 df.reset_index(level=df.index.names, inplace=True)
@@ -905,8 +916,6 @@ def to_sql(  # pylint: disable=too-many-locals
                 _logger.debug("sql: %s", sql)
                 cursor.executemany(sql, (parameters,))
             if table != created_table:  # upsert
-                if lock:
-                    _lock(cursor, [table], schema=schema)
                 _upsert(
                     cursor=cursor,
                     schema=schema,
@@ -1385,10 +1394,8 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
                 use_threads=use_threads,
                 boto3_session=boto3_session,
                 s3_additional_kwargs=s3_additional_kwargs,
+                lock=lock,
             )
-            if lock and table == created_table:
-                # Lock before copy if copying into target (not temp) table
-                _lock(cursor, [table], schema=schema)
             _copy(
                 cursor=cursor,
                 path=path,
@@ -1404,8 +1411,6 @@ def copy_from_files(  # pylint: disable=too-many-locals,too-many-arguments
                 manifest=manifest,
             )
             if table != created_table:  # upsert
-                if lock:
-                    _lock(cursor, [table], schema=schema)
                 _upsert(
                     cursor=cursor,
                     schema=schema,