driver/tests/test_dbapi.py

import types
from collections import namedtuple
from contextlib import contextmanager
import socket
from unittest.mock import patch

from clickhouse_driver import connect
from clickhouse_driver.dbapi import (
    ProgrammingError, InterfaceError, OperationalError
)
from clickhouse_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(
            'clickhouse://{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('clickhouse_driver.dbapi.extras.namedtuple') as nt_mock:
                nt_mock.side_effect = ValueError
                self.assertEqual(cursor.fetchone(), nt(1))