Skip to content
Closed
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
7 changes: 3 additions & 4 deletions bigquery_magics/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -467,10 +467,9 @@ def _query_with_bigframes(query: str, params: List[Any], args: Any):
raise ValueError("Dry run is not supported by bigframes engine.")

if bpd is None:
raise ValueError("Bigframes package is not installed.")

bpd.options.bigquery.project = context.project
bpd.options.bigquery.credentials = context.credentials
raise ValueError(
"Please install the 'bigframes' package (pip install bigframes) to use the bigframes engine."
)

max_results = int(args.max_results) if args.max_results else None

Expand Down
50 changes: 47 additions & 3 deletions bigquery_magics/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,9 @@
# See the License for the specific language governing permissions and
# limitations under the License.

from dataclasses import dataclass
from typing import Optional
import dataclasses
import enum
from typing import Optional, Union

import google.api_core.client_options as client_options
import google.cloud.bigquery as bigquery
Expand All @@ -26,7 +27,50 @@ def _get_default_credentials_with_project():
return pydata_google_auth.default(scopes=_SCOPES, use_local_webserver=False)


@dataclass
class MissingReason(enum.Enum):
"""Provides a way to disambiguate why an option is missing.

This is used instead of a None value to allow for custom validation and
docs generation.
"""

# These missing reasons are because the engine doesn't support the feature,
# for example, use_rest_api on the bigframes engine. In this case, raise
# if the user has set it, as the magics woult otherwise act in a way
# contrary to that in which the user explicitly requested.
NOT_SUPPORTED_BY_ENGINE_INFEASIBLE = enum.auto()

# This is like the above, but is theoretically possible. Include a call to
# action to reach out to the bigframes team if this is a feature they
# would like to use.
NOT_SUPPORTED_BY_ENGINE_BUT_POSSIBLE = enum.auto()

# This missing reason is for options that are magics-only and apply to all
# engines. For example, the destination variable is a magics-only setting
# and doesn't affect how queries are handled by the engine.
ENGINE_UNIVERSAL = enum.auto()


@dataclasses.dataclass(frozen=True)
class MagicsSetting:
"""Encapsulates information about settings and how to set them.

This is used to generate documentation, merge settings across the various
ways they are duplicated, and to validate settings provided by a user.
"""

description: str
cell_arg: Union[MissingReason, str]
magics_context: Union[MissingReason, str]
bigframes_option: Union[MissingReason, str]


magics_settings = [
# TODO: copy table from sheets to here
Copy link
Collaborator Author

@tswast tswast Jan 14, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

example:

Suggested change
# TODO: copy table from sheets to here
# TODO: copy table from sheets to here
MagicSetting("Limits the number of rows in the returned DataFrame", cell_arg="--max_results", magics_context=MissingReason.NOT_SUPPORTED_BY_ENGINE_BUT_POSSIBLE, bigframes_option="read_gbq_query.max_results"),

]


@dataclasses.dataclass
class Context(object):
"""Storage for objects to be used throughout an IPython notebook session.

Expand Down
4 changes: 3 additions & 1 deletion tests/unit/test_bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -2013,7 +2013,9 @@ def test_bigquery_magic_bigframes__bigframes_is_not_installed__should_raise_erro
ip.extension_manager.load_extension("bigquery_magics")
sql = "SELECT 0 AS something"

with pytest.raises(ValueError, match="Bigframes package is not installed."):
with pytest.raises(
ValueError, match=re.escape("Please install the 'bigframes' package")
):
ip.run_cell_magic("bigquery", "", sql)


Expand Down
Loading