mirror of
https://github.com/hanzo-ds/driver.git
synced 2026-10-30 20:32:40 +00:00
zeekay
073c1bf61e
The wire protocol was the last holdout and it was not cosmetic: the Hanzo Datastore server reads X-Datastore-User, and answers a request bearing X-ClickHouse-User with 401. Verified against the live server, both directions. So a client still sending the old headers is not merely off-brand, it cannot authenticate over HTTP. Everything else follows: types, dialects, error modules, docstrings and doc links now name Datastore, and doc URLs point at docs.hanzo.ai/datastore rather than the upstream vendor's site. Zero occurrences of the old name remain. Still cut from the pinned upstream version, so behaviour is unchanged.
497 lines
19 KiB
Python
497 lines
19 KiB
Python
import types
|
|
from collections import namedtuple
|
|
from contextlib import contextmanager
|
|
import socket
|
|
from unittest.mock import patch
|
|
|
|
from datastore_driver import connect
|
|
from datastore_driver.dbapi import (
|
|
ProgrammingError, InterfaceError, OperationalError
|
|
)
|
|
from datastore_driver.dbapi.extras import DictCursor, NamedTupleCursor
|
|
from tests.testcase import BaseTestCase
|
|
|
|
|
|
class DBAPITestCaseBase(BaseTestCase):
|
|
def create_connection(self, **kwargs):
|
|
kwargs.setdefault('user', self.user)
|
|
kwargs.setdefault('password', self.password)
|
|
kwargs.setdefault('host', self.host)
|
|
kwargs.setdefault('port', self.port)
|
|
kwargs.setdefault('database', self.database)
|
|
|
|
return connect(**kwargs)
|
|
|
|
@contextmanager
|
|
def created_connection(self, **kwargs):
|
|
connection = self.create_connection(**kwargs)
|
|
try:
|
|
yield connection
|
|
finally:
|
|
connection.close()
|
|
|
|
@contextmanager
|
|
def created_cursor(self, **kwargs):
|
|
cursor_kwargs = kwargs.pop('cursor_kwargs', {})
|
|
with self.created_connection(**kwargs) as connection:
|
|
cursor = connection.cursor(**cursor_kwargs)
|
|
|
|
try:
|
|
yield cursor
|
|
finally:
|
|
cursor.close()
|
|
|
|
|
|
class DBAPITestCase(DBAPITestCaseBase):
|
|
def test_no_host_and_dsn(self):
|
|
with self.assertRaises(ValueError) as e:
|
|
self.create_connection(host=None)
|
|
self.assertEqual(str(e.exception), 'host or dsn is required')
|
|
|
|
def test_simple(self):
|
|
with self.created_connection() as connection:
|
|
cursor = connection.cursor()
|
|
rv = cursor.execute('SELECT 1')
|
|
self.assertIsNone(rv)
|
|
self.assertEqual(cursor.fetchall(), [(1, )])
|
|
|
|
def test_from_dsn(self):
|
|
connection = connect(
|
|
'datastore://{user}:{password}@{host}:{port}/{database}'.format(
|
|
user=self.user, password=self.password,
|
|
host=self.host, port=self.port, database=self.database
|
|
)
|
|
)
|
|
cursor = connection.cursor()
|
|
rv = cursor.execute('SELECT 1')
|
|
self.assertIsNone(rv)
|
|
self.assertEqual(cursor.fetchall(), [(1, )])
|
|
connection.close()
|
|
|
|
def test_connect_default_params(self):
|
|
connection = connect(host=self.host)
|
|
cursor = connection.cursor()
|
|
|
|
rv = cursor.execute('SELECT 1')
|
|
self.assertIsNone(rv)
|
|
self.assertEqual(cursor.fetchall(), [(1, )])
|
|
connection.close()
|
|
|
|
def test_execute_fetchone(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, list)
|
|
|
|
self.assertEqual(cursor.fetchone(), (0, ))
|
|
self.assertEqual(cursor.fetchone(), (1, ))
|
|
self.assertEqual(cursor.fetchone(), (2, ))
|
|
self.assertEqual(cursor.fetchone(), (3, ))
|
|
self.assertEqual(cursor.fetchone(), None)
|
|
|
|
def test_execute_fetchmany(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, list)
|
|
self.assertEqual(cursor.fetchmany(), [(0, )])
|
|
self.assertEqual(cursor.fetchmany(None), [(1, )])
|
|
self.assertEqual(cursor.fetchmany(0), [])
|
|
self.assertEqual(cursor.fetchmany(-1), [(2, ), (3, )])
|
|
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.fetchmany(1), [(0, )])
|
|
self.assertEqual(cursor.fetchmany(2), [(1, ), (2, )])
|
|
self.assertEqual(cursor.fetchmany(3), [(3, )])
|
|
self.assertEqual(cursor.fetchmany(3), [])
|
|
|
|
cursor.arraysize = 2
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.fetchmany(), [(0,), (1,)])
|
|
self.assertEqual(cursor.fetchmany(), [(2,), (3,)])
|
|
|
|
def test_execute_fetchall(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.rowcount, 4)
|
|
self.assertEqual(cursor.fetchall(), [(0,), (1,), (2,), (3,)])
|
|
|
|
def test_executemany(self):
|
|
with self.created_cursor() as cursor, self.create_table('a UInt32'):
|
|
data = [(0, ), (1, ), (2, )]
|
|
rv = cursor.executemany('INSERT INTO test VALUES', data)
|
|
self.assertIsNone(rv, None)
|
|
self.assertEqual(cursor.rowcount, 3)
|
|
|
|
cursor.execute('SELECT * FROM test')
|
|
self.assertEqual(cursor.fetchall(), data)
|
|
|
|
def test_fake_transactions(self):
|
|
with self.created_connection() as connection:
|
|
connection.commit()
|
|
connection.rollback()
|
|
|
|
def test_exception_execute(self):
|
|
with self.created_cursor(user='wrong_user') as cursor:
|
|
with self.assertRaises(OperationalError) as e:
|
|
cursor.execute('SELECT 1')
|
|
code = 516 if self.server_version > (20, ) else 192
|
|
self.assertIn('Code: {}'.format(code), str(e.exception))
|
|
|
|
def test_exception_executemany(self):
|
|
with self.created_cursor(user='wrong_user') as cursor:
|
|
with self.assertRaises(OperationalError) as e:
|
|
cursor.executemany('INSERT INTO test VALUES', [(0, )])
|
|
code = 516 if self.server_version > (20, ) else 192
|
|
self.assertIn('Code: {}'.format(code), str(e.exception))
|
|
self.assertEqual(cursor.rowcount, -1)
|
|
|
|
def test_rowcount_insert_from_select(self):
|
|
with self.created_cursor() as cursor, self.create_table('a UInt8'):
|
|
cursor.execute(
|
|
'INSERT INTO test '
|
|
'SELECT number FROM system.numbers LIMIT 4'
|
|
)
|
|
self.assertEqual(cursor.rowcount, -1)
|
|
|
|
def test_execute_insert(self):
|
|
with self.created_cursor() as cursor, self.create_table('a UInt8'):
|
|
cursor.execute('INSERT INTO test VALUES', [[4]])
|
|
self.assertEqual(cursor.rowcount, 1)
|
|
|
|
def test_description(self):
|
|
with self.created_cursor() as cursor:
|
|
self.assertIsNone(cursor.description)
|
|
cursor.execute('SELECT CAST(1 AS UInt32) AS test')
|
|
desc = cursor.description
|
|
self.assertEqual(len(desc), 1)
|
|
self.assertEqual(desc[0].name, 'test')
|
|
self.assertEqual(desc[0].type_code, 'UInt32')
|
|
|
|
def test_pep249_sizes(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.setinputsizes(0)
|
|
cursor.setoutputsize(0)
|
|
|
|
def test_ddl(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.execute('DROP TABLE IF EXISTS test')
|
|
self.assertEqual(cursor.fetchall(), [])
|
|
self.assertEqual(cursor.rowcount, -1)
|
|
|
|
def test_remember_last_host(self):
|
|
self.n_calls = 0
|
|
getaddrinfo = socket.getaddrinfo
|
|
|
|
def side_getaddrinfo(host, *args, **kwargs):
|
|
if host == 'wrong_host':
|
|
self.n_calls += 1
|
|
raise socket.error(-2, 'Name or service not known')
|
|
return getaddrinfo(host, *args, **kwargs)
|
|
|
|
with patch('socket.getaddrinfo') as mocked_getaddrinfo:
|
|
mocked_getaddrinfo.side_effect = side_getaddrinfo
|
|
|
|
conn_kwargs = {
|
|
'host': 'wrong_host',
|
|
'port': 1234,
|
|
'alt_hosts': '{}:{}'.format(self.host, self.port)
|
|
}
|
|
with self.created_connection(**conn_kwargs) as connection:
|
|
cursor = connection.cursor()
|
|
cursor.execute('SELECT 1')
|
|
self.assertEqual(cursor.fetchall(), [(1, )])
|
|
cursor.close()
|
|
|
|
cursor = connection.cursor()
|
|
cursor.execute('SELECT 1')
|
|
self.assertEqual(cursor.fetchall(), [(1, )])
|
|
cursor.close()
|
|
|
|
# Last host must be remembered and getaddrinfo must call exactly
|
|
# once with host == 'wrong_host'.
|
|
self.assertEqual(self.n_calls, 1)
|
|
|
|
def test_remove_cursor_from_connection_on_closing(self):
|
|
with self.created_connection() as connection:
|
|
self.assertEqual(len(connection.cursors), 0)
|
|
cur = connection.cursor()
|
|
self.assertEqual(len(connection.cursors), 1)
|
|
cur.close()
|
|
self.assertEqual(len(connection.cursors), 0)
|
|
|
|
|
|
class StreamingTestCase(DBAPITestCaseBase):
|
|
def test_fetchone(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.set_stream_results(True, 2)
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, types.GeneratorType)
|
|
|
|
self.assertEqual(cursor.fetchone(), (0, ))
|
|
self.assertEqual(cursor.fetchone(), (1, ))
|
|
self.assertEqual(cursor.fetchone(), (2, ))
|
|
self.assertEqual(cursor.fetchone(), (3, ))
|
|
self.assertEqual(cursor.fetchone(), None)
|
|
|
|
def test_fetchmany(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.set_stream_results(True, 2)
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, types.GeneratorType)
|
|
self.assertEqual(cursor.fetchmany(), [(0, )])
|
|
self.assertEqual(cursor.fetchmany(None), [(1, )])
|
|
self.assertEqual(cursor.fetchmany(0), [])
|
|
self.assertEqual(cursor.fetchmany(-1), [(2, ), (3, )])
|
|
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.fetchmany(1), [(0, )])
|
|
self.assertEqual(cursor.fetchmany(2), [(1, ), (2, )])
|
|
self.assertEqual(cursor.fetchmany(3), [(3, )])
|
|
self.assertEqual(cursor.fetchmany(3), [])
|
|
|
|
def test_fetchall(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.set_stream_results(True, 2)
|
|
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
# Check rowcount before and after fetch.
|
|
self.assertEqual(cursor.rowcount, -1)
|
|
self.assertEqual(cursor.fetchall(), [(0, ), (1, ), (2, ), (3, )])
|
|
self.assertEqual(cursor.rowcount, -1)
|
|
|
|
|
|
class ReprTestCase(DBAPITestCaseBase):
|
|
def test_cursor_repr(self):
|
|
with self.created_cursor() as cursor:
|
|
self.assertIn('closed: False', str(cursor))
|
|
|
|
def test_connection_repr(self):
|
|
with self.created_connection() as conn:
|
|
self.assertIn('closed: False', str(conn))
|
|
|
|
|
|
class ExtrasTestCase(DBAPITestCaseBase):
|
|
def test_columns_with_types_select(self):
|
|
with self.created_cursor() as cursor:
|
|
self.assertIsNone(cursor.columns_with_types)
|
|
cursor.execute(
|
|
'SELECT CAST(number AS UInt64) AS x '
|
|
'FROM system.numbers LIMIT 4'
|
|
)
|
|
cursor.fetchall()
|
|
self.assertEqual(cursor.columns_with_types, [('x', 'UInt64')])
|
|
|
|
def test_columns_with_types_insert(self):
|
|
with self.created_cursor() as cursor, self.create_table('a UInt8'):
|
|
cursor.executemany('INSERT INTO test (a) VALUES', [(123, )])
|
|
self.assertIsNone(cursor.columns_with_types)
|
|
|
|
def test_columns_with_types_streaming(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.set_stream_results(True, 2)
|
|
cursor.execute(
|
|
'SELECT CAST(number AS UInt64) AS x '
|
|
'FROM system.numbers LIMIT 4'
|
|
)
|
|
self.assertEqual(cursor.columns_with_types, [('x', 'UInt64')])
|
|
list(cursor)
|
|
self.assertEqual(cursor.columns_with_types, [('x', 'UInt64')])
|
|
|
|
def test_set_external_tables(self):
|
|
with self.created_cursor() as cursor:
|
|
data = [(0, ), (1, ), (2, )]
|
|
cursor.set_external_table('table1', [('x', 'UInt32')], data)
|
|
cursor.execute('SELECT * FROM table1')
|
|
self.assertEqual(cursor.fetchall(), data)
|
|
|
|
def test_settings(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.set_settings({'max_threads': 42})
|
|
|
|
cursor.execute(
|
|
"SELECT name, value, changed FROM system.settings "
|
|
"WHERE name = 'max_threads'",
|
|
)
|
|
self.assertEqual(cursor.fetchall(), [('max_threads', '42', 1)])
|
|
|
|
def test_set_query_id(self):
|
|
with self.created_cursor() as cursor:
|
|
query_id = 'my_query_id'
|
|
cursor.set_query_id(query_id)
|
|
cursor.execute(
|
|
'SELECT query_id '
|
|
'FROM system.processes '
|
|
'WHERE query_id = %(query_id)s',
|
|
{'query_id': query_id}
|
|
)
|
|
self.assertEqual(cursor.fetchall(), [(query_id, )])
|
|
|
|
def test_types_check(self):
|
|
with self.created_cursor() as cursor, self.create_table('a UInt8'):
|
|
cursor.set_types_check(True)
|
|
|
|
data = [(300, )]
|
|
cursor.executemany('INSERT INTO test (a) VALUES', data)
|
|
cursor.execute('SELECT * FROM test')
|
|
self.assertEqual(cursor.fetchall(), [(44, )])
|
|
|
|
def test_cursor_iteration(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(list(cursor), [(0,), (1,), (2,), (3,)])
|
|
|
|
def test_context_managers(self):
|
|
with self.create_connection() as connection:
|
|
with connection.cursor() as cursor:
|
|
cursor.execute('SELECT 1')
|
|
self.assertEqual(cursor.fetchall(), [(1, )])
|
|
|
|
|
|
class InterfaceTestCase(DBAPITestCaseBase):
|
|
def test_execute_after_close(self):
|
|
with self.created_cursor() as cursor:
|
|
cursor.close()
|
|
with self.assertRaises(InterfaceError) as e:
|
|
cursor.execute('SELECT 1')
|
|
self.assertEqual(str(e.exception), 'cursor already closed')
|
|
|
|
def test_create_cursor_on_closed_connection(self):
|
|
connection = self.create_connection()
|
|
connection.close()
|
|
|
|
with self.assertRaises(InterfaceError) as e:
|
|
connection.cursor()
|
|
self.assertEqual(str(e.exception), 'connection already closed')
|
|
|
|
def test_execute_fetch_before_query(self):
|
|
with self.created_cursor() as cursor:
|
|
with self.assertRaises(ProgrammingError) as e:
|
|
cursor.fetchall()
|
|
self.assertEqual(str(e.exception), 'no results to fetch')
|
|
|
|
|
|
class DictCursorFactoryTestCase(DBAPITestCaseBase):
|
|
def test_execute_fetchone(self):
|
|
cursor_kwargs = {'cursor_factory': DictCursor}
|
|
|
|
with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, list)
|
|
self.assertEqual(cursor.fetchone(), {'number': 0})
|
|
self.assertEqual(cursor.fetchone(), {'number': 1})
|
|
self.assertEqual(cursor.fetchone(), {'number': 2})
|
|
self.assertEqual(cursor.fetchone(), {'number': 3})
|
|
self.assertEqual(cursor.fetchone(), None)
|
|
|
|
def test_execute_fetchmany(self):
|
|
cursor_kwargs = {'cursor_factory': DictCursor}
|
|
|
|
with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, list)
|
|
self.assertEqual(cursor.fetchmany(), [{'number': 0}])
|
|
self.assertEqual(cursor.fetchmany(None), [{'number': 1}])
|
|
self.assertEqual(cursor.fetchmany(0), [])
|
|
self.assertEqual(
|
|
cursor.fetchmany(-1), [{'number': 2}, {'number': 3}]
|
|
)
|
|
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.fetchmany(1), [{'number': 0}])
|
|
self.assertEqual(
|
|
cursor.fetchmany(2), [{'number': 1}, {'number': 2}]
|
|
)
|
|
self.assertEqual(cursor.fetchmany(3), [{'number': 3}])
|
|
self.assertEqual(cursor.fetchmany(3), [])
|
|
|
|
cursor.arraysize = 2
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(
|
|
cursor.fetchmany(), [{'number': 0}, {'number': 1}])
|
|
self.assertEqual(
|
|
cursor.fetchmany(), [{'number': 2}, {'number': 3}]
|
|
)
|
|
|
|
def test_execute_fetchall(self):
|
|
cursor_kwargs = {'cursor_factory': DictCursor}
|
|
|
|
with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.rowcount, 4)
|
|
self.assertEqual(
|
|
cursor.fetchall(), [
|
|
{'number': 0}, {'number': 1}, {'number': 2}, {'number': 3}
|
|
])
|
|
|
|
|
|
class NamedTupleCursorFactoryTestCase(DBAPITestCaseBase):
|
|
def test_execute_fetchone(self):
|
|
cursor_kwargs = {'cursor_factory': NamedTupleCursor}
|
|
|
|
with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, list)
|
|
nt = namedtuple('Record', cursor._columns)
|
|
|
|
self.assertEqual(cursor.fetchone(), nt(0))
|
|
self.assertEqual(cursor.fetchone(), nt(1))
|
|
self.assertEqual(cursor.fetchone(), nt(2))
|
|
self.assertEqual(cursor.fetchone(), nt(3))
|
|
self.assertEqual(cursor.fetchone(), None)
|
|
|
|
def test_execute_fetchmany(self):
|
|
cursor_kwargs = {'cursor_factory': NamedTupleCursor}
|
|
|
|
with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, list)
|
|
nt = namedtuple('Record', cursor._columns)
|
|
|
|
self.assertEqual(cursor.fetchmany(), [nt(0)])
|
|
self.assertEqual(cursor.fetchmany(None), [nt(1)])
|
|
self.assertEqual(cursor.fetchmany(0), [])
|
|
self.assertEqual(cursor.fetchmany(-1), [nt(2), nt(3)])
|
|
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.fetchmany(1), [nt(0)])
|
|
self.assertEqual(cursor.fetchmany(2), [nt(1), nt(2)])
|
|
self.assertEqual(cursor.fetchmany(3), [nt(3)])
|
|
self.assertEqual(cursor.fetchmany(3), [])
|
|
|
|
cursor.arraysize = 2
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.fetchmany(), [nt(0), nt(1)])
|
|
self.assertEqual(cursor.fetchmany(), [nt(2), nt(3)])
|
|
|
|
def test_execute_fetchall(self):
|
|
cursor_kwargs = {'cursor_factory': NamedTupleCursor}
|
|
|
|
with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
self.assertEqual(cursor.rowcount, 4)
|
|
nt = namedtuple('Record', cursor._columns)
|
|
|
|
self.assertEqual(cursor.fetchall(), [nt(0), nt(1), nt(2), nt(3)])
|
|
|
|
def test_make_nt_caching(self):
|
|
cursor_kwargs = {'cursor_factory': NamedTupleCursor}
|
|
|
|
with self.created_cursor(cursor_kwargs=cursor_kwargs) as cursor:
|
|
cursor.execute('SELECT number FROM system.numbers LIMIT 4')
|
|
|
|
self.assertIsInstance(cursor._rows, list)
|
|
nt = namedtuple('Record', cursor._columns)
|
|
|
|
self.assertEqual(cursor.fetchone(), nt(0))
|
|
|
|
with patch('datastore_driver.dbapi.extras.namedtuple') as nt_mock:
|
|
nt_mock.side_effect = ValueError
|
|
self.assertEqual(cursor.fetchone(), nt(1))
|