Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 20 additions & 20 deletions ibis/backends/bigquery/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -544,7 +544,7 @@
catalog: str | None = None,
force: bool = False,
collate: str | None = None,
job_id: str | None = None,
job_id_prefix: str | None = None,
**options: Any,
) -> None:
properties = [
Expand All @@ -564,7 +564,7 @@
properties=sge.Properties(expressions=properties),
)

self.raw_sql(stmt.sql(self.name), job_id=job_id)
self.raw_sql(stmt.sql(self.name), job_id_prefix=job_id_prefix)

Check warning on line 567 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L567

Added line #L567 was not covered by tests

def drop_database(
self,
Expand All @@ -574,7 +574,7 @@
catalog: str | None = None,
force: bool = False,
cascade: bool = False,
job_id: str | None = None,
job_id_prefix: str | None = None,
) -> None:
"""Drop a BigQuery dataset."""
stmt = sge.Drop(
Expand All @@ -584,7 +584,7 @@
cascade=cascade,
)

self.raw_sql(stmt.sql(self.name), job_id=job_id)
self.raw_sql(stmt.sql(self.name), job_id_prefix=job_id_prefix)

Check warning on line 587 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L587

Added line #L587 was not covered by tests

def table(
self,
Expand Down Expand Up @@ -670,7 +670,7 @@
return BigQuerySchema.to_ibis(job.schema)

def raw_sql(
self, query: str, params=None, job_id: str | None = None
self, query: str, params=None, job_id_prefix: str | None = None
) -> RowIterator:
query_parameters = [
bigquery_param(param.type(), value, param.get_name())
Expand All @@ -681,12 +681,12 @@

job_config = bq.job.QueryJobConfig(query_parameters=query_parameters or [])

if job_id is not None:
if job_id_prefix is not None:
return self.client.query(
query,
job_config=job_config,
project=self.billing_project,
job_id=job_id,
job_id_prefix=job_id_prefix,
).result()
else:
return self.client.query_and_wait(
Expand Down Expand Up @@ -781,14 +781,14 @@
*,
params: Mapping[ir.Scalar, Any] | None = None,
limit: int | str | None = None,
job_id: str | None = None,
job_id_prefix: str | None = None,
**kwargs: Any,
) -> RowIterator:
self._run_pre_execute_hooks(table_expr)
sql = self.compile(table_expr, limit=limit, params=params, **kwargs)
self._log(sql)

return self.raw_sql(sql, params=params, job_id=job_id)
return self.raw_sql(sql, params=params, job_id_prefix=job_id_prefix)

Check warning on line 791 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L791

Added line #L791 was not covered by tests

def to_pyarrow(
self,
Expand Down Expand Up @@ -982,7 +982,7 @@
partition_by: str | None = None,
cluster_by: Iterable[str] | None = None,
options: Mapping[str, Any] | None = None,
job_id: str | None = None,
job_id_prefix: str | None = None,
) -> ir.Table:
"""Create a table in BigQuery.

Expand Down Expand Up @@ -1015,9 +1015,9 @@
options
BigQuery-specific table options; see the BigQuery documentation for
details: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#table_option_list
job_id
Optional custom job ID; when specified, bigquery will use this job ID instead
of a randomly generated one
job_id_prefix
Optional custom job ID prefix; when specified, bigquery will use this as a
prefix for the job ID it generates

Returns
-------
Expand Down Expand Up @@ -1109,7 +1109,7 @@

sql = stmt.sql(self.name)

self.raw_sql(sql, job_id=job_id)
self.raw_sql(sql, job_id_prefix=job_id_prefix)

Check warning on line 1112 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L1112

Added line #L1112 was not covered by tests
return self.table(table.name, database=(table.catalog, table.db))

def drop_table(
Expand All @@ -1119,7 +1119,7 @@
*,
database: tuple[str | str] | str | None = None,
force: bool = False,
job_id: str | None = None,
job_id_prefix: str | None = None,
) -> None:
table_loc = self._to_sqlglot_table(database)
catalog, db = self._to_catalog_db_tuple(table_loc)
Expand All @@ -1132,7 +1132,7 @@
),
exists=force,
)
self.raw_sql(stmt.sql(self.name), job_id=job_id)
self.raw_sql(stmt.sql(self.name), job_id_prefix=job_id_prefix)

Check warning on line 1135 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L1135

Added line #L1135 was not covered by tests

def create_view(
self,
Expand All @@ -1142,7 +1142,7 @@
*,
database: str | None = None,
overwrite: bool = False,
job_id: str | None = None,
job_id_prefix: str | None = None,
) -> ir.Table:
table_loc = self._to_sqlglot_table(database)
catalog, db = self._to_catalog_db_tuple(table_loc)
Expand All @@ -1158,7 +1158,7 @@
replace=overwrite,
)
self._run_pre_execute_hooks(obj)
self.raw_sql(stmt.sql(self.name), job_id=job_id)
self.raw_sql(stmt.sql(self.name), job_id_prefix=job_id_prefix)

Check warning on line 1161 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L1161

Added line #L1161 was not covered by tests
return self.table(name, database=(catalog, database))

def drop_view(
Expand All @@ -1168,7 +1168,7 @@
*,
database: str | None = None,
force: bool = False,
job_id: str | None = None,
job_id_prefix: str | None = None,
) -> None:
table_loc = self._to_sqlglot_table(database)
catalog, db = self._to_catalog_db_tuple(table_loc)
Expand All @@ -1182,7 +1182,7 @@
),
exists=force,
)
self.raw_sql(stmt.sql(self.name), job_id=job_id)
self.raw_sql(stmt.sql(self.name), job_id_prefix=job_id_prefix)

Check warning on line 1185 in ibis/backends/bigquery/__init__.py

View check run for this annotation

Codecov / codecov/patch

ibis/backends/bigquery/__init__.py#L1185

Added line #L1185 was not covered by tests

def _drop_cached_table(self, name):
self.drop_table(
Expand Down
9 changes: 4 additions & 5 deletions ibis/backends/bigquery/tests/system/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import decimal
import os
from urllib.parse import urlparse
from uuid import uuid4

import pandas as pd
import pandas.testing as tm
Expand Down Expand Up @@ -549,8 +548,8 @@ def test_window_with_count_distinct(tmp_table, expr, query):
tm.assert_frame_equal(result, expected)


def test_create_job_with_custom_job_id(con):
job_id = f"test_job_id_{uuid4()}"
def test_create_job_with_custom_job_prefix(con):
job_id_prefix = "test_job_prefix_"
query = "SELECT 1"
result = con.raw_sql(query, job_id=job_id)
assert result.job_id == job_id
result = con.raw_sql(query, job_id_prefix=job_id_prefix)
assert result.job_id.startswith(job_id_prefix)
Loading