singlestoredb 0.4.0__py3-none-any.whl → 1.0.4__py3-none-any.whl

singlestoredb/{clients/pymysqlsv → mysql}/cursors.py

@@ -1,7 +1,10 @@
 # type: ignore
 import re
+from collections import namedtuple
 
 from . import err
+from ..connection import Cursor as BaseCursor
+from ..utils.debug import log_query
 
 
 #: Regular expression for :meth:`Cursor.executemany`.
@@ -15,15 +18,21 @@ RE_INSERT_VALUES = re.compile(
 )
 
 
-class Cursor:
+class Cursor(BaseCursor):
     """
     This is the object used to interact with the database.
 
     Do not create an instance of a Cursor yourself. Call
-    connections.Connection.cursor().
+    connection.Connection.cursor().
 
     See `Cursor <https://www.python.org/dev/peps/pep-0249/#cursor-objects>`_ in
     the specification.
+
+    Parameters
+    ----------
+    connection : Connection
+        The connection the cursor is associated with.
+
     """
 
     #: Max statement size which :meth:`executemany` generates.
@@ -33,27 +42,54 @@ class Cursor:
     max_stmt_length = 1024000
 
     def __init__(self, connection):
-        self.connection = connection
-        self.description = None
-        self.rownumber = 0
+        self._connection = connection
+        self.warning_count = 0
+        self._description = None
+        self._rownumber = 0
         self.rowcount = -1
         self.arraysize = 1
         self._executed = None
         self._result = None
         self._rows = None
+        self.lastrowid = None
+
+    @property
+    def messages(self):
+        # TODO
+        return []
+
+    @property
+    def description(self):
+        return self._description
+
+    @property
+    def connection(self):
+        return self._connection
+
+    @property
+    def rownumber(self):
+        return self._rownumber
 
     def close(self):
-        """
-        Closing a cursor just exhausts all remaining data.
-        """
-        conn = self.connection
+        """Closing a cursor just exhausts all remaining data."""
+        conn = self._connection
         if conn is None:
             return
         try:
             while self.nextset():
                 pass
         finally:
-            self.connection = None
+            self._connection = None
+
+    @property
+    def open(self) -> bool:
+        conn = self._connection
+        if conn is None:
+            return False
+        return True
+
+    def is_connected(self):
+        return self.open
 
     def __enter__(self):
         return self
@@ -63,9 +99,9 @@ class Cursor:
         self.close()
 
     def _get_db(self):
-        if not self.connection:
+        if not self._connection:
             raise err.ProgrammingError('Cursor closed')
-        return self.connection
+        return self._connection
 
     def _check_executed(self):
         if not self._executed:
@@ -99,87 +135,94 @@ class Cursor:
     def nextset(self):
         return self._nextset(False)
 
-    def _ensure_bytes(self, x, encoding=None):
-        if isinstance(x, str):
-            x = x.encode(encoding)
-        elif isinstance(x, (tuple, list)):
-            x = type(x)(self._ensure_bytes(v, encoding=encoding) for v in x)
-        return x
-
     def _escape_args(self, args, conn):
-        if isinstance(args, (tuple, list)):
-            return tuple(conn.literal(arg) for arg in args)
-        elif isinstance(args, dict):
-            return {key: conn.literal(val) for (key, val) in args.items()}
-        else:
-            # If it's not a dictionary let's try escaping it anyways.
-            # Worst case it will throw a Value error
-            return conn.escape(args)
+        dtype = type(args)
+        literal = conn.literal
+        if dtype is tuple or dtype is list or isinstance(args, (tuple, list)):
+            return tuple(literal(arg) for arg in args)
+        elif dtype is dict or isinstance(args, dict):
+            return {key: literal(val) for (key, val) in args.items()}
+        # If it's not a dictionary let's try escaping it anyways.
+        # Worst case it will throw a Value error
+        return conn.escape(args)
 
     def mogrify(self, query, args=None):
         """
-        Returns the exact string that would be sent to the database by calling the
-        execute() method.
+        Returns the exact string sent to the database by calling the execute() method.
 
-        :param query: Query to mogrify.
-        :type query: str
+        This method follows the extension to the DB API 2.0 followed by Psycopg.
 
-        :param args: Parameters used with query. (optional)
-        :type args: tuple, list or dict
+        Parameters
+        ----------
+        query : str
+            Query to mogrify.
+        args : Sequence[Any] or Dict[str, Any] or Any, optional
+            Parameters used with query. (optional)
 
-        :return: The query with argument binding applied.
-        :rtype: str
+        Returns
+        -------
+        str : The query with argument binding applied.
 
-        This method follows the extension to the DB API 2.0 followed by Psycopg.
         """
         conn = self._get_db()
 
-        if args is not None:
+        if args:
             query = query % self._escape_args(args, conn)
 
         return query
 
     def execute(self, query, args=None):
-        """Execute a query.
+        """
+        Execute a query.
 
-        :param query: Query to execute.
-        :type query: str
+        If args is a list or tuple, :1, :2, etc. can be used as a
+        placeholder in the query.  If args is a dict, :name can be used
+        as a placeholder in the query.
 
-        :param args: Parameters used with query. (optional)
-        :type args: tuple, list or dict
+        Parameters
+        ----------
+        query : str
+            Query to execute.
+        args : Sequence[Any] or Dict[str, Any] or Any, optional
+            Parameters used with query. (optional)
 
-        :return: Number of affected rows.
-        :rtype: int
+        Returns
+        -------
+        int : Number of affected rows.
 
-        If args is a list or tuple, %s can be used as a placeholder in the query.
-        If args is a dict, %(name)s can be used as a placeholder in the query.
         """
         while self.nextset():
             pass
 
+        log_query(query, args)
+
         query = self.mogrify(query, args)
 
         result = self._query(query)
         self._executed = query
         return result
 
-    def executemany(self, query, args):
-        """Run several data against one query.
-
-        :param query: Query to execute.
-        :type query: str
-
-        :param args: Sequence of sequences or mappings. It is used as parameter.
-        :type args: tuple or list
-
-        :return: Number of rows affected, if any.
-        :rtype: int or None
+    def executemany(self, query, args=None):
+        """
+        Run several data against one query.
 
         This method improves performance on multiple-row INSERT and
         REPLACE. Otherwise it is equivalent to looping over args with
         execute().
+
+        Parameters
+        ----------
+        query : str,
+            Query to execute.
+        args : Sequnce[Any], optional
+            Sequence of sequences or mappings. It is used as parameter.
+
+        Returns
+        -------
+        int : Number of rows affected, if any.
+
         """
-        if not args:
+        if args is None or len(args) == 0:
             return
 
         m = RE_INSERT_VALUES.match(query)
@@ -210,7 +253,11 @@ class Cursor:
         if isinstance(postfix, str):
             postfix = postfix.encode(encoding)
         sql = bytearray(prefix)
-        args = iter(args)
+        # Detect dataframes
+        if hasattr(args, 'itertuples'):
+            args = args.itertuples(index=False)
+        else:
+            args = iter(args)
         v = values % escape(next(args), conn)
         if isinstance(v, str):
             v = v.encode(encoding, 'surrogateescape')
@@ -218,7 +265,7 @@ class Cursor:
         rows = 0
         for arg in args:
             v = values % escape(arg, conn)
-            if isinstance(v, str):
+            if type(v) is str or isinstance(v, str):
                 v = v.encode(encoding, 'surrogateescape')
             if len(sql) + len(v) + len(postfix) + 1 > max_stmt_length:
                 rows += self.execute(sql + postfix)
@@ -231,15 +278,8 @@ class Cursor:
         return rows
 
     def callproc(self, procname, args=()):
-        """Execute stored procedure procname with args.
-
-        :param procname: Name of procedure to execute on server.
-        :type procname: str
-
-        :param args: Sequence of parameters to use with procedure.
-        :type args: tuple or list
-
-        Returns the original args.
+        """
+        Execute stored procedure procname with args.
 
         Compatibility warning: PEP-249 specifies that any modified
         parameters must be returned. This is currently impossible
@@ -259,6 +299,18 @@ class Cursor:
         behavior with respect to the DB-API. Be sure to use nextset()
         to advance through all result sets; otherwise you may get
         disconnected.
+
+        Parameters
+        ----------
+        procname : str
+            Name of procedure to execute on server.
+        args : Sequence[Any], optional
+            Sequence of parameters to use with procedure.
+
+        Returns
+        -------
+        Sequence[Any] : The original args.
+
         """
         conn = self._get_db()
         if args:
@@ -271,7 +323,7 @@ class Cursor:
             )
             self.nextset()
 
-        q = 'CALL %s(%s)' % (
+        q = 'CALL {}({})'.format(
             procname,
             ','.join(['@_%s_%d' % (procname, i) for i in range(len(args))]),
         )
@@ -286,20 +338,21 @@ class Cursor:
 
     def _unchecked_fetchone(self):
         """Fetch the next row."""
-        if self._rows is None or self.rownumber >= len(self._rows):
+        if self._rows is None or self._rownumber >= len(self._rows):
             return None
-        result = self._rows[self.rownumber]
-        self.rownumber += 1
+        result = self._rows[self._rownumber]
+        self._rownumber += 1
         return result
 
     def fetchmany(self, size=None):
         """Fetch several rows."""
         self._check_executed()
         if self._rows is None:
+            self.warning_count = self._result.warning_count
             return ()
-        end = self.rownumber + (size or self.arraysize)
-        result = self._rows[self.rownumber: end]
-        self.rownumber = min(end, len(self._rows))
+        end = self._rownumber + (size or self.arraysize)
+        result = self._rows[self._rownumber: end]
+        self._rownumber = min(end, len(self._rows))
         return result
 
     def fetchall(self):
@@ -307,17 +360,17 @@ class Cursor:
         self._check_executed()
         if self._rows is None:
             return ()
-        if self.rownumber:
-            result = self._rows[self.rownumber:]
+        if self._rownumber:
+            result = self._rows[self._rownumber:]
         else:
             result = self._rows
-        self.rownumber = len(self._rows)
+        self._rownumber = len(self._rows)
         return result
 
     def scroll(self, value, mode='relative'):
         self._check_executed()
         if mode == 'relative':
-            r = self.rownumber + value
+            r = self._rownumber + value
         elif mode == 'absolute':
             r = value
         else:
@@ -325,7 +378,7 @@ class Cursor:
 
         if not (0 <= r < len(self._rows)):
             raise IndexError('out of range')
-        self.rownumber = r
+        self._rownumber = r
 
     def _query(self, q):
         conn = self._get_db()
@@ -335,11 +388,12 @@ class Cursor:
         return self.rowcount
 
     def _clear_result(self):
-        self.rownumber = 0
+        self._rownumber = 0
         self._result = None
 
         self.rowcount = 0
-        self.description = None
+        self.warning_count = 0
+        self._description = None
         self.lastrowid = None
         self._rows = None
 
@@ -349,15 +403,19 @@ class Cursor:
         self._result = result = conn._result
 
         self.rowcount = result.affected_rows
-        self.description = result.description
+        self.warning_count = result.warning_count
+        # Affected rows is set to max int64 for compatibility with MySQLdb, but
+        # the DB-API requires this value to be -1. This happens in unbuffered mode.
+        if self.rowcount == 18446744073709551615:
+            self.rowcount = -1
+        self._description = result.description
         self.lastrowid = result.insert_id
         self._rows = result.rows
 
     def __iter__(self):
         self._check_executed()
-        _unchecked_fetchone = self._unchecked_fetchone
 
-        def fetchall_unbuffered_gen():
+        def fetchall_unbuffered_gen(_unchecked_fetchone=self._unchecked_fetchone):
             while True:
                 out = _unchecked_fetchone()
                 if out is not None:
@@ -378,6 +436,10 @@ class Cursor:
     NotSupportedError = err.NotSupportedError
 
 
+class CursorSV(Cursor):
+    """Cursor class for C extension."""
+
+
 class DictCursorMixin:
     # You can override this to use OrderedDict or other dict-like types.
     dict_type = dict
@@ -385,7 +447,7 @@ class DictCursorMixin:
     def _do_get_result(self):
         super(DictCursorMixin, self)._do_get_result()
         fields = []
-        if self.description:
+        if self._description:
             for f in self._result.fields:
                 name = f.name
                 if name in fields:
@@ -403,7 +465,42 @@ class DictCursorMixin:
 
 
 class DictCursor(DictCursorMixin, Cursor):
-    """A cursor which returns results as a dictionary"""
+    """A cursor which returns results as a dictionary."""
+
+
+class DictCursorSV(Cursor):
+    """A cursor which returns results as a dictionary for C extension."""
+
+
+class NamedtupleCursorMixin:
+
+    def _do_get_result(self):
+        super(NamedtupleCursorMixin, self)._do_get_result()
+        fields = []
+        if self._description:
+            for f in self._result.fields:
+                name = f.name
+                if name in fields:
+                    name = f.table_name + '.' + name
+                fields.append(name)
+            self._fields = fields
+            self._namedtuple = namedtuple('Row', self._fields, rename=True)
+
+        if fields and self._rows:
+            self._rows = [self._conv_row(r) for r in self._rows]
+
+    def _conv_row(self, row):
+        if row is None:
+            return None
+        return self._namedtuple(*row)
+
+
+class NamedtupleCursor(NamedtupleCursorMixin, Cursor):
+    """A cursor which returns results in a named tuple."""
+
+
+class NamedtupleCursorSV(Cursor):
+    """A cursor which returns results as a named tuple for C extension."""
 
 
 class SSCursor(Cursor):
@@ -420,13 +517,14 @@ class SSCursor(Cursor):
     returning the total number of rows, so the only way to tell how many rows
     there are is to iterate over every row returned. Also, it currently isn't
     possible to scroll backwards, as only the current row is held in memory.
+
     """
 
     def _conv_row(self, row):
         return row
 
     def close(self):
-        conn = self.connection
+        conn = self._connection
         if conn is None:
             return
 
@@ -437,7 +535,7 @@ class SSCursor(Cursor):
             while self.nextset():
                 pass
         finally:
-            self.connection = None
+            self._connection = None
 
     __del__ = close
 
@@ -464,28 +562,34 @@ class SSCursor(Cursor):
         """Fetch next row."""
         row = self.read_next()
         if row is None:
+            self.warning_count = self._result.warning_count
             return None
-        self.rownumber += 1
+        self._rownumber += 1
         return row
 
     def fetchall(self):
         """
-        Fetch all, as per MySQLdb. Pretty useless for large queries, as
-        it is buffered. See fetchall_unbuffered(), if you want an unbuffered
+        Fetch all, as per MySQLdb.
+
+        Pretty useless for large queries, as it is buffered.
+        See fetchall_unbuffered(), if you want an unbuffered
         generator version of this method.
+
         """
         return list(self.fetchall_unbuffered())
 
     def fetchall_unbuffered(self):
         """
-        Fetch all, implemented as a generator, which isn't to standard,
-        however, it doesn't make sense to return everything in a list, as that
-        would use ridiculous memory for large result sets.
+        Fetch all, implemented as a generator.
+
+        This is not a standard DB-API operation, however, it doesn't make
+        sense to return everything in a list, as that would use ridiculous
+        memory for large result sets.
+
         """
         self._check_executed()
-        _unchecked_fetchone = self._unchecked_fetchone
 
-        def fetchall_unbuffered_gen():
+        def fetchall_unbuffered_gen(_unchecked_fetchone=self._unchecked_fetchone):
             while True:
                 out = _unchecked_fetchone()
                 if out is not None:
@@ -507,9 +611,10 @@ class SSCursor(Cursor):
         for i in range(size):
             row = self.read_next()
             if row is None:
+                self.warning_count = self._result.warning_count
                 break
             rows.append(row)
-            self.rownumber += 1
+            self._rownumber += 1
         return rows
 
     def scroll(self, value, mode='relative'):
@@ -523,17 +628,17 @@ class SSCursor(Cursor):
 
             for _ in range(value):
                 self.read_next()
-            self.rownumber += value
+            self._rownumber += value
         elif mode == 'absolute':
-            if value < self.rownumber:
+            if value < self._rownumber:
                 raise err.NotSupportedError(
                     'Backwards scrolling not supported by this cursor',
                 )
 
-            end = value - self.rownumber
+            end = value - self._rownumber
             for _ in range(end):
                 self.read_next()
-            self.rownumber = value
+            self._rownumber = value
         else:
             raise err.ProgrammingError('unknown scroll mode %s' % mode)
 
@@ -546,7 +651,7 @@ class SSCursorSV(SSCursor):
         row = self._result._read_rowdata_packet_unbuffered(1)
         if row is None:
             return None
-        self.rownumber += 1
+        self._rownumber += 1
         return row
 
     def fetchone(self):
@@ -562,7 +667,10 @@ class SSCursorSV(SSCursor):
         out = self._result._read_rowdata_packet_unbuffered(size)
         if out is None:
             return []
-        self.rownumber += len(out)
+        if size == 1:
+            self._rownumber += 1
+            return [out]
+        self._rownumber += len(out)
         return out
 
     def scroll(self, value, mode='relative'):
@@ -575,23 +683,31 @@ class SSCursorSV(SSCursor):
                 )
 
             self._result._read_rowdata_packet_unbuffered(value)
-            self.rownumber += value
+            self._rownumber += value
         elif mode == 'absolute':
-            if value < self.rownumber:
+            if value < self._rownumber:
                 raise err.NotSupportedError(
                     'Backwards scrolling not supported by this cursor',
                 )
 
-            end = value - self.rownumber
+            end = value - self._rownumber
             self._result._read_rowdata_packet_unbuffered(end)
-            self.rownumber = value
+            self._rownumber = value
         else:
             raise err.ProgrammingError('unknown scroll mode %s' % mode)
 
 
+class SSDictCursor(DictCursorMixin, SSCursor):
+    """An unbuffered cursor, which returns results as a dictionary"""
+
+
 class SSDictCursorSV(SSCursorSV):
-    """An unbuffered cursor for use with PyMySQLsv, which returns results as a dictionary"""
+    """An unbuffered cursor for the C extension, which returns a dictionary"""
 
 
-class SSDictCursor(DictCursorMixin, SSCursor):
-    """An unbuffered cursor, which returns results as a dictionary"""
+class SSNamedtupleCursor(NamedtupleCursorMixin, SSCursor):
+    """An unbuffered cursor, which returns results as a named tuple"""
+
+
+class SSNamedtupleCursorSV(SSCursorSV):
+    """An unbuffered cursor for the C extension, which returns results as a named tuple"""