Skip to content

Commit 85f3672

Browse files
tswasttswast
authored
feat: add bigframes.bigquery.deterministic_decrypt* and bigframes.bigquery.deterministic_encrypt functions (#17212)
🦕
1 parent 994a22d commit 85f3672

7 files changed

Lines changed: 339 additions & 0 deletions

File tree

‎packages/bigframes/bigframes/bigquery/__init__.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,11 @@
7272
st_regionstats,
7373
st_simplify,
7474
)
75+
from bigframes.bigquery._operations.global_namespace.aead_encryption import (
76+
deterministic_decrypt_bytes,
77+
deterministic_decrypt_string,
78+
deterministic_encrypt,
79+
)
7580
from bigframes.bigquery._operations.io import load_data
7681
from bigframes.bigquery._operations.json import (
7782
json_extract,
@@ -121,6 +126,10 @@
121126
st_length,
122127
st_regionstats,
123128
st_simplify,
129+
# deterministic encryption ops
130+
deterministic_decrypt_bytes,
131+
deterministic_decrypt_string,
132+
deterministic_encrypt,
124133
# json ops
125134
json_extract,
126135
json_extract_array,
@@ -179,6 +188,10 @@
179188
"st_length",
180189
"st_regionstats",
181190
"st_simplify",
191+
# deterministic encryption ops
192+
"deterministic_decrypt_bytes",
193+
"deterministic_decrypt_string",
194+
"deterministic_encrypt",
182195
# json ops
183196
"json_extract",
184197
"json_extract_array",

‎packages/bigframes/bigframes/bigquery/_operations/global_namespace/aead_encryption.py‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# DO NOT MODIFY THIS FILE DIRECTLY.
16+
# This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml
17+
# by the script: scripts/generate_bigframes_bigquery.py
18+
19+
from __future__ import annotations
20+
21+
import datetime
22+
from typing import Any, Literal, Optional, TypeVar, Union
23+
24+
import bigframes.bigquery._googlesql
25+
import bigframes.core.col
26+
import bigframes.core.expression as ex
27+
import bigframes.core.sentinels as sentinels
28+
import bigframes.operations as ops
29+
import bigframes.series as series
30+
from bigframes import dtypes
31+
from bigframes.operations import googlesql
32+
33+
T = TypeVar("T", series.Series, bigframes.core.col.Expression)
34+
35+
_DETERMINISTIC_DECRYPT_BYTES_OP = googlesql.GoogleSqlScalarOp(
36+
"DETERMINISTIC_DECRYPT_BYTES",
37+
args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()),
38+
signature=lambda *args: dtypes.BYTES_DTYPE,
39+
)
40+
_DETERMINISTIC_DECRYPT_STRING_OP = googlesql.GoogleSqlScalarOp(
41+
"DETERMINISTIC_DECRYPT_STRING",
42+
args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()),
43+
signature=lambda *args: dtypes.STRING_DTYPE,
44+
)
45+
_DETERMINISTIC_ENCRYPT_OP = googlesql.GoogleSqlScalarOp(
46+
"DETERMINISTIC_ENCRYPT",
47+
args=(googlesql.ArgSpec(), googlesql.ArgSpec(), googlesql.ArgSpec()),
48+
signature=lambda *args: dtypes.BYTES_DTYPE,
49+
)
50+
51+
52+
def deterministic_decrypt_bytes(
53+
keyset: Union[
54+
T,
55+
bigframes.core.col.Expression,
56+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict],
57+
],
58+
ciphertext: Union[
59+
T,
60+
bigframes.core.col.Expression,
61+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes],
62+
],
63+
additional_data: Union[
64+
T,
65+
bigframes.core.col.Expression,
66+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes],
67+
],
68+
) -> T:
69+
"""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."""
70+
return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
71+
_DETERMINISTIC_DECRYPT_BYTES_OP,
72+
keyset,
73+
ciphertext,
74+
additional_data,
75+
) # type: ignore
76+
77+
78+
def deterministic_decrypt_string(
79+
keyset: Union[
80+
T,
81+
bigframes.core.col.Expression,
82+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict],
83+
],
84+
ciphertext: Union[
85+
T,
86+
bigframes.core.col.Expression,
87+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes],
88+
],
89+
additional_data: Union[
90+
T,
91+
bigframes.core.col.Expression,
92+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], str],
93+
],
94+
) -> T:
95+
"""Like `DETERMINISTIC_DECRYPT_BYTES`, but where plaintext is of type STRING."""
96+
return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
97+
_DETERMINISTIC_DECRYPT_STRING_OP,
98+
keyset,
99+
ciphertext,
100+
additional_data,
101+
) # type: ignore
102+
103+
104+
def deterministic_encrypt(
105+
keyset: Union[
106+
T,
107+
bigframes.core.col.Expression,
108+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, dict],
109+
],
110+
plaintext: Union[
111+
T,
112+
bigframes.core.col.Expression,
113+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str],
114+
],
115+
additional_data: Union[
116+
T,
117+
bigframes.core.col.Expression,
118+
Union[Literal[sentinels.Sentinel.ARGUMENT_DEFAULT], bytes, str],
119+
],
120+
) -> T:
121+
"""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`."""
122+
return bigframes.bigquery._googlesql.apply_googlesql_scalar_op(
123+
_DETERMINISTIC_ENCRYPT_OP,
124+
keyset,
125+
plaintext,
126+
additional_data,
127+
) # type: ignore

‎packages/bigframes/scripts/data/sql-functions/global_namespace/aead_encryption.yaml‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,131 @@
1+
urn: extension:google:bq_scalar_functions
2+
scalar_functions:
3+
- name: "deterministic_decrypt_bytes"
4+
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."
5+
impls:
6+
# Signature: deterministic_decrypt_bytes:vbin_vbin_vbin
7+
- args:
8+
- name: "keyset"
9+
value: binary
10+
optional: false
11+
keyword_only: false
12+
- name: "ciphertext"
13+
value: binary
14+
optional: false
15+
keyword_only: false
16+
- name: "additional_data"
17+
value: binary
18+
optional: false
19+
keyword_only: false
20+
return: binary
21+
# Signature: deterministic_decrypt_bytes:struct_vbin_vbin
22+
- args:
23+
- name: "keyset"
24+
value: struct
25+
optional: false
26+
keyword_only: false
27+
- name: "ciphertext"
28+
value: binary
29+
optional: false
30+
keyword_only: false
31+
- name: "additional_data"
32+
value: binary
33+
optional: false
34+
keyword_only: false
35+
return: binary
36+
- name: "deterministic_decrypt_string"
37+
description: "Like `DETERMINISTIC_DECRYPT_BYTES`, but where plaintext is of type STRING."
38+
impls:
39+
# Signature: deterministic_decrypt_string:vbin_vbin_str
40+
- args:
41+
- name: "keyset"
42+
value: binary
43+
optional: false
44+
keyword_only: false
45+
- name: "ciphertext"
46+
value: binary
47+
optional: false
48+
keyword_only: false
49+
- name: "additional_data"
50+
value: string
51+
optional: false
52+
keyword_only: false
53+
return: string
54+
# Signature: deterministic_decrypt_string:struct_vbin_str
55+
- args:
56+
- name: "keyset"
57+
value: struct
58+
optional: false
59+
keyword_only: false
60+
- name: "ciphertext"
61+
value: binary
62+
optional: false
63+
keyword_only: false
64+
- name: "additional_data"
65+
value: string
66+
optional: false
67+
keyword_only: false
68+
return: string
69+
- name: "deterministic_encrypt"
70+
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`."
71+
impls:
72+
# Signature: deterministic_encrypt:vbin_str_str
73+
- args:
74+
- name: "keyset"
75+
value: binary
76+
optional: false
77+
keyword_only: false
78+
- name: "plaintext"
79+
value: string
80+
optional: false
81+
keyword_only: false
82+
- name: "additional_data"
83+
value: string
84+
optional: false
85+
keyword_only: false
86+
return: binary
87+
# Signature: deterministic_encrypt:vbin_vbin_vbin
88+
- args:
89+
- name: "keyset"
90+
value: binary
91+
optional: false
92+
keyword_only: false
93+
- name: "plaintext"
94+
value: binary
95+
optional: false
96+
keyword_only: false
97+
- name: "additional_data"
98+
value: binary
99+
optional: false
100+
keyword_only: false
101+
return: binary
102+
# Signature: deterministic_encrypt:struct_str_str
103+
- args:
104+
- name: "keyset"
105+
value: struct
106+
optional: false
107+
keyword_only: false
108+
- name: "plaintext"
109+
value: string
110+
optional: false
111+
keyword_only: false
112+
- name: "additional_data"
113+
value: string
114+
optional: false
115+
keyword_only: false
116+
return: binary
117+
# Signature: deterministic_encrypt:struct_vbin_vbin
118+
- args:
119+
- name: "keyset"
120+
value: struct
121+
optional: false
122+
keyword_only: false
123+
- name: "plaintext"
124+
value: binary
125+
optional: false
126+
keyword_only: false
127+
- name: "additional_data"
128+
value: binary
129+
optional: false
130+
keyword_only: false
131+
return: binary

‎packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_bytes/out.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
SELECT
2+
`rowindex`,
3+
DETERMINISTIC_DECRYPT_BYTES(`bytes_col`, `bytes_col`, `bytes_col`) AS `0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_decrypt_string/out.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
SELECT
2+
`rowindex`,
3+
DETERMINISTIC_DECRYPT_STRING(`bytes_col`, `bytes_col`, `string_col`) AS `0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/bigquery/_operations/global_namespace/snapshots/test_aead_encryption/test_deterministic_encrypt/out.sql‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
SELECT
2+
`rowindex`,
3+
DETERMINISTIC_ENCRYPT(`bytes_col`, `bytes_col`, `bytes_col`) AS `0`
4+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/bigquery/_operations/global_namespace/test_aead_encryption.py‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
# DO NOT MODIFY THIS FILE DIRECTLY.
16+
# This file was generated from: scripts/data/sql-functions/global_namespace/aead_encryption.yaml
17+
# by the script: scripts/generate_bigframes_bigquery.py
18+
19+
from typing import cast
20+
21+
import pytest
22+
23+
import bigframes.bigquery._operations.global_namespace.aead_encryption as aead_encryption
24+
import bigframes.pandas as bpd
25+
26+
pytest.importorskip("pytest_snapshot")
27+
28+
29+
def test_deterministic_decrypt_bytes(scalar_types_df: bpd.DataFrame, snapshot):
30+
result = aead_encryption.deterministic_decrypt_bytes(
31+
cast(bpd.Series, scalar_types_df["bytes_col"]),
32+
cast(bpd.Series, scalar_types_df["bytes_col"]),
33+
cast(bpd.Series, scalar_types_df["bytes_col"]),
34+
).to_frame()
35+
36+
snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql")
37+
38+
39+
def test_deterministic_decrypt_string(scalar_types_df: bpd.DataFrame, snapshot):
40+
result = aead_encryption.deterministic_decrypt_string(
41+
cast(bpd.Series, scalar_types_df["bytes_col"]),
42+
cast(bpd.Series, scalar_types_df["bytes_col"]),
43+
cast(bpd.Series, scalar_types_df["string_col"]),
44+
).to_frame()
45+
46+
snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql")
47+
48+
49+
def test_deterministic_encrypt(scalar_types_df: bpd.DataFrame, snapshot):
50+
result = aead_encryption.deterministic_encrypt(
51+
cast(bpd.Series, scalar_types_df["bytes_col"]),
52+
cast(bpd.Series, scalar_types_df["bytes_col"]),
53+
cast(bpd.Series, scalar_types_df["bytes_col"]),
54+
).to_frame()
55+
56+
snapshot.assert_match(result.sql.rstrip() + "\n", "out.sql")

0 commit comments

Comments
 (0)