| from sqlalchemy.engine import create_engine |
| from pyhive.tests.sqlalchemy_test_case import SqlAlchemyTestCase |
| from pyhive.tests.sqlalchemy_test_case import with_engine_connection |
| from sqlalchemy.exc import NoSuchTableError, DatabaseError |
| from sqlalchemy.schema import MetaData, Table, Column |
| from sqlalchemy.types import String |
| from sqlalchemy.sql import text |
| from sqlalchemy import types |
| from decimal import Decimal |
| |
| import unittest |
| import contextlib |
| |
| |
| class TestSqlAlchemyTrino(unittest.TestCase, SqlAlchemyTestCase): |
| def create_engine(self): |
| return create_engine('trino+pyhive://localhost:18080/hive/default?source={}'.format(self.id())) |
| |
| def test_bad_format(self): |
| self.assertRaises( |
| ValueError, |
| lambda: create_engine('trino+pyhive://localhost:18080/hive/default/what'), |
| ) |
| |
| @with_engine_connection |
| def test_reflect_select(self, engine, connection): |
| """reflecttable should be able to fill in a table from the name""" |
| one_row_complex = Table('one_row_complex', MetaData(), autoload_with=engine) |
| # Presto ignores the union column |
| self.assertEqual(len(one_row_complex.c), 15 - 1) |
| self.assertIsInstance(one_row_complex.c.string, Column) |
| rows = connection.execute(one_row_complex.select()).fetchall() |
| self.assertEqual(len(rows), 1) |
| self.assertEqual(list(rows[0]), [ |
| True, |
| 127, |
| 32767, |
| 2147483647, |
| 9223372036854775807, |
| 0.5, |
| 0.25, |
| 'a string', |
| '1970-01-01 00:00:00.000', |
| b'123', |
| [1, 2], |
| {"1": 2, "3": 4}, |
| [1, 2], |
| Decimal('0.1'), |
| ]) |
| |
| self.assertIsInstance(one_row_complex.c.boolean.type, types.Boolean) |
| self.assertIsInstance(one_row_complex.c.tinyint.type, types.Integer) |
| self.assertIsInstance(one_row_complex.c.smallint.type, types.Integer) |
| self.assertIsInstance(one_row_complex.c.int.type, types.Integer) |
| self.assertIsInstance(one_row_complex.c.bigint.type, types.BigInteger) |
| self.assertIsInstance(one_row_complex.c.float.type, types.Float) |
| self.assertIsInstance(one_row_complex.c.double.type, types.Float) |
| self.assertIsInstance(one_row_complex.c.string.type, String) |
| self.assertIsInstance(one_row_complex.c.timestamp.type, types.NullType) |
| self.assertIsInstance(one_row_complex.c.binary.type, types.VARBINARY) |
| self.assertIsInstance(one_row_complex.c.array.type, types.NullType) |
| self.assertIsInstance(one_row_complex.c.map.type, types.NullType) |
| self.assertIsInstance(one_row_complex.c.struct.type, types.NullType) |
| self.assertIsInstance(one_row_complex.c.decimal.type, types.NullType) |
| |
| @with_engine_connection |
| def test_reflect_no_such_table(self, engine, connection): |
| """reflecttable should throw an exception on an invalid table""" |
| self.assertRaises( |
| NoSuchTableError, |
| lambda: Table('this_does_not_exist', MetaData(), autoload_with=engine)) |
| self.assertRaises( |
| DatabaseError, |
| lambda: Table('this_does_not_exist', MetaData(schema="also_does_not_exist"), autoload_with=engine)) |
| |
| def test_url_default(self): |
| engine = create_engine('trino+pyhive://localhost:18080/hive') |
| try: |
| with contextlib.closing(engine.connect()) as connection: |
| self.assertEqual(connection.execute(text('SELECT 1 AS foobar FROM one_row')).scalar(), 1) |
| finally: |
| engine.dispose() |
| |
| @with_engine_connection |
| def test_reserved_words(self, engine, connection): |
| """Trino uses double quotes, not backticks""" |
| # Use keywords for the table/column name |
| fake_table = Table('select', MetaData(), Column('current_timestamp', String)) |
| query = str(fake_table.select().where(fake_table.c.current_timestamp == 'a').compile(engine)) |
| self.assertIn('"select"', query) |
| self.assertIn('"current_timestamp"', query) |
| self.assertNotIn('`select`', query) |
| self.assertNotIn('`current_timestamp`', query) |
