wehub-resource-sync
134 lines
5.1 KiB
Python
134 lines
5.1 KiB
Python
import logging
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import uuid
|
|
from typing import TYPE_CHECKING, Iterable, Optional
|
|
|
|
import pyarrow.parquet as pq
|
|
|
|
if TYPE_CHECKING:
|
|
import pyarrow as pa
|
|
|
|
import ray
|
|
from ray.data._internal.datasource import bigquery_datasource
|
|
from ray.data._internal.execution.interfaces import TaskContext
|
|
from ray.data._internal.remote_fn import cached_remote_fn
|
|
from ray.data._internal.util import _check_import
|
|
from ray.data.block import Block, BlockAccessor
|
|
from ray.data.datasource.datasink import Datasink
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_MAX_RETRY_CNT = 10
|
|
RATE_LIMIT_EXCEEDED_SLEEP_TIME = 11
|
|
|
|
|
|
class BigQueryDatasink(Datasink[None]):
|
|
def __init__(
|
|
self,
|
|
project_id: str,
|
|
dataset: str,
|
|
max_retry_cnt: int = DEFAULT_MAX_RETRY_CNT,
|
|
overwrite_table: Optional[bool] = True,
|
|
) -> None:
|
|
_check_import(self, module="google.cloud", package="bigquery")
|
|
_check_import(self, module="google.cloud", package="bigquery_storage")
|
|
_check_import(self, module="google.api_core", package="exceptions")
|
|
|
|
self.project_id = project_id
|
|
self.dataset = dataset
|
|
self.max_retry_cnt = max_retry_cnt
|
|
self.overwrite_table = overwrite_table
|
|
|
|
def on_write_start(self, schema: Optional["pa.Schema"] = None) -> None:
|
|
from google.api_core import exceptions
|
|
|
|
if self.project_id is None or self.dataset is None:
|
|
raise ValueError("project_id and dataset are required args")
|
|
|
|
# Set up datasets to write
|
|
client = bigquery_datasource._create_client(project_id=self.project_id)
|
|
dataset_id = self.dataset.split(".", 1)[0]
|
|
try:
|
|
client.get_dataset(dataset_id)
|
|
except exceptions.NotFound:
|
|
client.create_dataset(f"{self.project_id}.{dataset_id}", timeout=30)
|
|
logger.info("Created dataset " + dataset_id)
|
|
|
|
# Delete table if overwrite_table is True
|
|
if self.overwrite_table:
|
|
logger.info(
|
|
f"Attempting to delete table {self.dataset}"
|
|
+ " if it already exists since kwarg overwrite_table = True."
|
|
)
|
|
client.delete_table(f"{self.project_id}.{self.dataset}", not_found_ok=True)
|
|
else:
|
|
logger.info(
|
|
f"The write will append to table {self.dataset}"
|
|
+ " if it already exists since kwarg overwrite_table = False."
|
|
)
|
|
|
|
def write(
|
|
self,
|
|
blocks: Iterable[Block],
|
|
ctx: TaskContext,
|
|
) -> None:
|
|
def _write_single_block(block: Block, project_id: str, dataset: str) -> None:
|
|
from google.api_core import exceptions
|
|
from google.cloud import bigquery
|
|
|
|
block = BlockAccessor.for_block(block).to_arrow()
|
|
|
|
client = bigquery_datasource._create_client(project_id=project_id)
|
|
job_config = bigquery.LoadJobConfig(autodetect=True)
|
|
job_config.source_format = bigquery.SourceFormat.PARQUET
|
|
job_config.write_disposition = bigquery.WriteDisposition.WRITE_APPEND
|
|
|
|
with tempfile.TemporaryDirectory() as temp_dir:
|
|
fp = os.path.join(temp_dir, f"block_{uuid.uuid4()}.parquet")
|
|
pq.write_table(block, fp, compression="SNAPPY")
|
|
|
|
retry_cnt = 0
|
|
while retry_cnt <= self.max_retry_cnt:
|
|
with open(fp, "rb") as source_file:
|
|
job = client.load_table_from_file(
|
|
source_file, dataset, job_config=job_config
|
|
)
|
|
try:
|
|
logger.info(job.result())
|
|
break
|
|
except (exceptions.Forbidden, exceptions.TooManyRequests) as e:
|
|
retry_cnt += 1
|
|
if retry_cnt > self.max_retry_cnt:
|
|
break
|
|
logger.info(
|
|
"A block write encountered a rate limit exceeded error"
|
|
+ f" {retry_cnt} time(s). Sleeping to try again."
|
|
)
|
|
logging.debug(e)
|
|
time.sleep(RATE_LIMIT_EXCEEDED_SLEEP_TIME)
|
|
|
|
# Raise exception if retry_cnt exceeds max_retry_cnt
|
|
if retry_cnt > self.max_retry_cnt:
|
|
logger.info(
|
|
f"Maximum ({self.max_retry_cnt}) retry count exceeded. Ray"
|
|
" will attempt to retry the block write via fault tolerance."
|
|
)
|
|
raise RuntimeError(
|
|
f"Write failed due to {retry_cnt}"
|
|
" repeated API rate limit exceeded responses. Consider"
|
|
" specifying the max_retry_cnt kwarg with a higher value."
|
|
)
|
|
|
|
_write_single_block = cached_remote_fn(_write_single_block)
|
|
|
|
# Launch a remote task for each block within this write task
|
|
ray.get(
|
|
[
|
|
_write_single_block.remote(block, self.project_id, self.dataset)
|
|
for block in blocks
|
|
if BlockAccessor.for_block(block).num_rows() > 0
|
|
]
|
|
)
|