Support write mode in ZonalFile and override related methods

Author: suni72

gcsfs/zb_hns_utils.py

@@ -11,3 +11,18 @@ async def download_range(offset, length, mrd):
     buffer = BytesIO()
     await mrd.download_ranges([(offset, length, buffer)])
     return buffer.getvalue()
+
+
+async def init_aaow(grpc_client, bucket_name, object_name, generation=None):
+    """
+    Creates and opens the AsyncAppendableObjectWriter.
+    """
+
+    writer = AsyncAppendableObjectWriter(
+        client=grpc_client,
+        bucket_name=bucket_name,
+        object_name=object_name,
+        generation=generation,
+    )
+    await writer.open()
+    return writer

gcsfs/zonal_file.py

@@ -22,6 +22,10 @@ def __init__(self, *args, **kwargs):
             self.mrd = asyn.sync(
                 self.gcsfs.loop, self._init_mrd, self.bucket, self.key, self.generation
             )
+        elif "w" in self.mode:
+            self.aaow = asyn.sync(
+                self.gcsfs.loop, self._init_aaow, self.bucket, self.key, self.generation
+            )
         else:
             raise NotImplementedError(
                 "Only read operations are currently supported for Zonal buckets."
@@ -47,10 +51,76 @@ def _fetch_range(self, start, end):
                 return b""
             raise
 
+    async def _init_aaow(
+        self, bucket_name, object_name, generation=None, overwrite=True
+    ):
+        """
+        Initializes the AsyncAppendableObjectWriter.
+        """
+        if generation is None and await self.gcsfs._exists(self.path):
+            info = self.gcsfs.info(self.path)
+            generation = info.get("generation")
+
+        return await zb_hns_utils.init_aaow(
+            self.gcsfs.grpc_client, bucket_name, object_name, generation, overwrite
+        )
+
+    def flush(self, force=False):
+        raise NotImplementedError(
+            "Write operations are not yet implemented for Zonal buckets."
+        )
+
+    def commit(self):
+        raise NotImplementedError(
+            "Write operations are not yet implemented for Zonal buckets."
+        )
+
+    def discard(self):
+        raise NotImplementedError(
+            "Write operations are not yet implemented for Zonal buckets."
+        )
+
+    async def initiate_upload(
+        fs,
+        bucket,
+        key,
+        content_type="application/octet-stream",
+        metadata=None,
+        fixed_key_metadata=None,
+        mode="overwrite",
+        kms_key_name=None,
+    ):
+        raise NotImplementedError(
+            "Write operations are not yet implemented for Zonal buckets."
+        )
+
+    async def simple_upload(
+        fs,
+        bucket,
+        key,
+        datain,
+        metadatain=None,
+        consistency=None,
+        content_type="application/octet-stream",
+        fixed_key_metadata=None,
+        mode="overwrite",
+        kms_key_name=None,
+    ):
+        raise NotImplementedError(
+            "Write operations are not yet implemented for Zonal buckets."
+        )
+
+    async def upload_chunk(fs, location, data, offset, size, content_type):
+        raise NotImplementedError(
+            "Write operations are not yet implemented for Zonal buckets."
+        )
+
     def close(self):
         """
         Closes the ZonalFile and the underlying AsyncMultiRangeDownloader.
         """
         if self.mrd:
             asyn.sync(self.gcsfs.loop, self.mrd.close)
+        if self.aaow:
+            asyn.sync(self.gcsfs.loop, self.aaow.close)
         super().close()