Handle exceptions during writes to Azure

smart_open/azure.py

@@ -417,6 +417,18 @@ def __init__(
     def flush(self):
         pass
 
+    def terminate(self):
+        """Do not commit block list on abort.
+
+        Uploaded (uncommitted) blocks will be garbage collected after 7 days.
+
+        See also https://stackoverflow.com/a/69673084/5511061."""
+        logger.debug("terminating")
+        if not self.closed:
+            self._block_list = []
+            self._is_closed = True
+        logger.debug("successfully terminated")
+
     #
     # Override some methods from io.IOBase.
     #
@@ -509,7 +521,10 @@ def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        self.close()
+        if exc_type is not None:
+            self.terminate()
+        else:
+            self.close()
 
     def __str__(self):
         return "(%s, %r, %r)" % (

smart_open/tests/test_azure.py

@@ -58,6 +58,7 @@ def commit_block_list(self, block_list, metadata=None):
         metadata.update({"size": len(data)})
         self.set_blob_metadata(metadata)
         self._container_client.register_blob_client(self)
+        self._staged_contents = {}
 
     def delete_blob(self):
         self._container_client.delete_blob(self)
@@ -618,6 +619,26 @@ def test_write_blob_client(self):
         ))
         self.assertEqual(output, [test_string])
 
+    def test_abort_upload(self):
+        """Does aborted upload skip commit_block_list?"""
+        test_string = b"42" * 42
+        blob_name = "test_abort_upload_%s" % BLOB_NAME
+
+        container_client = CLIENT.get_container_client(CONTAINER_NAME)
+        blob_client = container_client.get_blob_client(blob_name)
+
+        try:
+            with smart_open.open(
+                "azure://%s/%s" % (CONTAINER_NAME, blob_name),
+                "wb",
+                transport_params={"client": blob_client, "min_part_size": 42},
+            ) as fout:
+                fout.write(test_string)
+                raise ValueError
+        except ValueError:
+            # FakeBlobClient.commit_block_list was not called
+            self.assertGreater(len(blob_client._staged_contents), 0)
+
     def test_incorrect_input(self):
         """Does azure write fail on incorrect input?"""
         blob_name = "test_incorrect_input_%s" % BLOB_NAME