| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -21,7 +21,7 @@ | |||
| 21 | 21 | if you want to use these Rapid Storage APIs. | |
| 22 | 22 | ||
| 23 | 23 | """ | |
| 24 | - from typing import Optional | ||
| 24 | + from typing import Optional, Union | ||
| 25 | 25 | from google.cloud import _storage_v2 | |
| 26 | 26 | from google.cloud.storage._experimental.asyncio.async_grpc_client import ( | |
| 27 | 27 | AsyncGrpcClient, | |
@@ -31,6 +31,10 @@ | |||
| 31 | 31 | ) | |
| 32 | 32 | ||
| 33 | 33 | ||
| 34 | + _MAX_CHUNK_SIZE_BYTES = 2 * 1024 * 1024 # 2 MiB | ||
| 35 | + _MAX_BUFFER_SIZE_BYTES = 16 * 1024 * 1024 # 16 MiB | ||
| 36 | + | ||
| 37 | + | ||
| 34 | 38 | class AsyncAppendableObjectWriter: | |
| 35 | 39 | """Class for appending data to a GCS Appendable Object asynchronously.""" | |
| 36 | 40 | ||
@@ -118,7 +122,13 @@ async def state_lookup(self) -> int: | |||
| 118 | 122 | ||
| 119 | 123 | :rtype: int | |
| 120 | 124 | :returns: persisted size. | |
| 125 | + | ||
| 126 | + :raises ValueError: If the stream is not open (i.e., `open()` has not | ||
| 127 | + been called). | ||
| 121 | 128 | """ | |
| 129 | + if not self._is_stream_open: | ||
| 130 | + raise ValueError("Stream is not open. Call open() before state_lookup().") | ||
| 131 | + | ||
| 122 | 132 | await self.write_obj_stream.send( | |
| 123 | 133 | _storage_v2.BidiWriteObjectRequest( | |
| 124 | 134 | state_lookup=True, | |
@@ -129,7 +139,11 @@ async def state_lookup(self) -> int: | |||
| 129 | 139 | return self.persisted_size | |
| 130 | 140 | ||
| 131 | 141 | async def open(self) -> None: | |
| 132 | - """Opens the underlying bidi-gRPC stream.""" | ||
| 142 | + """Opens the underlying bidi-gRPC stream. | ||
| 143 | + | ||
| 144 | + :raises ValueError: If the stream is already open. | ||
| 145 | + | ||
| 146 | + """ | ||
| 133 | 147 | if self._is_stream_open: | |
| 134 | 148 | raise ValueError("Underlying bidi-gRPC stream is already open") | |
| 135 | 149 | ||
@@ -142,15 +156,65 @@ async def open(self) -> None: | |||
| 142 | 156 | # Update self.persisted_size | |
| 143 | 157 | _ = await self.state_lookup() | |
| 144 | 158 | ||
| 145 | - async def append(self, data: bytes): | ||
| 146 | - raise NotImplementedError("append is not implemented yet.") | ||
| 159 | + async def append(self, data: bytes) -> None: | ||
| 160 | + """Appends data to the Appendable object. | ||
| 161 | + | ||
| 162 | + This method sends the provided data to the GCS server in chunks. It | ||
| 163 | + maintains an internal threshold `_MAX_BUFFER_SIZE_BYTES` and will | ||
| 164 | + automatically flush the data to make it visible to readers when that | ||
| 165 | + threshold has reached. | ||
| 166 | + | ||
| 167 | + :type data: bytes | ||
| 168 | + :param data: The bytes to append to the object. | ||
| 169 | + | ||
| 170 | + :rtype: None | ||
| 171 | + | ||
| 172 | + :raises ValueError: If the stream is not open (i.e., `open()` has not | ||
| 173 | + been called). | ||
| 174 | + """ | ||
| 175 | + | ||
| 176 | + if not self._is_stream_open: | ||
| 177 | + raise ValueError("Stream is not open. Call open() before append().") | ||
| 178 | + total_bytes = len(data) | ||
| 179 | + if total_bytes == 0: | ||
| 180 | + # TODO: add warning. | ||
| 181 | + return | ||
| 182 | + if self.offset is None: | ||
| 183 | + assert self.persisted_size is not None | ||
| 184 | + self.offset = self.persisted_size | ||
| 185 | + | ||
| 186 | + start_idx = 0 | ||
| 187 | + bytes_to_flush = 0 | ||
| 188 | + while start_idx < total_bytes: | ||
| 189 | + end_idx = min(start_idx + _MAX_CHUNK_SIZE_BYTES, total_bytes) | ||
| 190 | + await self.write_obj_stream.send( | ||
| 191 | + _storage_v2.BidiWriteObjectRequest( | ||
| 192 | + write_offset=self.offset, | ||
| 193 | + checksummed_data=_storage_v2.ChecksummedData( | ||
| 194 | + content=data[start_idx:end_idx] | ||
| 195 | + ), | ||
| 196 | + ) | ||
| 197 | + ) | ||
| 198 | + chunk_size = end_idx - start_idx | ||
| 199 | + self.offset += chunk_size | ||
| 200 | + bytes_to_flush += chunk_size | ||
| 201 | + if bytes_to_flush >= _MAX_BUFFER_SIZE_BYTES: | ||
| 202 | + await self.flush() | ||
| 203 | + bytes_to_flush = 0 | ||
| 204 | + start_idx = end_idx | ||
| 147 | 205 | ||
| 148 | 206 | async def flush(self) -> int: | |
| 149 | 207 | """Flushes the data to the server. | |
| 150 | 208 | ||
| 151 | 209 | :rtype: int | |
| 152 | 210 | :returns: The persisted size after flush. | |
| 211 | + | ||
| 212 | + :raises ValueError: If the stream is not open (i.e., `open()` has not | ||
| 213 | + been called). | ||
| 153 | 214 | """ | |
| 215 | + if not self._is_stream_open: | ||
| 216 | + raise ValueError("Stream is not open. Call open() before flush().") | ||
| 217 | + | ||
| 154 | 218 | await self.write_obj_stream.send( | |
| 155 | 219 | _storage_v2.BidiWriteObjectRequest( | |
| 156 | 220 | flush=True, | |
@@ -162,14 +226,34 @@ async def flush(self) -> int: | |||
| 162 | 226 | self.offset = self.persisted_size | |
| 163 | 227 | return self.persisted_size | |
| 164 | 228 | ||
| 165 | - async def close(self, finalize_on_close=False) -> int: | ||
| 166 | - """Returns persisted_size""" | ||
| 229 | + async def close(self, finalize_on_close=False) -> Union[int, _storage_v2.Object]: | ||
| 230 | + """Closes the underlying bidi-gRPC stream. | ||
| 231 | + | ||
| 232 | + :type finalize_on_close: bool | ||
| 233 | + :param finalize_on_close: Finalizes the Appendable Object. No more data | ||
| 234 | + can be appended. | ||
| 235 | + | ||
| 236 | + rtype: Union[int, _storage_v2.Object] | ||
| 237 | + returns: Updated `self.persisted_size` by default after closing the | ||
| 238 | + bidi-gRPC stream. However, if `finalize_on_close=True` is passed, | ||
| 239 | + returns the finalized object resource. | ||
| 240 | + | ||
| 241 | + :raises ValueError: If the stream is not open (i.e., `open()` has not | ||
| 242 | + been called). | ||
| 243 | + | ||
| 244 | + """ | ||
| 245 | + if not self._is_stream_open: | ||
| 246 | + raise ValueError("Stream is not open. Call open() before close().") | ||
| 247 | + | ||
| 167 | 248 | if finalize_on_close: | |
| 168 | 249 | await self.finalize() | |
| 250 | + else: | ||
| 251 | + await self.flush() | ||
| 252 | + await self.write_obj_stream.close() | ||
| 169 | 253 | ||
| 170 | - await self.write_obj_stream.close() | ||
| 171 | 254 | self._is_stream_open = False | |
| 172 | 255 | self.offset = None | |
| 256 | + return self.object_resource if finalize_on_close else self.persisted_size | ||
| 173 | 257 | ||
| 174 | 258 | async def finalize(self) -> _storage_v2.Object: | |
| 175 | 259 | """Finalizes the Appendable Object. | |
@@ -178,12 +262,20 @@ async def finalize(self) -> _storage_v2.Object: | |||
| 178 | 262 | ||
| 179 | 263 | rtype: google.cloud.storage_v2.types.Object | |
| 180 | 264 | returns: The finalized object resource. | |
| 265 | + | ||
| 266 | + :raises ValueError: If the stream is not open (i.e., `open()` has not | ||
| 267 | + been called). | ||
| 181 | 268 | """ | |
| 269 | + if not self._is_stream_open: | ||
| 270 | + raise ValueError("Stream is not open. Call open() before finalize().") | ||
| 271 | + | ||
| 182 | 272 | await self.write_obj_stream.send( | |
| 183 | 273 | _storage_v2.BidiWriteObjectRequest(finish_write=True) | |
| 184 | 274 | ) | |
| 185 | 275 | response = await self.write_obj_stream.recv() | |
| 186 | 276 | self.object_resource = response.resource | |
| 277 | + self.persisted_size = self.object_resource.size | ||
| 278 | + return self.object_resource | ||
| 187 | 279 | ||
| 188 | 280 | # helper methods. | |
| 189 | 281 | async def append_from_string(self, data: str): | |
| Back | FazBrowse Home | New Git URL |
0 commit comments