dataframe.py

# 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,
        )


@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()

    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__)
    )

    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_fallback(self, *args, **kwargs):
        return _return_type_wrapper(self._df._repr_html_fallback, StreamingDataFrame)(
            *args, **kwargs
        )

    _repr_html_fallback.__doc__ = _curate_df_doc(
        inspect.getdoc(dataframe.DataFrame._repr_html_fallback)
    )

    @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