Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 140ca3a

Browse filesBrowse files
r1br1b
committed
feat: add JSON type, bindparam support
1 parent 736d4a2 commit 140ca3a
Copy full SHA for 140ca3a

File tree

Expand file treeCollapse file tree

5 files changed

+93
-3
lines changed
Filter options
Expand file treeCollapse file tree

5 files changed

+93
-3
lines changed

‎sqlalchemy_bigquery/__init__.py

Copy file name to clipboardExpand all lines: sqlalchemy_bigquery/__init__.py
+2-2Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
import warnings
2424

2525
from .version import __version__
26-
2726
from .base import BigQueryDialect, dialect
2827
from ._types import (
2928
ARRAY,
@@ -37,14 +36,14 @@
3736
FLOAT64,
3837
INT64,
3938
INTEGER,
39+
JSON,
4040
NUMERIC,
4141
RECORD,
4242
STRING,
4343
STRUCT,
4444
TIME,
4545
TIMESTAMP,
4646
)
47-
4847
from . import _versions_helpers
4948

5049
sys_major, sys_minor, sys_micro = _versions_helpers.extract_runtime_version()
@@ -74,6 +73,7 @@
7473
"FLOAT64",
7574
"INT64",
7675
"INTEGER",
76+
"JSON",
7777
"NUMERIC",
7878
"RECORD",
7979
"STRING",

‎sqlalchemy_bigquery/_json.py

Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
import sqlalchemy
2+
3+
4+
class JSON(sqlalchemy.sql.sqltypes.JSON):
5+
def bind_expression(self, bindvalue):
6+
# JSON query parameters have type STRING
7+
# This hook ensures that the rendered expression has type JSON
8+
return sqlalchemy.func.PARSE_JSON(bindvalue, type_=self)

‎sqlalchemy_bigquery/_types.py

Copy file name to clipboardExpand all lines: sqlalchemy_bigquery/_types.py
+3Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
except ImportError: # pragma: NO COVER
2828
pass
2929

30+
from ._json import JSON
3031
from ._struct import STRUCT
3132

3233
_type_map = {
@@ -41,6 +42,7 @@
4142
"FLOAT": sqlalchemy.types.Float,
4243
"INT64": sqlalchemy.types.Integer,
4344
"INTEGER": sqlalchemy.types.Integer,
45+
"JSON": JSON,
4446
"NUMERIC": sqlalchemy.types.Numeric,
4547
"RECORD": STRUCT,
4648
"STRING": sqlalchemy.types.String,
@@ -61,6 +63,7 @@
6163
FLOAT = _type_map["FLOAT"]
6264
INT64 = _type_map["INT64"]
6365
INTEGER = _type_map["INTEGER"]
66+
JSON = _type_map["JSON"]
6467
NUMERIC = _type_map["NUMERIC"]
6568
RECORD = _type_map["RECORD"]
6669
STRING = _type_map["STRING"]

‎sqlalchemy_bigquery/base.py

Copy file name to clipboardExpand all lines: sqlalchemy_bigquery/base.py
+16-1Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@
5959
import re
6060

6161
from .parse_url import parse_url
62-
from . import _helpers, _struct, _types
62+
from . import _helpers, _json, _struct, _types
6363
import sqlalchemy_bigquery_vendored.sqlalchemy.postgresql.base as vendored_postgresql
6464

6565
# Illegal characters is intended to be all characters that are not explicitly
@@ -547,6 +547,13 @@ def visit_bindparam(
547547
bq_type = self.dialect.type_compiler.process(type_)
548548
bq_type = self.__remove_type_parameter(bq_type)
549549

550+
if bq_type == "JSON":
551+
# FIXME: JSON is not a member of `SqlParameterScalarTypes` in the DBAPI
552+
# For now, we hack around this by:
553+
# - Rewriting the bindparam type to STRING
554+
# - Applying a bind expression that converts the parameter back to JSON
555+
bq_type = "STRING"
556+
550557
assert_(param != "%s", f"Unexpected param: {param}")
551558

552559
if bindparam.expanding: # pragma: NO COVER
@@ -641,6 +648,9 @@ def visit_NUMERIC(self, type_, **kw):
641648

642649
visit_DECIMAL = visit_NUMERIC
643650

651+
def visit_JSON(self, type_, **kw):
652+
return "JSON"
653+
644654

645655
class BigQueryDDLCompiler(DDLCompiler):
646656
option_datatype_mapping = {
@@ -1076,6 +1086,7 @@ class BigQueryDialect(DefaultDialect):
10761086
sqlalchemy.sql.sqltypes.TIMESTAMP: BQTimestamp,
10771087
sqlalchemy.sql.sqltypes.ARRAY: BQArray,
10781088
sqlalchemy.sql.sqltypes.Enum: sqlalchemy.sql.sqltypes.Enum,
1089+
sqlalchemy.sql.sqltypes.JSON: _json.JSON,
10791090
}
10801091

10811092
def __init__(
@@ -1086,6 +1097,8 @@ def __init__(
10861097
credentials_info=None,
10871098
credentials_base64=None,
10881099
list_tables_page_size=1000,
1100+
json_serializer=None,
1101+
json_deserializer=None,
10891102
*args,
10901103
**kwargs,
10911104
):
@@ -1098,6 +1111,8 @@ def __init__(
10981111
self.identifier_preparer = self.preparer(self)
10991112
self.dataset_id = None
11001113
self.list_tables_page_size = list_tables_page_size
1114+
self._json_serializer = json_serializer
1115+
self._json_deserializer = json_deserializer
11011116

11021117
@classmethod
11031118
def dbapi(cls):

‎tests/unit/test__json.py

Copy file name to clipboard
+64Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import json
2+
from unittest import mock
3+
4+
import pytest
5+
import sqlalchemy
6+
7+
8+
@pytest.fixture
9+
def json_table(metadata):
10+
from sqlalchemy_bigquery import JSON
11+
12+
return sqlalchemy.Table("json_table", metadata, sqlalchemy.Column("json", JSON))
13+
14+
15+
@pytest.fixture
16+
def json_data():
17+
return {"foo": "bar"}
18+
19+
20+
def test_set_json_serde(faux_conn, metadata, json_table, json_data):
21+
from sqlalchemy_bigquery import JSON
22+
23+
json_serializer = mock.Mock(side_effect=json.dumps)
24+
json_deserializer = mock.Mock(side_effect=json.loads)
25+
26+
engine = sqlalchemy.create_engine(
27+
"bigquery://myproject/mydataset",
28+
json_serializer=json_serializer,
29+
json_deserializer=json_deserializer,
30+
)
31+
32+
json_column = json_table.c.json
33+
34+
process_bind = json_column.type.bind_processor(engine.dialect)
35+
process_bind(json_data)
36+
assert json_serializer.mock_calls == [mock.call(json_data)]
37+
38+
process_result = json_column.type.result_processor(engine.dialect, JSON)
39+
process_result(json.dumps(json_data))
40+
assert json_deserializer.mock_calls == [mock.call(json.dumps(json_data))]
41+
42+
43+
def test_json_create(faux_conn, metadata, json_table, json_data):
44+
expr = sqlalchemy.schema.CreateTable(json_table)
45+
sql = expr.compile(faux_conn.engine).string
46+
assert sql == ("\nCREATE TABLE `json_table` (\n" "\t`json` JSON\n" ") \n\n")
47+
48+
49+
def test_json_insert(faux_conn, metadata, json_table, json_data):
50+
expr = sqlalchemy.insert(json_table).values(json=json_data)
51+
sql = expr.compile(faux_conn.engine).string
52+
assert (
53+
sql == "INSERT INTO `json_table` (`json`) VALUES (PARSE_JSON(%(json:STRING)s))"
54+
)
55+
56+
57+
def test_json_where(faux_conn, metadata, json_table, json_data):
58+
expr = sqlalchemy.select(json_table.c.json).where(json_table.c.json == json_data)
59+
sql = expr.compile(faux_conn.engine).string
60+
assert sql == (
61+
"SELECT `json_table`.`json` \n"
62+
"FROM `json_table` \n"
63+
"WHERE `json_table`.`json` = PARSE_JSON(%(json_1:STRING)s)"
64+
)

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.