test_SSCursor.py

# type: ignore
import sys

try:
    import singlestoredb.mysql as sv
    from singlestoredb.mysql.tests import base
    import singlestoredb.mysql.cursors as cursors
    from singlestoredb.mysql.constants import CLIENT
except Exception:
    # For local testing from top-level directory, without installing
    sys.path.append('../../..')
    import singlestoredb.mysql as sv  # noqa: F401
    from singlestoredb.mysql.tests import base
    import singlestoredb.mysql.cursors as cursors
    from singlestoredb.mysql.constants import CLIENT


class TestSSCursor(base.PyMySQLTestCase):

    def test_SSCursor(self):
        # affected_rows = 18446744073709551615
        affected_rows = -1

        conn = self.connect(
            client_flag=CLIENT.MULTI_STATEMENTS,
            cursorclass=cursors.SSCursor,
        )
        data = [
            ('America', '', 'America/Jamaica'),
            ('America', '', 'America/Los_Angeles'),
            ('America', '', 'America/Lima'),
            ('America', '', 'America/New_York'),
            ('America', '', 'America/Menominee'),
            ('America', '', 'America/Havana'),
            ('America', '', 'America/El_Salvador'),
            ('America', '', 'America/Costa_Rica'),
            ('America', '', 'America/Denver'),
            ('America', '', 'America/Detroit'),
        ]

        cursor = conn.cursor()

        cursor.execute('DROP TABLE IF EXISTS tz_data')

        # Create table
        cursor.execute(
            'CREATE TABLE tz_data ('
            'region VARCHAR(64),'
            'zone VARCHAR(64),'
            'name VARCHAR(64))',
        )

        conn.begin()
        # Test INSERT
        for i in data:
            cursor.execute('INSERT INTO tz_data VALUES (%s, %s, %s)', i)
            self.assertEqual(conn.affected_rows(), 1, 'affected_rows does not match')
        conn.commit()

        # Test fetchone()
        iter = 0
        cursor.execute('SELECT * FROM tz_data')
        while True:
            row = cursor.fetchone()
            if row is None:
                break
            iter += 1

            # Test cursor.rowcount
            self.assertEqual(
                cursor.rowcount,
                affected_rows,
                'cursor.rowcount != %s' % (str(affected_rows)),
            )

            # Test cursor.rownumber
            self.assertEqual(
                cursor.rownumber, iter, 'cursor.rowcount != %s' % (str(iter)),
            )

            # Test row came out the same as it went in
            self.assertEqual((row in data), True, 'Row not found in source data')

        # Test fetchall
        cursor.execute('SELECT * FROM tz_data')
        self.assertEqual(
            len(cursor.fetchall()),
            len(data),
            'fetchall failed. Number of rows does not match',
        )

        # Test fetchmany
        cursor.execute('SELECT * FROM tz_data')
        self.assertEqual(
            len(cursor.fetchmany(2)),
            2,
            'fetchmany failed. Number of rows does not match',
        )

        # So MySQLdb won't throw "Commands out of sync"
        while True:
            res = cursor.fetchone()
            if res is None:
                break

        # Test update, affected_rows()
        cursor.execute('UPDATE tz_data SET zone = %s', ['Foo'])
        conn.commit()
        self.assertEqual(
            cursor.rowcount,
            len(data),
            'Update failed. affected_rows != %s' % (str(len(data))),
        )

        # Test executemany
        cursor.executemany('INSERT INTO tz_data VALUES (%s, %s, %s)', data)
        self.assertEqual(
            cursor.rowcount,
            len(data),
            'executemany failed. cursor.rowcount != %s' % (str(len(data))),
        )

        # Test multiple datasets
        cursor.execute('SELECT 1; SELECT 2; SELECT 3')
        self.assertListEqual(list(cursor), [(1,)])
        self.assertTrue(cursor.nextset())
        self.assertListEqual(list(cursor), [(2,)])
        self.assertTrue(cursor.nextset())
        self.assertListEqual(list(cursor), [(3,)])
        self.assertFalse(cursor.nextset())

        cursor.execute('DROP TABLE IF EXISTS tz_data')
        cursor.close()


__all__ = ['TestSSCursor']

if __name__ == '__main__':
    import unittest

    unittest.main()