diff --git a/packages/bigframes/bigframes/bigquery/__init__.py b/packages/bigframes/bigframes/bigquery/__init__.py index 86a45546b748..d643b986ac8e 100644 --- a/packages/bigframes/bigframes/bigquery/__init__.py +++ b/packages/bigframes/bigframes/bigquery/__init__.py @@ -72,6 +72,11 @@ st_regionstats, st_simplify, ) +from bigframes.bigquery._operations.global_namespace.aead_encryption import ( + deterministic_decrypt_bytes, + deterministic_decrypt_string, + deterministic_encrypt, +) from bigframes.bigquery._operations.io import load_data from bigframes.bigquery._operations.json import ( json_extract, @@ -121,6 +126,10 @@ st_length, st_regionstats, st_simplify, + # deterministic encryption ops + deterministic_decrypt_bytes, + deterministic_decrypt_string, + deterministic_encrypt, # json ops json_extract, json_extract_array, @@ -179,6 +188,10 @@ "st_length", "st_regionstats", "st_simplify", + # deterministic encryption ops + "deterministic_decrypt_bytes", + "deterministic_decrypt_string", + "deterministic_encrypt", # json ops "json_extract", "json_extract_array", diff --git a/packages/bigframes/bigframes/bigquery/_operations/global_namespace/aead_encryption.py b/packages/bigframes/bigframes/bigquery/_operations/global_namespace/aead_encryption.py new file mode 100644 index 000000000000..b1522a70d73e --- /dev/null +++ b/packages/bigframes/bigframes/bigquery/_operations/global_namespace/aead_encryption.py @@ -0,0 +1,127 @@ +# Copyright 2026 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. +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml +# by the script: scripts/generate_bigframes_bigquery.py + +from __future__ import annotations + +import datetime +from typing import Any, Literal, Optional, TypeVar, Union + +import bigframes.bigquery._googlesql +import bigframes.core.col +import bigframes.core.expression as ex +import bigframes.core.sentinels as sentinels +import bigframes.operations as ops +import bigframes.series as series +from bigframes import dtypes +from bigframes.operations import googlesql + +T = TypeVar("T", series.Series, bigframes.core.col.Expression) + +_DETERMINISTIC_DECRYPT_BYTES_OP = googlesql.GoogleSqlScalarOp( + "DETERMINISTIC_DECRYPT_BYTES", + args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()), + signature=lambda *args: dtypes.BYTES_DTYPE, +) +_DETERMINISTIC_DECRYPT_STRING_OP = googlesql.GoogleSqlScalarOp( + "DETERMINISTIC_DECRYPT_STRING", + args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()), + signature=lambda *args: dtypes.STRING_DTYPE, +) +_DETERMINISTIC_ENCRYPT_OP = googlesql.GoogleSqlScalarOp( + "DETERMINISTIC_ENCRYPT", + args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()), + signature=lambda *args: dtypes.BYTES_DTYPE, +) + + +def deterministic_decrypt_bytes( + keyset: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict], + ], + ciphertext: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + additional_data: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], +) -> T: + """Uses the matching key from `keyset` to decrypt `ciphertext` and verifies the integrity of the data using `additional_data`. Returns an error if decryption fails.""" + return bigframes.bigquery._googlesql.apply_googlesql_scalar_op( + _DETERMINISTIC_DECRYPT_BYTES_OP, + keyset, + ciphertext, + additional_data, + ) # type: ignore + + +def deterministic_decrypt_string( + keyset: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict], + ], + ciphertext: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes], + ], + additional_data: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], str], + ], +) -> T: + """Like `DETERMINISTIC_DECRYPT_BYTES`, but where plaintext is of type STRING.""" + return bigframes.bigquery._googlesql.apply_googlesql_scalar_op( + _DETERMINISTIC_DECRYPT_STRING_OP, + keyset, + ciphertext, + additional_data, + ) # type: ignore + + +def deterministic_encrypt( + keyset: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict], + ], + plaintext: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ], + additional_data: Union[ + T, + bigframes.core.col.Expression, + Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str], + ], +) -> T: + """Encrypts `plaintext` using the primary cryptographic key in `keyset` using deterministic AEAD. The algorithm of the primary key must be `DETERMINISTIC_AEAD_AES_SIV_CMAC_256`. Binds the ciphertext to the context defined by `additional_data`. Returns `NULL` if any input is `NULL`.""" + return bigframes.bigquery._googlesql.apply_googlesql_scalar_op( + _DETERMINISTIC_ENCRYPT_OP, + keyset, + plaintext, + additional_data, + ) # type: ignore diff --git a/packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml b/packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml new file mode 100644 index 000000000000..ffd26e5e0e7b --- /dev/null +++ b/packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml @@ -0,0 +1,131 @@ +urn: extension:google:bq_scalar_functions +scalar_functions: + - name: "deterministic_decrypt_bytes" + description: "Uses the matching key from `keyset` to decrypt `ciphertext` and verifies the integrity of the data using `additional_data`. Returns an error if decryption fails." + impls: + # Signature: deterministic_decrypt_bytes:vbin_vbin_vbin + - args: + - name: "keyset" + value: binary + optional: false + keyword_only: false + - name: "ciphertext" + value: binary + optional: false + keyword_only: false + - name: "additional_data" + value: binary + optional: false + keyword_only: false + return: binary + # Signature: deterministic_decrypt_bytes:struct_vbin_vbin + - args: + - name: "keyset" + value: struct + optional: false + keyword_only: false + - name: "ciphertext" + value: binary + optional: false + keyword_only: false + - name: "additional_data" + value: binary + optional: false + keyword_only: false + return: binary + - name: "deterministic_decrypt_string" + description: "Like `DETERMINISTIC_DECRYPT_BYTES`, but where plaintext is of type STRING." + impls: + # Signature: deterministic_decrypt_string:vbin_vbin_str + - args: + - name: "keyset" + value: binary + optional: false + keyword_only: false + - name: "ciphertext" + value: binary + optional: false + keyword_only: false + - name: "additional_data" + value: string + optional: false + keyword_only: false + return: string + # Signature: deterministic_decrypt_string:struct_vbin_str + - args: + - name: "keyset" + value: struct + optional: false + keyword_only: false + - name: "ciphertext" + value: binary + optional: false + keyword_only: false + - name: "additional_data" + value: string + optional: false + keyword_only: false + return: string + - name: "deterministic_encrypt" + description: "Encrypts `plaintext` using the primary cryptographic key in `keyset` using deterministic AEAD. The algorithm of the primary key must be `DETERMINISTIC_AEAD_AES_SIV_CMAC_256`. Binds the ciphertext to the context defined by `additional_data`. Returns `NULL` if any input is `NULL`." + impls: + # Signature: deterministic_encrypt:vbin_str_str + - args: + - name: "keyset" + value: binary + optional: false + keyword_only: false + - name: "plaintext" + value: string + optional: false + keyword_only: false + - name: "additional_data" + value: string + optional: false + keyword_only: false + return: binary + # Signature: deterministic_encrypt:vbin_vbin_vbin + - args: + - name: "keyset" + value: binary + optional: false + keyword_only: false + - name: "plaintext" + value: binary + optional: false + keyword_only: false + - name: "additional_data" + value: binary + optional: false + keyword_only: false + return: binary + # Signature: deterministic_encrypt:struct_str_str + - args: + - name: "keyset" + value: struct + optional: false + keyword_only: false + - name: "plaintext" + value: string + optional: false + keyword_only: false + - name: "additional_data" + value: string + optional: false + keyword_only: false + return: binary + # Signature: deterministic_encrypt:struct_vbin_vbin + - args: + - name: "keyset" + value: struct + optional: false + keyword_only: false + - name: "plaintext" + value: binary + optional: false + keyword_only: false + - name: "additional_data" + value: binary + optional: false + keyword_only: false + return: binary diff --git a/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_bytes/out.sql b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_bytes/out.sql new file mode 100644 index 000000000000..1a567e3b002b --- /dev/null +++ b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_bytes/out.sql @@ -0,0 +1,4 @@ +SELECT + `rowindex`, + DETERMINISTIC_DECRYPT_BYTES(`bytes_col`, `bytes_col`, `bytes_col`) AS `0` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` diff --git a/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_string/out.sql b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_string/out.sql new file mode 100644 index 000000000000..84625951e6da --- /dev/null +++ b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_string/out.sql @@ -0,0 +1,4 @@ +SELECT + `rowindex`, + DETERMINISTIC_DECRYPT_STRING(`bytes_col`, `bytes_col`, `string_col`) AS `0` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` diff --git a/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_encrypt/out.sql b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_encrypt/out.sql new file mode 100644 index 000000000000..659097ded956 --- /dev/null +++ b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_encrypt/out.sql @@ -0,0 +1,4 @@ +SELECT + `rowindex`, + DETERMINISTIC_ENCRYPT(`bytes_col`, `bytes_col`, `bytes_col`) AS `0` +FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0` diff --git a/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/test_aead_encryption.py b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/test_aead_encryption.py new file mode 100644 index 000000000000..d7e14ccb4bef --- /dev/null +++ b/packages/bigframes/tests/unit/bigquery/_operations/global_namespace/test_aead_encryption.py @@ -0,0 +1,56 @@ +# Copyright 2026 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. +# +# DO NOT MODIFY THIS FILE DIRECTLY. +# This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml +# by the script: scripts/generate_bigframes_bigquery.py + +from typing import cast + +import pytest + +import bigframes.bigquery._operations.global_namespace.aead_encryption as aead_encryption +import bigframes.pandas as bpd + +pytest.importorskip("pytest_snapshot") + + +def test_deterministic_decrypt_bytes(scalar_types_df: bpd.DataFrame, snapshot): + result = aead_encryption.deterministic_decrypt_bytes( + cast(bpd.Series, scalar_types_df["bytes_col"]), + cast(bpd.Series, scalar_types_df["bytes_col"]), + cast(bpd.Series, scalar_types_df["bytes_col"]), + ).to_frame() + + snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql") + + +def test_deterministic_decrypt_string(scalar_types_df: bpd.DataFrame, snapshot): + result = aead_encryption.deterministic_decrypt_string( + cast(bpd.Series, scalar_types_df["bytes_col"]), + cast(bpd.Series, scalar_types_df["bytes_col"]), + cast(bpd.Series, scalar_types_df["string_col"]), + ).to_frame() + + snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql") + + +def test_deterministic_encrypt(scalar_types_df: bpd.DataFrame, snapshot): + result = aead_encryption.deterministic_encrypt( + cast(bpd.Series, scalar_types_df["bytes_col"]), + cast(bpd.Series, scalar_types_df["bytes_col"]), + cast(bpd.Series, scalar_types_df["bytes_col"]), + ).to_frame() + + snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql")