alamashir commented on code in PR #65499: URL: https://github.com/apache/airflow/pull/65499#discussion_r3114553973
########## providers/google/src/airflow/providers/google/cloud/hooks/bigquery.py: ########## @@ -96,6 +97,21 @@ BigQueryJob = CopyJob | QueryJob | LoadJob | ExtractJob Review Comment: Good eye — the two uses of `fields` serve different layers and the mismatch is intentional (driven by the `google-cloud-bigquery` client's API contract): 1. **User-facing `fields` parameter** + merge loop — REST API camelCase names (e.g. `"definitionBody"`, `"description"`). Users interact with the hook using the BigQuery REST API's JSON shape because `routine` itself is also a REST-style dict. Keeping these consistent means the user never has to translate between two naming schemes. 2. **`_ROUTINE_WRITABLE_PROPERTIES` → `client.update_routine(merged, fields)`** — Python attribute names (`"body"`, `"type_"`). This one is forced by `google-cloud-bigquery`'s internal `_build_resource_from_properties` (see `google/cloud/bigquery/_helpers.py`), whose docstring states: *"filter_fields, which follow the name of the Python object"*. The client translates these back to REST names via `Routine._PROPERTY_TO_API_FIELD` before sending. Passing `"definitionBody"` here would silently drop the field. Since `routines.update` is currently a full-body PUT (not PATCH — see [the TODO in the client](https://github.com/googleapis/python-bigquery/blob/main/google/cloud/bigquery/client.py#L1431)), we pass the full writable-property list here to ensure the entire routine is sent. I could rename the internal constant to `_ROUTINE_WRITABLE_PROPERTY_NAMES` to make the distinction more visible in the code — would that help? Or I could add a short comment at the call site explaining the two naming schemes. Let me know which you prefer. ########## providers/google/src/airflow/providers/google/cloud/operators/bigquery.py: ########## @@ -2563,3 +2564,494 @@ def on_kill(self) -> None: ) else: self.log.info("Skipping to cancel job: %s:%s.%s", self.project_id, self.location, self.job_id) + + +class BigQueryCreateRoutineOperator(GoogleCloudBaseOperator): + """ + Create a BigQuery routine (UDF, stored procedure, table-valued function, or aggregate). + + The routine is defined by a ``Routine`` resource as documented at + https://cloud.google.com/bigquery/docs/reference/rest/v2/routines#Routine. The full resource + may be passed via ``routine_resource``, or individual fields can be passed as keyword + arguments, which are merged into ``routine_resource`` at execute time. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:BigQueryCreateRoutineOperator` + + :param dataset_id: The dataset that will own the routine. + :param routine_id: The identifier of the routine to create. + :param routine_resource: The routine resource as a dict or + :class:`~google.cloud.bigquery.routine.Routine`. If omitted, the resource is assembled + from the individual keyword arguments below. + :param project_id: Optional. The project that owns the dataset. Falls back to the connection's + default. + :param routine_type: Optional. One of ``"SCALAR_FUNCTION"``, ``"PROCEDURE"``, + ``"TABLE_VALUED_FUNCTION"``, ``"AGGREGATE_FUNCTION"``. + :param language: Optional. ``"SQL"`` or ``"JAVASCRIPT"``. + :param definition_body: Optional. The body of the routine (SQL or JavaScript). + :param arguments: Optional. Sequence of argument dicts in the Google API representation + (see the ``Routine.Argument`` schema). + :param return_type: Optional. The return type of the routine as a ``StandardSqlDataType`` + dict. + :param return_table_type: Optional. The return table type for table-valued functions as a + ``StandardSqlTableType`` dict. + :param imported_libraries: Optional. URIs of Cloud Storage libraries to import (JavaScript + UDFs only). + :param determinism_level: Optional. Determinism level for the routine. + :param security_mode: Optional. Security mode (``"DEFINER"`` or ``"INVOKER"``). + :param data_governance_type: Optional. Data governance type for the routine. + :param description: Optional. Description of the routine. + :param remote_function_options: Optional. Options for remote functions. + :param spark_options: Optional. Options for Spark stored procedures. + :param if_exists: What to do if a routine with the same identifier already exists. + ``"fail"`` (default) raises, ``"skip"`` returns the existing routine, ``"replace"`` + deletes the existing routine and creates the new one. + :param gcp_conn_id: The connection ID used to connect to Google Cloud. + :param location: The location of the BigQuery dataset. + :param impersonation_chain: Optional service account to impersonate using short-term + credentials, or chained list of accounts. + :param retry: A retry object used to retry requests. + :param timeout: The amount of time, in seconds, to wait for the request. + """ + + template_fields: Sequence[str] = ( + "project_id", + "dataset_id", + "routine_id", + "definition_body", + "arguments", + "routine_resource", + "gcp_conn_id", + "impersonation_chain", + ) + template_fields_renderers = {"routine_resource": "json", "arguments": "json"} + template_ext: Sequence[str] = (".sql",) + ui_color = BigQueryUIColors.TABLE.value + + def __init__( + self, + *, + dataset_id: str, + routine_id: str, + routine_resource: dict[str, Any] | Routine | None = None, + project_id: str = PROVIDE_PROJECT_ID, + routine_type: str | None = None, + language: str | None = None, + definition_body: str | None = None, + arguments: Sequence[dict[str, Any]] | None = None, + return_type: dict[str, Any] | None = None, + return_table_type: dict[str, Any] | None = None, + imported_libraries: Sequence[str] | None = None, + determinism_level: str | None = None, + security_mode: str | None = None, + data_governance_type: str | None = None, + description: str | None = None, + remote_function_options: dict[str, Any] | None = None, + spark_options: dict[str, Any] | None = None, + if_exists: str = "fail", + gcp_conn_id: str = "google_cloud_default", + location: str | None = None, + impersonation_chain: str | Sequence[str] | None = None, + retry: Retry = DEFAULT_RETRY, + timeout: float | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + if if_exists not in {"fail", "skip", "replace"}: + raise ValueError(f"`if_exists` must be one of 'fail', 'skip', 'replace'; got {if_exists!r}") + self.project_id = project_id + self.dataset_id = dataset_id + self.routine_id = routine_id + self.routine_resource = routine_resource + self.routine_type = routine_type + self.language = language + self.definition_body = definition_body + self.arguments = arguments + self.return_type = return_type + self.return_table_type = return_table_type + self.imported_libraries = imported_libraries + self.determinism_level = determinism_level + self.security_mode = security_mode + self.data_governance_type = data_governance_type + self.description = description + self.remote_function_options = remote_function_options + self.spark_options = spark_options + self.if_exists = if_exists + self.gcp_conn_id = gcp_conn_id + self.location = location + self.impersonation_chain = impersonation_chain + self.retry = retry + self.timeout = timeout + + def _build_resource(self) -> dict[str, Any]: + if isinstance(self.routine_resource, Routine): + resource = self.routine_resource.to_api_repr() + elif self.routine_resource is not None: + resource = dict(self.routine_resource) + else: + resource = {} + + field_map = { + "routineType": self.routine_type, + "language": self.language, + "definitionBody": self.definition_body, + "arguments": list(self.arguments) if self.arguments is not None else None, + "returnType": self.return_type, + "returnTableType": self.return_table_type, + "importedLibraries": ( + list(self.imported_libraries) if self.imported_libraries is not None else None + ), + "determinismLevel": self.determinism_level, + "securityMode": self.security_mode, + "dataGovernanceType": self.data_governance_type, + "description": self.description, + "remoteFunctionOptions": self.remote_function_options, + "sparkOptions": self.spark_options, + } + for key, value in field_map.items(): + if value is not None: + resource.setdefault(key, value) + return resource + + def execute(self, context: Context) -> dict[str, Any]: + hook = BigQueryHook( + gcp_conn_id=self.gcp_conn_id, + location=self.location, + impersonation_chain=self.impersonation_chain, + ) + resource = self._build_resource() + self.log.info( + "Creating routine %s.%s.%s (if_exists=%s)", + self.project_id or hook.project_id, + self.dataset_id, + self.routine_id, + self.if_exists, + ) + routine = hook.create_routine( + routine=resource, + dataset_id=self.dataset_id, + routine_id=self.routine_id, + project_id=self.project_id, + if_exists=self.if_exists, + retry=self.retry, + timeout=self.timeout, + ) + return routine.to_api_repr() + + +class BigQueryUpdateRoutineOperator(GoogleCloudBaseOperator): + """ + Patch selected fields of an existing BigQuery routine. + + Only the fields listed in ``fields`` are updated. Any field listed but left unset in + ``routine_resource`` is cleared on the server. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:BigQueryUpdateRoutineOperator` + + :param dataset_id: The dataset that owns the routine. + :param routine_id: The identifier of the routine to update. + :param routine_resource: The routine resource (dict or + :class:`~google.cloud.bigquery.routine.Routine`) with the new values for the fields + listed in ``fields``. + :param fields: Properties to update (e.g. ``["body", "description"]``). + :param project_id: Optional. The project that owns the dataset. + :param gcp_conn_id: The connection ID used to connect to Google Cloud. + :param location: The location of the BigQuery dataset. + :param impersonation_chain: Optional service account to impersonate. + :param retry: A retry object used to retry requests. + :param timeout: The amount of time, in seconds, to wait for the request. + """ + + template_fields: Sequence[str] = ( + "project_id", + "dataset_id", + "routine_id", + "routine_resource", + "gcp_conn_id", + "impersonation_chain", + ) + template_fields_renderers = {"routine_resource": "json"} + template_ext: Sequence[str] = (".json", ".sql") + ui_color = BigQueryUIColors.TABLE.value + + def __init__( + self, + *, + dataset_id: str, + routine_id: str, + routine_resource: dict[str, Any] | Routine, + fields: Sequence[str], + project_id: str = PROVIDE_PROJECT_ID, + gcp_conn_id: str = "google_cloud_default", + location: str | None = None, + impersonation_chain: str | Sequence[str] | None = None, + retry: Retry = DEFAULT_RETRY, + timeout: float | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + if not fields: + raise ValueError("`fields` must be a non-empty sequence of routine properties to update.") + self.project_id = project_id + self.dataset_id = dataset_id + self.routine_id = routine_id + self.routine_resource = routine_resource + self.fields = list(fields) + self.gcp_conn_id = gcp_conn_id + self.location = location + self.impersonation_chain = impersonation_chain + self.retry = retry + self.timeout = timeout + + def execute(self, context: Context) -> dict[str, Any]: + hook = BigQueryHook( + gcp_conn_id=self.gcp_conn_id, + location=self.location, + impersonation_chain=self.impersonation_chain, + ) + self.log.info( + "Updating routine %s.%s.%s (fields=%s)", + self.project_id or hook.project_id, + self.dataset_id, + self.routine_id, + self.fields, + ) + routine = hook.update_routine( + routine=self.routine_resource, + fields=self.fields, + dataset_id=self.dataset_id, + routine_id=self.routine_id, + project_id=self.project_id, + retry=self.retry, + timeout=self.timeout, + ) + return routine.to_api_repr() + + +class BigQueryDeleteRoutineOperator(GoogleCloudBaseOperator): + """ + Delete a BigQuery routine. + + .. seealso:: + For more information on how to use this operator, take a look at the guide: + :ref:`howto/operator:BigQueryDeleteRoutineOperator` + + :param dataset_id: The dataset that owns the routine. + :param routine_id: The identifier of the routine to delete. + :param project_id: Optional. The project that owns the dataset. + :param ignore_if_missing: If ``True``, do not fail when the routine does not exist. + Defaults to ``False``. + :param gcp_conn_id: The connection ID used to connect to Google Cloud. + :param location: The location of the BigQuery dataset. + :param impersonation_chain: Optional service account to impersonate. + :param retry: A retry object used to retry requests. + :param timeout: The amount of time, in seconds, to wait for the request. + """ + + template_fields: Sequence[str] = ( + "project_id", + "dataset_id", + "routine_id", + "gcp_conn_id", + "impersonation_chain", + ) + ui_color = BigQueryUIColors.TABLE.value + + def __init__( + self, + *, + dataset_id: str, + routine_id: str, + project_id: str = PROVIDE_PROJECT_ID, Review Comment: The opposite defaults are intentional — they mirror the existing Airflow convention for BigQuery delete operations (`BigQueryDeleteTableOperator` / `BigQueryHook.delete_table`): - **Hook default (`not_found_ok=True`)**: forgiving, so the hook can be composed safely into teardown flows, retries, and idempotent cleanup without needing a try/except around every call. - **Operator default (`ignore_if_missing=False`)**: strict, so DAG authors get a loud failure when their assumptions about routine existence are wrong. Silently ignoring a missing routine in a DAG is usually a bug, not a feature. The operator's `ignore_if_missing=True` flag lets users explicitly opt in to soft-delete when that's what they want — matching how `BigQueryDeleteTableOperator.ignore_if_missing` behaves. Reference (same repo): [`BigQueryDeleteTableOperator`](https://github.com/apache/airflow/blob/main/providers/google/src/airflow/providers/google/cloud/operators/bigquery.py) also defaults to `ignore_if_missing=False` while calling `hook.delete_table(not_found_ok=True)` under the hood. Happy to flip it if the project preference has changed, but I'd want to align both operators (routine + table) rather than diverging from the existing pattern. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
