Skip to content

Commit

Permalink
[qob] allow writing to requester pays buckets
Browse files Browse the repository at this point in the history
CHANGELOG: In Query-on-Batch, allow writing to requester pays buckets, which was broken before this release.
  • Loading branch information
Dan King committed Apr 7, 2023
1 parent a980f1c commit 8d31848
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 2 deletions.
12 changes: 12 additions & 0 deletions hail/python/test/hail/fs/test_worker_driver_fs.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,18 @@ def test_requester_pays_write_no_settings():
assert False


@skip_in_azure
def test_requester_pays_write_with_project():
hl.stop()
hl.init(gcs_requester_pays_configuration='hail-vdc')
random_filename = 'gs://hail-services-requester-pays/test_requester_pays_on_worker_driver_' + secret_alnum_string(10)
try:
hl.utils.range_table(4, n_partitions=4).write(random_filename, overwrite=True)
finally:
hl.current_backend().fs.rmtree(random_filename)
assert False


@skip_in_azure
def test_requester_pays_with_project():
hl.stop()
Expand Down
19 changes: 17 additions & 2 deletions hail/src/main/scala/is/hail/io/fs/GoogleStorageFS.scala
Original file line number Diff line number Diff line change
Expand Up @@ -261,13 +261,28 @@ class GoogleStorageFS(
.build()

val os: PositionedOutputStream = new FSPositionedOutputStream(8 * 1024 * 1024) {
private[this] val write: WriteChannel = storage.writer(blobInfo)
private[this] var writer: WriteChannel = null

private[this] def writeHandlingRequesterPays(): Int = {
if (writer != null) {
writer.write(bb)
} else {
handleRequesterPays(
{ (options: Seq[BlobSourceOption]) =>
writer = storage.writer(blobInfo, options:_*)
writer.write(bb)
},
BlobSourceOption.userProject _,
bucket
)
}
}

override def flush(): Unit = {
bb.flip()

while (bb.remaining() > 0)
write.write(bb)
writeHandlingRequesterPays()

bb.clear()
}
Expand Down

0 comments on commit 8d31848

Please sign in to comment.