Add JSON handling to dbapi and dialect

tests/integration/test_sqlalchemy_integration.py

@@ -13,6 +13,8 @@
 import sqlalchemy as sqla
 from sqlalchemy.sql import and_, or_, not_
 
+from trino.sqlalchemy.datatype import JSON
+
 
 @pytest.fixture
 def trino_connection(run_trino, request):
@@ -255,3 +257,41 @@ def test_cte(trino_connection):
  result = conn.execute(s)
  rows = result.fetchall()
  assert len(rows) == 15
+
+
+@pytest.mark.parametrize(
+ 'trino_connection,json_object',
+ [
+ ('memory', None),
+ ('memory', 1),
+ ('memory', 'test'),
+ ('memory', [1, 'test']),
+ ('memory', {'test': 1}),
+ ],
+ indirect=['trino_connection']
+)
+def test_json_column(trino_connection, json_object):
+ engine, conn = trino_connection
+
+ if not engine.dialect.has_schema(engine, "test"):
+ engine.execute(sqla.schema.CreateSchema("test"))
+ metadata = sqla.MetaData()
+
+ try:
+ table_with_json = sqla.Table(
+ 'table_with_json',
+ metadata,
+ sqla.Column('id', sqla.Integer),
+ sqla.Column('json_column', JSON),
+ schema="test"
+ )
+ metadata.create_all(engine)
+ ins = table_with_json.insert()
+ conn.execute(ins, {"id": 1, "json_column": json_object})
+ query = sqla.select(table_with_json)
+ result = conn.execute(query)
+ rows = result.fetchall()
+ assert len(rows) == 1
+ assert rows[0] == (1, json_object)
+ finally:
+ metadata.drop_all(engine)

trino/sqlalchemy/compiler.py

@@ -206,6 +206,9 @@ def visit_TIME(self, type_, **kw):
  datatype += " WITH TIME ZONE"
  return datatype
 
+ def visit_JSON(self, type_, **kw):
+ return 'JSON'
+
 
 class TrinoIdentifierPreparer(compiler.IdentifierPreparer):
  reserved_words = RESERVED_WORDS

trino/sqlalchemy/datatype.py

@@ -9,12 +9,14 @@
 # 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.
+import json
 import re
 from typing import Iterator, List, Optional, Tuple, Type, Union, Dict, Any
 
 from sqlalchemy import util
 from sqlalchemy.sql import sqltypes
-from sqlalchemy.sql.type_api import TypeEngine
+from sqlalchemy.sql.type_api import TypeDecorator, TypeEngine
+from sqlalchemy.types import String
 
 SQLType = Union[TypeEngine, Type[TypeEngine]]
 
@@ -71,6 +73,19 @@ def __init__(self, precision=None, timezone=False):
  self.precision = precision
 
 
+class JSON(TypeDecorator):
+ impl = String
+
+ def process_bind_param(self, value, dialect):
+ return json.dumps(value)
+
+ def process_result_value(self, value, dialect):
+ return json.loads(value)
+
+ def get_col_spec(self, **kw):
+ return 'JSON'
+
+
 # https://trino.io/docs/current/language/types.html
 _type_map = {
  # === Boolean ===
@@ -90,7 +105,7 @@ def __init__(self, precision=None, timezone=False):
  "varchar": sqltypes.VARCHAR,
  "char": sqltypes.CHAR,
  "varbinary": sqltypes.VARBINARY,
- "json": sqltypes.JSON,
+ "json": JSON,
  # === Date and time ===
  "date": sqltypes.DATE,
  "time": TIME,