Support UUID for SQLAlchemy

Support UUID for SQLAlchemy

Author: lpoulain

tests/integration/test_sqlalchemy_integration.py

@@ -9,6 +9,8 @@
 # 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 uuid
+
 import pytest
 import sqlalchemy as sqla
 from sqlalchemy.sql import and_, not_, or_
@@ -133,6 +135,60 @@ def test_insert(trino_connection):
         metadata.drop_all(engine)
 
 
+@pytest.mark.skipif(
+    sqlalchemy_version() < "2.0",
+    reason="sqlalchemy.Uuid only exists with SQLAlchemy 2.0 and above"
+)
+@pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
+def test_define_and_create_table_uuid(trino_connection):
+    engine, conn = trino_connection
+    if not engine.dialect.has_schema(conn, "test"):
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
+    metadata = sqla.MetaData()
+    try:
+        sqla.Table('users',
+                   metadata,
+                   sqla.Column('guid', sqla.Uuid),
+                   schema="test")
+        metadata.create_all(engine)
+        assert sqla.inspect(engine).has_table('users', schema="test")
+        users = sqla.Table('users', metadata, schema='test', autoload_with=conn)
+        assert_column(users, "guid", sqla.sql.sqltypes.Uuid)
+    finally:
+        metadata.drop_all(engine)
+
+
+@pytest.mark.skipif(
+    sqlalchemy_version() < "2.0",
+    reason="sqlalchemy.Uuid only exists with SQLAlchemy 2.0 and above"
+)
+@pytest.mark.parametrize('trino_connection', ['memory'], indirect=True)
+def test_insert_uuid(trino_connection):
+    engine, conn = trino_connection
+
+    if not engine.dialect.has_schema(conn, "test"):
+        with engine.begin() as connection:
+            connection.execute(sqla.schema.CreateSchema("test"))
+    metadata = sqla.MetaData()
+    try:
+        users = sqla.Table('users',
+                           metadata,
+                           sqla.Column('guid', sqla.Uuid),
+                           schema="test")
+        metadata.create_all(engine)
+        ins = users.insert()
+        guid = uuid.uuid4()
+        conn.execute(ins, {"guid": guid})
+        query = sqla.select(users)
+        result = conn.execute(query)
+        rows = result.fetchall()
+        assert len(rows) == 1
+        assert rows[0] == (guid,)
+    finally:
+        metadata.drop_all(engine)
+
+
 @pytest.mark.skipif(
     sqlalchemy_version() < "1.4",
     reason="columns argument to select() must be a Python list or other iterable"

trino/sqlalchemy/datatype.py

@@ -13,6 +13,7 @@
 import re
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Type, Union
 
+import sqlalchemy
 from sqlalchemy import util
 from sqlalchemy.sql import sqltypes
 from sqlalchemy.sql.type_api import TypeDecorator, TypeEngine
@@ -129,6 +130,9 @@ def get_col_spec(self, **kw):
     # 'tdigest': TDIGEST,
 }
 
+if hasattr(sqlalchemy, "Uuid"):
+    _type_map["uuid"] = sqlalchemy.Uuid
+
 
 def unquote(string: str, quote: str = '"', escape: str = "\\") -> str:
     """