# Copyright 2017 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  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 unittest

import mock


class TestCursor(unittest.TestCase):

    @staticmethod
    def _get_target_class():
        from google.cloud.bigquery.dbapi import Cursor
        return Cursor

    def _make_one(self, *args, **kw):
        return self._get_target_class()(*args, **kw)

    def _mock_client(
            self, rows=None, schema=None, num_dml_affected_rows=None):
        from google.cloud.bigquery import client

        if rows is None:
            total_rows = 0
        else:
            total_rows = len(rows)

        mock_client = mock.create_autospec(client.Client)
        mock_client.query.return_value = self._mock_job(
            total_rows=total_rows,
            schema=schema,
            num_dml_affected_rows=num_dml_affected_rows)
        mock_client.list_rows.return_value = rows
        return mock_client

    def _mock_job(
            self, total_rows=0, schema=None, num_dml_affected_rows=None):
        from google.cloud.bigquery import job
        mock_job = mock.create_autospec(job.QueryJob)
        mock_job.error_result = None
        mock_job.state = 'DONE'
        mock_job.result.return_value = mock_job

        if num_dml_affected_rows is None:
            mock_job.statement_type = None  # API sends back None for SELECT
        else:
            mock_job.statement_type = 'UPDATE'

        mock_job.query_results.return_value = self._mock_results(
            total_rows=total_rows, schema=schema,
            num_dml_affected_rows=num_dml_affected_rows)
        return mock_job

    def _mock_results(
            self, total_rows=0, schema=None, num_dml_affected_rows=None):
        from google.cloud.bigquery import query
        mock_results = mock.create_autospec(query.QueryResults)
        mock_results.schema = schema
        mock_results.num_dml_affected_rows = num_dml_affected_rows
        mock_results.total_rows = total_rows
        return mock_results

    def test_ctor(self):
        from google.cloud.bigquery.dbapi import connect
        from google.cloud.bigquery.dbapi import Cursor
        connection = connect(self._mock_client())
        cursor = self._make_one(connection)
        self.assertIsInstance(cursor, Cursor)
        self.assertIs(cursor.connection, connection)

    def test_close(self):
        from google.cloud.bigquery.dbapi import connect
        connection = connect(self._mock_client())
        cursor = connection.cursor()
        # close() is a no-op, there is nothing to test.
        cursor.close()

    def test_fetchone_wo_execute_raises_error(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(self._mock_client())
        cursor = connection.cursor()
        self.assertRaises(dbapi.Error, cursor.fetchone)

    def test_fetchone_w_row(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(
            self._mock_client(rows=[(1,)]))
        cursor = connection.cursor()
        cursor.execute('SELECT 1;')
        row = cursor.fetchone()
        self.assertEqual(row, (1,))
        self.assertIsNone(cursor.fetchone())

    def test_fetchmany_wo_execute_raises_error(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(self._mock_client())
        cursor = connection.cursor()
        self.assertRaises(dbapi.Error, cursor.fetchmany)

    def test_fetchmany_w_row(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(
            self._mock_client(rows=[(1,)]))
        cursor = connection.cursor()
        cursor.execute('SELECT 1;')
        rows = cursor.fetchmany()
        self.assertEqual(len(rows), 1)
        self.assertEqual(rows[0], (1,))

    def test_fetchmany_w_size(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(
            self._mock_client(
                rows=[
                    (1, 2, 3),
                    (4, 5, 6),
                    (7, 8, 9),
                ]))
        cursor = connection.cursor()
        cursor.execute('SELECT a, b, c;')
        rows = cursor.fetchmany(size=2)
        self.assertEqual(len(rows), 2)
        self.assertEqual(rows[0], (1, 2, 3))
        self.assertEqual(rows[1], (4, 5, 6))
        second_page = cursor.fetchmany(size=2)
        self.assertEqual(len(second_page), 1)
        self.assertEqual(second_page[0], (7, 8, 9))
        third_page = cursor.fetchmany(size=2)
        self.assertEqual(third_page, [])

    def test_fetchmany_w_arraysize(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(
            self._mock_client(
                rows=[
                    (1, 2, 3),
                    (4, 5, 6),
                    (7, 8, 9),
                ]))
        cursor = connection.cursor()
        cursor.execute('SELECT a, b, c;')
        cursor.arraysize = 2
        rows = cursor.fetchmany()
        self.assertEqual(len(rows), 2)
        self.assertEqual(rows[0], (1, 2, 3))
        self.assertEqual(rows[1], (4, 5, 6))
        second_page = cursor.fetchmany()
        self.assertEqual(len(second_page), 1)
        self.assertEqual(second_page[0], (7, 8, 9))
        third_page = cursor.fetchmany()
        self.assertEqual(third_page, [])

    def test_fetchall_wo_execute_raises_error(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(self._mock_client())
        cursor = connection.cursor()
        self.assertRaises(dbapi.Error, cursor.fetchall)

    def test_fetchall_w_row(self):
        from google.cloud.bigquery import dbapi
        connection = dbapi.connect(
            self._mock_client(rows=[(1,)]))
        cursor = connection.cursor()
        cursor.execute('SELECT 1;')
        self.assertIsNone(cursor.description)
        self.assertEqual(cursor.rowcount, 1)
        rows = cursor.fetchall()
        self.assertEqual(len(rows), 1)
        self.assertEqual(rows[0], (1,))

    def test_execute_custom_job_id(self):
        from google.cloud.bigquery.dbapi import connect
        client = self._mock_client(rows=[], num_dml_affected_rows=0)
        connection = connect(client)
        cursor = connection.cursor()
        cursor.execute('SELECT 1;', job_id='foo')
        args, kwargs = client.query.call_args
        self.assertEqual(args[0], 'SELECT 1;')
        self.assertEqual(kwargs['job_id'], 'foo')

    def test_execute_w_dml(self):
        from google.cloud.bigquery.dbapi import connect
        connection = connect(
            self._mock_client(rows=[], num_dml_affected_rows=12))
        cursor = connection.cursor()
        cursor.execute('DELETE FROM UserSessions WHERE user_id = \'test\';')
        rows = cursor.fetchall()
        self.assertIsNone(cursor.description)
        self.assertEqual(cursor.rowcount, 12)
        self.assertEqual(rows, [])

    def test_execute_w_query(self):
        from google.cloud.bigquery.schema import SchemaField
        from google.cloud.bigquery import dbapi

        connection = dbapi.connect(self._mock_client(
            rows=[('hello', 'world', 1), ('howdy', 'y\'all', 2)],
            schema=[
                SchemaField('a', 'STRING', mode='NULLABLE'),
                SchemaField('b', 'STRING', mode='REQUIRED'),
                SchemaField('c', 'INTEGER', mode='NULLABLE')]))
        cursor = connection.cursor()
        cursor.execute('SELECT a, b, c FROM hello_world WHERE d > 3;')

        # Verify the description.
        self.assertEqual(len(cursor.description), 3)
        a_name, a_type, _, _, _, _, a_null_ok = cursor.description[0]
        self.assertEqual(a_name, 'a')
        self.assertEqual(a_type, 'STRING')
        self.assertEqual(a_type, dbapi.STRING)
        self.assertTrue(a_null_ok)
        b_name, b_type, _, _, _, _, b_null_ok = cursor.description[1]
        self.assertEqual(b_name, 'b')
        self.assertEqual(b_type, 'STRING')
        self.assertEqual(b_type, dbapi.STRING)
        self.assertFalse(b_null_ok)
        c_name, c_type, _, _, _, _, c_null_ok = cursor.description[2]
        self.assertEqual(c_name, 'c')
        self.assertEqual(c_type, 'INTEGER')
        self.assertEqual(c_type, dbapi.NUMBER)
        self.assertTrue(c_null_ok)

        # Verify the results.
        self.assertEqual(cursor.rowcount, 2)
        row = cursor.fetchone()
        self.assertEqual(row, ('hello', 'world', 1))
        row = cursor.fetchone()
        self.assertEqual(row, ('howdy', 'y\'all', 2))
        row = cursor.fetchone()
        self.assertIsNone(row)

    def test_execute_raises_if_result_raises(self):
        import google.cloud.exceptions

        from google.cloud.bigquery import client
        from google.cloud.bigquery import job
        from google.cloud.bigquery.dbapi import connect
        from google.cloud.bigquery.dbapi import exceptions

        job = mock.create_autospec(job.QueryJob)
        job.result.side_effect = google.cloud.exceptions.GoogleCloudError('')
        client = mock.create_autospec(client.Client)
        client.query.return_value = job
        connection = connect(client)
        cursor = connection.cursor()

        with self.assertRaises(exceptions.DatabaseError):
            cursor.execute('SELECT 1')

    def test_executemany_w_dml(self):
        from google.cloud.bigquery.dbapi import connect
        connection = connect(
            self._mock_client(rows=[], num_dml_affected_rows=12))
        cursor = connection.cursor()
        cursor.executemany(
            'DELETE FROM UserSessions WHERE user_id = %s;',
            (('test',), ('anothertest',)))
        self.assertIsNone(cursor.description)
        self.assertEqual(cursor.rowcount, 12)

    def test__format_operation_w_dict(self):
        from google.cloud.bigquery.dbapi import cursor
        formatted_operation = cursor._format_operation(
            'SELECT %(somevalue)s, %(a `weird` one)s;',
            {
                'somevalue': 'hi',
                'a `weird` one': 'world',
            })
        self.assertEqual(
            formatted_operation, 'SELECT @`somevalue`, @`a \\`weird\\` one`;')

    def test__format_operation_w_wrong_dict(self):
        from google.cloud.bigquery import dbapi
        from google.cloud.bigquery.dbapi import cursor
        self.assertRaises(
            dbapi.ProgrammingError,
            cursor._format_operation,
            'SELECT %(somevalue)s, %(othervalue)s;',
            {
                'somevalue-not-here': 'hi',
                'othervalue': 'world',
            })

    def test__format_operation_w_sequence(self):
        from google.cloud.bigquery.dbapi import cursor
        formatted_operation = cursor._format_operation(
            'SELECT %s, %s;', ('hello', 'world'))
        self.assertEqual(formatted_operation, 'SELECT ?, ?;')

    def test__format_operation_w_too_short_sequence(self):
        from google.cloud.bigquery import dbapi
        from google.cloud.bigquery.dbapi import cursor
        self.assertRaises(
            dbapi.ProgrammingError,
            cursor._format_operation,
            'SELECT %s, %s;',
            ('hello',))
