add logic to route upload methods to zonal implementation

Author: suni72

gcsfs/core.py

@@ -1864,6 +1864,7 @@ def __init__(
         self.bucket = bucket
         self.key = key
         self.acl = acl
+        self.consistency = consistency
         self.checker = get_consistency_checker(consistency)
 
         if "a" in self.mode:
@@ -2073,6 +2074,18 @@ def _convert_fixed_key_metadata(metadata, *, from_google=False):
 
 
 async def upload_chunk(fs, location, data, offset, size, content_type):
+    from google.cloud.storage._experimental.asyncio.async_appendable_object_writer import (
+        AsyncAppendableObjectWriter,
+    )
+
+    from .extended_gcsfs import ExtendedGcsFileSystem
+    from .extended_gcsfs import upload_chunk as ext_upload_chunk
+
+    if isinstance(fs, ExtendedGcsFileSystem) and isinstance(
+        location, AsyncAppendableObjectWriter
+    ):
+
+        return await ext_upload_chunk(fs, location, data, offset, size, content_type)
     head = {}
     l = len(data)
     range = "bytes %i-%i/%i" % (offset, offset + l - 1, size)
@@ -2101,6 +2114,22 @@ async def initiate_upload(
     mode="overwrite",
     kms_key_name=None,
 ):
+    from .extended_gcsfs import ExtendedGcsFileSystem
+    from .extended_gcsfs import initiate_upload as ext_initiate_upload
+
+    if isinstance(fs, ExtendedGcsFileSystem) and await fs._is_zonal_bucket(bucket):
+
+        return await ext_initiate_upload(
+            fs,
+            bucket,
+            key,
+            content_type,
+            metadata,
+            fixed_key_metadata,
+            mode,
+            kms_key_name,
+        )
+
     j = {"name": key}
     if metadata:
         j["metadata"] = metadata
@@ -2135,6 +2164,24 @@ async def simple_upload(
     mode="overwrite",
     kms_key_name=None,
 ):
+    from .extended_gcsfs import ExtendedGcsFileSystem
+    from .extended_gcsfs import simple_upload as ext_simple_upload
+
+    if isinstance(fs, ExtendedGcsFileSystem) and await fs._is_zonal_bucket(bucket):
+
+        return await ext_simple_upload(
+            fs,
+            bucket,
+            key,
+            datain,
+            metadatain,
+            consistency,
+            content_type,
+            fixed_key_metadata,
+            mode,
+            kms_key_name,
+        )
+
     checker = get_consistency_checker(consistency)
     path = f"{fs._location}/upload/storage/v1/b/{quote(bucket)}/o"
     metadata = {"name": key}

gcsfs/extended_gcsfs.py

@@ -252,3 +252,41 @@ async def _cat_file(self, path, start=None, end=None, mrd=None, **kwargs):
             # Explicit cleanup if we created the MRD
             if mrd_created:
                 await mrd.close()
+
+
+async def upload_chunk(fs, location, data, offset, size, content_type):
+    raise NotImplementedError(
+        "upload_chunk is not implemented yet for ExtendedGcsFileSystem. Please use write() instead."
+    )
+
+
+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(
+        "initiate_upload is not implemented yet for ExtendedGcsFileSystem. Please use write() instead."
+    )
+
+
+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(
+        "simple_upload is not implemented yet for ExtendedGcsFileSystem. Please use write() instead."
+    )

gcsfs/zonal_file.py

@@ -75,10 +75,10 @@ def write(self, data):
         """
         Writes data using AsyncAppendableObjectWriter.
         """
-        if not self.writable():
-            raise ValueError("File not in write mode")
         if self.closed:
             raise ValueError("I/O operation on closed file.")
+        if not self.writable():
+            raise ValueError("File not in write mode")
         if self.forced:
             raise ValueError("This file has been force-flushed, can only close")
 
@@ -118,39 +118,49 @@ def discard(self):
             Data is uploaded via streaming and cannot be cancelled."
         )
 
-    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(
-            "Initiate_upload operation is not implemented yet for Zonal buckets. Please use write() instead."
+    def _initiate_upload(self):
+        """Initiates the upload for Zonal buckets using gRPC."""
+        from gcsfs.extended_gcsfs import initiate_upload
+
+        self.location = asyn.sync(
+            self.gcsfs.loop,
+            initiate_upload,
+            self.gcsfs,
+            self.bucket,
+            self.key,
+            self.content_type,
+            self.metadata,
+            self.fixed_key_metadata,
+            mode="create" if "x" in self.mode else "overwrite",
+            kms_key_name=self.kms_key_name,
+            timeout=self.timeout,
         )
 
-    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(
-            "Simple_upload operation is not implemented yet for Zonal buckets. Please use write() instead."
+    def _simple_upload(self):
+        """Performs a simple upload for Zonal buckets using gRPC."""
+        from gcsfs.extended_gcsfs import simple_upload
+
+        self.buffer.seek(0)
+        data = self.buffer.read()
+        asyn.sync(
+            self.gcsfs.loop,
+            simple_upload,
+            self.gcsfs,
+            self.bucket,
+            self.key,
+            data,
+            self.metadata,
+            self.consistency,
+            self.content_type,
+            self.fixed_key_metadata,
+            mode="create" if "x" in self.mode else "overwrite",
+            kms_key_name=self.kms_key_name,
+            timeout=self.timeout,
         )
 
-    async def upload_chunk(fs, location, data, offset, size, content_type):
+    def _upload_chunk(self, final=False):
         raise NotImplementedError(
-            "Upload_chunk operation is not implemented yet for Zonal buckets. Please use write() instead."
+            "_upload_chunk is not implemented yet for ZonalFile. Please use write() instead."
         )
 
     def close(self):