# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module for bigquery continuous queries"""
from __future__ import annotations
from abc import abstractmethod
from datetime import date, datetime
import functools
import inspect
import json
from typing import Optional, Union
import warnings
from google.cloud import bigquery
import pandas as pd
from bigframes import dataframe
from bigframes.core import log_adapter, nodes
import bigframes.exceptions as bfe
import bigframes.session
def _return_type_wrapper(method, cls):
@functools.wraps(method)
def wrapper(*args, **kwargs):
return_value = method(*args, **kwargs)
if isinstance(return_value, dataframe.DataFrame):
return cls._from_table_df(return_value)
return return_value
return wrapper
def _curate_df_doc(doc: Optional[str]):
if not doc:
return doc
# Remove examples, some are not applicable to StreamingDataFrame
doc = doc[: doc.find("**Examples:**")] + doc[doc.find("Args:") :]
doc = doc.replace("dataframe.DataFrame", "streaming.StreamingDataFrame")
doc = doc.replace(" DataFrame", " StreamingDataFrame")
return doc
class StreamingBase:
_session: bigframes.session.Session
@abstractmethod
def _appends_sql(
self, start_timestamp: Optional[Union[int, float, str, datetime, date]]
) -> str:
pass
def to_bigtable(
self,
*,
instance: str,
table: str,
service_account_email: Optional[str] = None,
app_profile: Optional[str] = None,
truncate: bool = False,
overwrite: bool = False,
auto_create_column_families: bool = False,
bigtable_options: Optional[dict] = None,
job_id: Optional[str] = None,
job_id_prefix: Optional[str] = None,
start_timestamp: Optional[Union[int, float, str, datetime, date]] = None,
end_timestamp: Optional[Union[int, float, str, datetime, date]] = None,
) -> bigquery.QueryJob:
"""
Export the StreamingDataFrame as a continue job and returns a
QueryJob object for some management functionality.
This method requires an existing bigtable preconfigured to
accept the continuous query export statement. For instructions
on export to bigtable, see
https://cloud.google.com/bigquery/docs/export-to-bigtable.
Args:
instance (str):
The name of the bigtable instance to export to.
table (str):
The name of the bigtable table to export to.
service_account_email (str):
Full name of the service account to run the continuous query.
Example: accountname@projectname.gserviceaccounts.com
If not provided, the user account will be used, but this
limits the lifetime of the continuous query.
app_profile (str, default None):
The bigtable app profile to export to. If None, no app
profile will be used.
truncate (bool, default False):
The export truncate option, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
overwrite (bool, default False):
The export overwrite option, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
auto_create_column_families (bool, default False):
The auto_create_column_families option, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
bigtable_options (dict, default None):
The bigtable options dict, which will be converted to JSON
using json.dumps, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
If None, no bigtable_options parameter will be passed.
job_id (str, default None):
If specified, replace the default job id for the query,
see job_id parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
job_id_prefix (str, default None):
If specified, a job id prefix for the query, see
job_id_prefix parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
start_timestamp (int, float, str, datetime, date, default None):
The starting timestamp for the query. Possible values are to 7 days in the past. If don't specify a timestamp (None), the query will default to the earliest possible time, 7 days ago. If provide a time-zone-naive timestamp, it will be treated as UTC.
Returns:
google.cloud.bigquery.QueryJob:
See https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob
The ongoing query job can be managed using this object.
For example, the job can be cancelled or its error status
can be examined.
"""
if not isinstance(
start_timestamp, (int, float, str, datetime, date, type(None))
):
raise ValueError(
f"Unsupported start_timestamp type {type(start_timestamp)}"
)
return _to_bigtable(
self._appends_sql(start_timestamp),
instance=instance,
table=table,
service_account_email=service_account_email,
session=self._session,
app_profile=app_profile,
truncate=truncate,
overwrite=overwrite,
auto_create_column_families=auto_create_column_families,
bigtable_options=bigtable_options,
job_id=job_id,
job_id_prefix=job_id_prefix,
)
def to_pubsub(
self,
*,
topic: str,
service_account_email: str,
job_id: Optional[str] = None,
job_id_prefix: Optional[str] = None,
start_timestamp: Optional[Union[int, float, str, datetime, date]] = None,
) -> bigquery.QueryJob:
"""
Export the StreamingDataFrame as a continue job and returns a
QueryJob object for some management functionality.
This method requires an existing pubsub topic. For instructions
on creating a pubsub topic, see
https://cloud.google.com/pubsub/docs/samples/pubsub-quickstart-create-topic?hl=en
Note that a service account is a requirement for continuous queries
exporting to pubsub.
Args:
topic (str):
The name of the pubsub topic to export to.
For example: "taxi-rides"
service_account_email (str):
Full name of the service account to run the continuous query.
Example: accountname@projectname.gserviceaccounts.com
job_id (str, default None):
If specified, replace the default job id for the query,
see job_id parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
job_id_prefix (str, default None):
If specified, a job id prefix for the query, see
job_id_prefix parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
start_timestamp (int, float, str, datetime, date, default None):
The starting timestamp for the query. Possible values are to 7 days in the past. If don't specify a timestamp (None), the query will default to the earliest possible time, 7 days ago. If provide a time-zone-naive timestamp, it will be treated as UTC.
Returns:
google.cloud.bigquery.QueryJob:
See https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob
The ongoing query job can be managed using this object.
For example, the job can be cancelled or its error status
can be examined.
"""
if not isinstance(
start_timestamp, (int, float, str, datetime, date, type(None))
):
raise ValueError(
f"Unsupported start_timestamp type {type(start_timestamp)}"
)
return _to_pubsub(
self._appends_sql(start_timestamp),
topic=topic,
service_account_email=service_account_email,
session=self._session,
job_id=job_id,
job_id_prefix=job_id_prefix,
)
[docs]
@log_adapter.class_logger
class StreamingDataFrame(StreamingBase):
__doc__ = (
_curate_df_doc(dataframe.DataFrame.__doc__)
+ """
.. note::
The bigframes.streaming module is a preview feature, and subject to change.
Currently only supports basic projection, filtering and preview operations.
"""
)
# Private constructor
_create_key = object()
[docs]
def __init__(self, df: dataframe.DataFrame, *, create_key=0):
if create_key is not StreamingDataFrame._create_key:
raise ValueError(
"StreamingDataFrame class shouldn't be created through constructor. Call bigframes.pandas.read_gbq_table_streaming method to create."
)
self._df = df
self._df._disable_cache_override = True
@classmethod
def _from_table_df(cls, df: dataframe.DataFrame) -> StreamingDataFrame:
return cls(df, create_key=cls._create_key)
@property
def _original_table(self):
def traverse(node: nodes.BigFrameNode):
if isinstance(node, nodes.ReadTableNode):
return f"{node.source.table.project_id}.{node.source.table.dataset_id}.{node.source.table.table_id}"
for child in node.child_nodes:
original_table = traverse(child)
if original_table:
return original_table
return None
return traverse(self._df._block._expr.node)
def __getitem__(self, *args, **kwargs):
return _return_type_wrapper(self._df.__getitem__, StreamingDataFrame)(
*args, **kwargs
)
__getitem__.__doc__ = _curate_df_doc(
inspect.getdoc(dataframe.DataFrame.__getitem__)
)
def __setitem__(self, *args, **kwargs):
return _return_type_wrapper(self._df.__setitem__, StreamingDataFrame)(
*args, **kwargs
)
__setitem__.__doc__ = _curate_df_doc(
inspect.getdoc(dataframe.DataFrame.__setitem__)
)
[docs]
def rename(self, *args, **kwargs):
return _return_type_wrapper(self._df.rename, StreamingDataFrame)(
*args, **kwargs
)
rename.__doc__ = _curate_df_doc(inspect.getdoc(dataframe.DataFrame.rename))
def __repr__(self, *args, **kwargs):
return _return_type_wrapper(self._df.__repr__, StreamingDataFrame)(
*args, **kwargs
)
__repr__.__doc__ = _curate_df_doc(inspect.getdoc(dataframe.DataFrame.__repr__))
def _repr_html_(self, *args, **kwargs):
return _return_type_wrapper(self._df._repr_html_, StreamingDataFrame)(
*args, **kwargs
)
_repr_html_.__doc__ = _curate_df_doc(
inspect.getdoc(dataframe.DataFrame._repr_html_)
)
@property
def sql(self):
sql_str, _, _ = self._df._to_sql_query(include_index=False, enable_cache=False)
return sql_str
sql.__doc__ = _curate_df_doc(inspect.getdoc(dataframe.DataFrame.sql))
# Patch for the required APPENDS clause
def _appends_sql(
self, start_timestamp: Optional[Union[int, float, str, datetime, date]]
) -> str:
sql_str = self.sql
original_table = self._original_table
assert original_table is not None
# TODO(b/405691193): set start time back to NULL. Now set it slightly after 7 days max interval to avoid the bug.
start_ts_str = (
str(f"TIMESTAMP('{pd.to_datetime(start_timestamp)}')")
if start_timestamp
else "CURRENT_TIMESTAMP() - (INTERVAL 7 DAY - INTERVAL 5 MINUTE)"
)
appends_clause = f"APPENDS(TABLE `{original_table}`, {start_ts_str})"
sql_str = sql_str.replace(f"`{original_table}`", appends_clause)
return sql_str
@property
def _session(self):
return self._df._session
_session.__doc__ = _curate_df_doc(inspect.getdoc(dataframe.DataFrame._session))
def _to_bigtable(
query: str,
*,
instance: str,
table: str,
service_account_email: Optional[str] = None,
session: Optional[bigframes.session.Session] = None,
app_profile: Optional[str] = None,
truncate: bool = False,
overwrite: bool = False,
auto_create_column_families: bool = False,
bigtable_options: Optional[dict] = None,
job_id: Optional[str] = None,
job_id_prefix: Optional[str] = None,
) -> bigquery.QueryJob:
"""Launches a BigQuery continuous query and returns a
QueryJob object for some management functionality.
This method requires an existing bigtable preconfigured to
accept the continuous query export statement. For instructions
on export to bigtable, see
https://cloud.google.com/bigquery/docs/export-to-bigtable.
Args:
query (str):
The sql statement to execute as a continuous function.
For example: "SELECT * FROM dataset.table"
This will be wrapped in an EXPORT DATA statement to
launch a continuous query writing to bigtable.
instance (str):
The name of the bigtable instance to export to.
table (str):
The name of the bigtable table to export to.
service_account_email (str):
Full name of the service account to run the continuous query.
Example: accountname@projectname.gserviceaccounts.com
If not provided, the user account will be used, but this
limits the lifetime of the continuous query.
session (bigframes.session.Session, default None):
The session object to use for the query. This determines
the project id and location of the query. If None, will
default to the bigframes global session.
app_profile (str, default None):
The bigtable app profile to export to. If None, no app
profile will be used.
truncate (bool, default False):
The export truncate option, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
overwrite (bool, default False):
The export overwrite option, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
auto_create_column_families (bool, default False):
The auto_create_column_families option, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
bigtable_options (dict, default None):
The bigtable options dict, which will be converted to JSON
using json.dumps, see
https://cloud.google.com/bigquery/docs/reference/standard-sql/other-statements#bigtable_export_option
If None, no bigtable_options parameter will be passed.
job_id (str, default None):
If specified, replace the default job id for the query,
see job_id parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
job_id_prefix (str, default None):
If specified, a job id prefix for the query, see
job_id_prefix parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
Returns:
google.cloud.bigquery.QueryJob:
See https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob
The ongoing query job can be managed using this object.
For example, the job can be cancelled or its error status
can be examined.
"""
msg = bfe.format_message(
"The bigframes.streaming module is a preview feature, and subject to change."
)
warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)
# get default client if not passed
if session is None:
session = bigframes.get_global_session()
bq_client = session.bqclient
# build export string from parameters
project = bq_client.project
app_profile_url_string = ""
if app_profile is not None:
app_profile_url_string = f"appProfiles/{app_profile}/"
bigtable_options_parameter_string = ""
if bigtable_options is not None:
bigtable_options_parameter_string = (
'bigtable_options = """' + json.dumps(bigtable_options) + '""",\n'
)
sql = (
"EXPORT DATA\n"
"OPTIONS (\n"
"format = 'CLOUD_BIGTABLE',\n"
f"{bigtable_options_parameter_string}"
f"truncate = {str(truncate)},\n"
f"overwrite = {str(overwrite)},\n"
f"auto_create_column_families = {str(auto_create_column_families)},\n"
f'uri = "https://bigtable.googleapis.com/projects/{project}/instances/{instance}/{app_profile_url_string}tables/{table}"\n'
")\n"
"AS (\n"
f"{query});"
)
# override continuous http parameter
job_config = bigquery.job.QueryJobConfig()
job_config_dict: dict = {"query": {"continuous": True}}
if service_account_email is not None:
job_config_dict["query"]["connectionProperties"] = {
"key": "service_account",
"value": service_account_email,
}
job_config_filled = job_config.from_api_repr(job_config_dict)
job_config_filled.labels = {"bigframes-api": "streaming_to_bigtable"}
# begin the query job
query_job = bq_client.query(
sql,
job_config=job_config_filled, # type:ignore
# typing error above is in bq client library
# (should accept abstract job_config, only takes concrete)
job_id=job_id,
job_id_prefix=job_id_prefix,
)
# return the query job to the user for lifetime management
return query_job
def _to_pubsub(
query: str,
*,
topic: str,
service_account_email: str,
session: Optional[bigframes.session.Session] = None,
job_id: Optional[str] = None,
job_id_prefix: Optional[str] = None,
) -> bigquery.QueryJob:
"""Launches a BigQuery continuous query and returns a
QueryJob object for some management functionality.
This method requires an existing pubsub topic. For instructions
on creating a pubsub topic, see
https://cloud.google.com/pubsub/docs/samples/pubsub-quickstart-create-topic?hl=en
Note that a service account is a requirement for continuous queries
exporting to pubsub.
Args:
query (str):
The sql statement to execute as a continuous function.
For example: "SELECT * FROM dataset.table"
This will be wrapped in an EXPORT DATA statement to
launch a continuous query writing to pubsub.
topic (str):
The name of the pubsub topic to export to.
For example: "taxi-rides"
service_account_email (str):
Full name of the service account to run the continuous query.
Example: accountname@projectname.gserviceaccounts.com
session (bigframes.session.Session, default None):
The session object to use for the query. This determines
the project id and location of the query. If None, will
default to the bigframes global session.
job_id (str, default None):
If specified, replace the default job id for the query,
see job_id parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
job_id_prefix (str, default None):
If specified, a job id prefix for the query, see
job_id_prefix parameter of
https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.client.Client#google_cloud_bigquery_client_Client_query
Returns:
google.cloud.bigquery.QueryJob:
See https://cloud.google.com/python/docs/reference/bigquery/latest/google.cloud.bigquery.job.QueryJob
The ongoing query job can be managed using this object.
For example, the job can be cancelled or its error status
can be examined.
"""
msg = bfe.format_message(
"The bigframes.streaming module is a preview feature, and subject to change."
)
warnings.warn(msg, stacklevel=1, category=bfe.PreviewWarning)
# get default client if not passed
if session is None:
session = bigframes.get_global_session()
bq_client = session.bqclient
# build export string from parameters
sql = (
"EXPORT DATA\n"
"OPTIONS (\n"
"format = 'CLOUD_PUBSUB',\n"
f'uri = "https://pubsub.googleapis.com/projects/{bq_client.project}/topics/{topic}"\n'
")\n"
"AS (\n"
f"{query});"
)
# override continuous http parameter
job_config = bigquery.job.QueryJobConfig()
job_config_filled = job_config.from_api_repr(
{
"query": {
"continuous": True,
"connectionProperties": {
"key": "service_account",
"value": service_account_email,
},
}
}
)
job_config_filled.labels = {"bigframes-api": "streaming_to_pubsub"}
# begin the query job
query_job = bq_client.query(
sql,
job_config=job_config_filled, # type:ignore
# typing error above is in bq client library
# (should accept abstract job_config, only takes concrete)
job_id=job_id,
job_id_prefix=job_id_prefix,
)
# return the query job to the user for lifetime management
return query_job