chdb 3.3.0__cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.whl

chdb/dbapi/connections.py

@@ -0,0 +1,100 @@
+from . import err
+from .cursors import Cursor
+from . import converters
+from ..state import sqlitelike as chdb_stateful
+
+DEBUG = False
+VERBOSE = False
+
+
+class Connection(object):
+    """
+    Representation of a connection with chdb.
+    """
+
+    def __init__(self, path=None):
+        self._closed = False
+        self.encoding = "utf8"
+        self._affected_rows = 0
+        self._resp = None
+
+        # Initialize sqlitelike connection
+        connection_string = ":memory:" if path is None else f"file:{path}"
+        self._conn = chdb_stateful.Connection(connection_string)
+
+        # Test connection with a simple query
+        cursor = self._conn.cursor()
+        cursor.execute("SELECT 1")
+        cursor.close()
+
+    def close(self):
+        """Send the quit message and close the socket."""
+        if self._closed:
+            raise err.Error("Already closed")
+        self._closed = True
+        self._conn.close()
+
+    @property
+    def open(self):
+        """Return True if the connection is open"""
+        return not self._closed
+
+    def commit(self):
+        """Commit changes to stable storage."""
+        # No-op for ClickHouse
+        pass
+
+    def rollback(self):
+        """Roll back the current transaction."""
+        # No-op for ClickHouse
+        pass
+
+    def cursor(self, cursor=None):
+        """Create a new cursor to execute queries with."""
+        if self._closed:
+            raise err.Error("Connection closed")
+        if cursor:
+            return Cursor(self)
+        return Cursor(self)
+
+    def query(self, sql, fmt="CSV"):
+        """Execute a query and return the raw result."""
+        if self._closed:
+            raise err.InterfaceError("Connection closed")
+
+        if isinstance(sql, str):
+            sql = sql.encode(self.encoding, "surrogateescape")
+
+        try:
+            result = self._conn.query(sql.decode(), fmt)
+            self._resp = result
+            return result
+        except Exception as error:
+            raise err.InterfaceError(f"Query error: {error}")
+
+    def escape(self, obj, mapping=None):
+        """Escape whatever value you pass to it."""
+        return converters.escape_item(obj, mapping)
+
+    def escape_string(self, s):
+        return converters.escape_string(s)
+
+    def _quote_bytes(self, s):
+        return converters.escape_bytes(s)
+
+    def __enter__(self):
+        """Context manager that returns a Cursor"""
+        return self.cursor()
+
+    def __exit__(self, exc, value, traceback):
+        """On successful exit, commit. On exception, rollback"""
+        if exc:
+            self.rollback()
+        else:
+            self.commit()
+        self.close()
+
+    @property
+    def resp(self):
+        """Return the last query response"""
+        return self._resp

chdb/dbapi/constants/FIELD_TYPE.py

@@ -0,0 +1,31 @@
+DECIMAL = 0
+TINY = 1
+SHORT = 2
+LONG = 3
+FLOAT = 4
+DOUBLE = 5
+NULL = 6
+TIMESTAMP = 7
+LONGLONG = 8
+INT24 = 9
+DATE = 10
+TIME = 11
+DATETIME = 12
+YEAR = 13
+NEWDATE = 14
+VARCHAR = 15
+BIT = 16
+JSON = 245
+NEWDECIMAL = 246
+ENUM = 247
+SET = 248
+TINY_BLOB = 249
+MEDIUM_BLOB = 250
+LONG_BLOB = 251
+BLOB = 252
+VAR_STRING = 253
+STRING = 254
+GEOMETRY = 255
+
+CHAR = TINY
+INTERVAL = ENUM

chdb/dbapi/converters.py

@@ -0,0 +1,293 @@
+import datetime
+from decimal import Decimal
+from .err import DataError
+import re
+import time
+
+
+def escape_item(val, mapping=None):
+    if mapping is None:
+        mapping = encoders
+    encoder = mapping.get(type(val))
+
+    # Fallback to default when no encoder found
+    if not encoder:
+        try:
+            encoder = mapping[str]
+        except KeyError:
+            raise TypeError("no default type converter defined")
+
+    val = encoder(val, mapping)
+    return val
+
+
+def escape_dict(val, mapping=None):
+    n = {}
+    for k, v in val.items():
+        quoted = escape_item(v, mapping)
+        n[k] = quoted
+    return n
+
+
+def escape_sequence(val, mapping=None):
+    n = []
+    for item in val:
+        quoted = escape_item(item, mapping)
+        n.append(quoted)
+    return "(" + ",".join(n) + ")"
+
+
+def escape_set(val, mapping=None):
+    return ','.join([escape_item(x, mapping) for x in val])
+
+
+def escape_bool(value, mapping=None):
+    return str(int(value))
+
+
+def escape_object(value, mapping=None):
+    return str(value)
+
+
+def escape_int(value, mapping=None):
+    return str(value)
+
+
+def escape_float(value, mapping=None):
+    return '%.15g' % value
+
+
+_escape_table = [chr(x) for x in range(128)]
+_escape_table[ord("'")] = u"''"
+_escape_table[ord("\\")] = "\\\\"
+
+
+def _escape_unicode(value, mapping=None):
+    """escapes *value* with adding single quote.
+
+    Value should be unicode
+    """
+    return value.translate(_escape_table)
+
+
+escape_string = _escape_unicode
+
+# On Python ~3.5, str.decode('ascii', 'surrogateescape') is slow.
+# (fixed in Python 3.6, http://bugs.python.org/issue24870)
+# Workaround is str.decode('latin1') then translate 0x80-0xff into 0udc80-0udcff.
+# We can escape special chars and surrogateescape at once.
+_escape_bytes_table = _escape_table + [chr(i) for i in range(0xdc80, 0xdd00)]
+
+
+def escape_bytes(value, mapping=None):
+    return "'%s'" % value.decode('latin1').translate(_escape_bytes_table)
+
+
+def escape_unicode(value, mapping=None):
+    return u"'%s'" % _escape_unicode(value)
+
+
+def escape_str(value, mapping=None):
+    return "'%s'" % escape_string(str(value), mapping)
+
+
+def escape_None(value, mapping=None):
+    return 'NULL'
+
+
+def escape_timedelta(obj, mapping=None):
+    seconds = int(obj.seconds) % 60
+    minutes = int(obj.seconds // 60) % 60
+    hours = int(obj.seconds // 3600) % 24 + int(obj.days) * 24
+    if obj.microseconds:
+        fmt = "'{0:02d}:{1:02d}:{2:02d}.{3:06d}'"
+    else:
+        fmt = "'{0:02d}:{1:02d}:{2:02d}'"
+    return fmt.format(hours, minutes, seconds, obj.microseconds)
+
+
+def escape_time(obj, mapping=None):
+    return "'{}'".format(obj.isoformat(timespec='microseconds'))
+
+
+def escape_datetime(obj, mapping=None):
+    return "'{}'".format(obj.isoformat(sep=' ', timespec='microseconds'))
+    # if obj.microsecond:
+    #    fmt = "'{0.year:04}-{0.month:02}-{0.day:02} {0.hour:02}:{0.minute:02}:{0.second:02}.{0.microsecond:06}'"
+    # else:
+    #    fmt = "'{0.year:04}-{0.month:02}-{0.day:02} {0.hour:02}:{0.minute:02}:{0.second:02}'"
+    # return fmt.format(obj)
+
+
+def escape_date(obj, mapping=None):
+    return "'{}'".format(obj.isoformat())
+
+
+def escape_struct_time(obj, mapping=None):
+    return escape_datetime(datetime.datetime(*obj[:6]))
+
+
+def _convert_second_fraction(s):
+    if not s:
+        return 0
+    # Pad zeros to ensure the fraction length in microseconds
+    s = s.ljust(6, '0')
+    return int(s[:6])
+
+
+def convert_datetime(obj):
+    """Returns a DATETIME or TIMESTAMP column value as a datetime object:
+
+      >>> datetime_or_None('2007-02-25 23:06:20')
+      datetime.datetime(2007, 2, 25, 23, 6, 20)
+
+    Illegal values are raise DataError
+
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+
+    try:
+        time_obj = datetime.datetime.strptime(obj, '%Y-%m-%d %H:%M:%S')
+        return time_obj
+    except Exception as err:
+        raise DataError("Not valid datetime struct: %s" % err)
+
+
+TIMEDELTA_RE = re.compile(r"(-)?(\d{1,3}):(\d{1,2}):(\d{1,2})(?:.(\d{1,6}))?")
+
+
+def convert_timedelta(obj):
+    """Returns a TIME column as a timedelta object:
+
+      >>> timedelta_or_None('25:06:17')
+      datetime.timedelta(1, 3977)
+      >>> timedelta_or_None('-25:06:17')
+      datetime.timedelta(-2, 83177)
+
+    Illegal values are returned as None:
+
+      >>> timedelta_or_None('random crap') is None
+      True
+
+    Note that MySQL always returns TIME columns as (+|-)HH:MM:SS, but
+    can accept values as (+|-)DD HH:MM:SS. The latter format will not
+    be parsed correctly by this function.
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+
+    m = TIMEDELTA_RE.match(obj)
+    if not m:
+        return obj
+
+    try:
+        groups = list(m.groups())
+        groups[-1] = _convert_second_fraction(groups[-1])
+        negate = -1 if groups[0] else 1
+        hours, minutes, seconds, microseconds = groups[1:]
+
+        tdelta = datetime.timedelta(
+            hours=int(hours),
+            minutes=int(minutes),
+            seconds=int(seconds),
+            microseconds=int(microseconds)
+        ) * negate
+        return tdelta
+    except ValueError as err:
+        raise DataError("Not valid time or timedelta struct: %s" % err)
+
+
+def convert_time(obj):
+    """Returns a TIME column as a time object:
+
+      >>> time_or_None('15:06:17')
+      datetime.time(15, 6, 17)
+
+    Illegal values are returned DataError:
+
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+
+    try:
+        time_obj = datetime.datetime.strptime(obj, '%H:%M:%S')
+        return time_obj.time()
+    except Exception:
+        return convert_timedelta(obj)
+
+
+def convert_date(obj):
+    """Returns a DATE column as a date object:
+
+      >>> date_or_None('2007-02-26')
+      datetime.date(2007, 2, 26)
+
+    Illegal values are returned as None:
+
+      >>> date_or_None('2007-02-31') is None
+      True
+      >>> date_or_None('0000-00-00') is None
+      True
+
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
+    try:
+        time_obj = datetime.datetime.strptime(obj, '%Y-%m-%d')
+        return time_obj.date()
+    except Exception as err:
+        raise DataError("Not valid date struct: %s" % err)
+
+
+def convert_set(s):
+    if isinstance(s, (bytes, bytearray)):
+        return set(s.split(b","))
+    return set(s.split(","))
+
+
+def convert_characters(connection, data):
+    if connection.use_unicode:
+        data = data.decode("utf8")
+    return data
+
+
+def convert_column_data(column_type, column_data):
+    data = column_data
+
+    # Null
+    if data is None:
+        return data
+
+    if not isinstance(column_type, str):
+        return data
+
+    column_type = column_type.lower().strip()
+    if column_type == 'time':
+        data = convert_time(column_data)
+    elif column_type == 'date':
+        data = convert_date(column_data)
+    elif column_type == 'datetime':
+        data = convert_datetime(column_data)
+
+    return data
+
+
+encoders = {
+    bool: escape_bool,
+    int: escape_int,
+    float: escape_float,
+    str: escape_unicode,
+    tuple: escape_sequence,
+    list: escape_sequence,
+    set: escape_sequence,
+    frozenset: escape_sequence,
+    dict: escape_dict,
+    type(None): escape_None,
+    datetime.date: escape_date,
+    datetime.datetime: escape_datetime,
+    datetime.timedelta: escape_timedelta,
+    datetime.time: escape_time,
+    time.struct_time: escape_struct_time,
+    Decimal: escape_object,
+}

chdb/dbapi/cursors.py

@@ -0,0 +1,247 @@
+from . import err
+import re
+
+# Regular expression for :meth:`Cursor.executemany`.
+# executemany only supports simple bulk insert.
+# You can use it to load large dataset.
+RE_INSERT_VALUES = re.compile(
+    r"\s*((?:INSERT|REPLACE)\b.+\bVALUES?\s*)"
+    + r"(\(\s*(?:%s|%\(.+\)s)\s*(?:,\s*(?:%s|%\(.+\)s)\s*)*\))"
+    + r"(\s*(?:ON DUPLICATE.*)?);?\s*\Z",
+    re.IGNORECASE | re.DOTALL,
+)
+
+
+class Cursor(object):
+    """
+    This is the object you use to interact with the database.
+
+    Do not create an instance of a Cursor yourself. Call
+    connections.Connection.cursor().
+
+    See `Cursor <https://www.python.org/dev/peps/pep-0249/#cursor-objects>`_ in
+    the specification.
+    """
+
+    #: Max statement size which :meth:`executemany` generates.
+    #:
+    #: Default value is 1024000.
+    max_stmt_length = 1024000
+
+    def __init__(self, connection):
+        self.connection = connection
+        self._cursor = connection._conn.cursor()
+        self.description = None
+        self.rowcount = -1
+        self.arraysize = 1
+        self.lastrowid = None
+        self._executed = None
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *exc_info):
+        del exc_info
+        self.close()
+
+    def __iter__(self):
+        return iter(self.fetchone, None)
+
+    def callproc(self, procname, args=()):
+        """Execute stored procedure procname with args
+
+        procname -- string, name of procedure to execute on server
+
+        args -- Sequence of parameters to use with procedure
+
+        Returns the original args.
+
+        Compatibility warning: PEP-249 specifies that any modified
+        parameters must be returned. This is currently impossible
+        as they are only available by storing them in a server
+        variable and then retrieved by a query. Since stored
+        procedures return zero or more result sets, there is no
+        reliable way to get at OUT or INOUT parameters via callproc.
+        The server variables are named @_procname_n, where procname
+        is the parameter above and n is the position of the parameter
+        (from zero). Once all result sets generated by the procedure
+        have been fetched, you can issue a SELECT @_procname_0, ...
+        query using .execute() to get any OUT or INOUT values.
+
+        Compatibility warning: The act of calling a stored procedure
+        itself creates an empty result set. This appears after any
+        result sets generated by the procedure. This is non-standard
+        behavior with respect to the DB-API. Be sure to use nextset()
+        to advance through all result sets; otherwise you may get
+        disconnected.
+        """
+
+        return args
+
+    def close(self):
+        """
+        Closing a cursor just exhausts all remaining data.
+        """
+        self._cursor.close()
+
+    def _get_db(self):
+        if not self.connection:
+            raise err.ProgrammingError("Cursor closed")
+        return self.connection
+
+    def _escape_args(self, args, conn):
+        if isinstance(args, (tuple, list)):
+            return tuple(conn.escape(arg) for arg in args)
+        elif isinstance(args, dict):
+            return {key: conn.escape(val) for (key, val) in args.items()}
+        else:
+            # If it's not a dictionary let's try escaping it anyway.
+            # Worst case it will throw a Value error
+            return conn.escape(args)
+
+    def mogrify(self, query, args=None):
+        """
+        Returns the exact string that is sent to the database by calling the
+        execute() method.
+
+        This method follows the extension to the DB API 2.0 followed by Psycopg.
+        """
+        conn = self._get_db()
+
+        if args is not None:
+            query = query % self._escape_args(args, conn)
+
+        return query
+
+    def execute(self, query, args=None):
+        """Execute a query
+
+        :param str query: Query to execute.
+
+        :param args: parameters used with query. (optional)
+        :type args: tuple, list or dict
+
+        :return: Number of affected rows
+        :rtype: int
+
+        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.
+        """
+        if args is not None:
+            query = query % self._escape_args(args, self.connection)
+
+        self._cursor.execute(query)
+
+        # Get description from column names and types
+        if hasattr(self._cursor, "_column_names") and self._cursor._column_names:
+            self.description = [
+                (name, type_info, None, None, None, None, None)
+                for name, type_info in zip(
+                    self._cursor._column_names, self._cursor._column_types
+                )
+            ]
+            self.rowcount = (
+                len(self._cursor._current_table) if self._cursor._current_table else -1
+            )
+        else:
+            self.description = None
+            self.rowcount = -1
+
+        self._executed = query
+        return self.rowcount
+
+    def executemany(self, query, args):
+        # type: (str, list) -> int
+        """Run several data against one query
+
+        :param query: query to execute on server
+        :param args:  Sequence of sequences or mappings.  It is used as parameter.
+        :return: Number of rows affected, if any.
+
+        This method improves performance on multiple-row INSERT and
+        REPLACE. Otherwise, it is equivalent to looping over args with
+        execute().
+        """
+        if not args:
+            return 0
+
+        m = RE_INSERT_VALUES.match(query)
+        if m:
+            q_prefix = m.group(1) % ()
+            q_values = m.group(2).rstrip()
+            q_postfix = m.group(3) or ""
+            assert q_values[0] == "(" and q_values[-1] == ")"
+            return self._do_execute_many(
+                q_prefix,
+                q_values,
+                q_postfix,
+                args,
+                self.max_stmt_length,
+                self._get_db().encoding,
+            )
+
+        self.rowcount = sum(self.execute(query, arg) for arg in args)
+        return self.rowcount
+
+    def _do_execute_many(
+        self, prefix, values, postfix, args, max_stmt_length, encoding
+    ):
+        conn = self._get_db()
+        escape = self._escape_args
+        if isinstance(prefix, str):
+            prefix = prefix.encode(encoding)
+        if isinstance(postfix, str):
+            postfix = postfix.encode(encoding)
+        sql = prefix
+        args = iter(args)
+        v = values % escape(next(args), conn)
+        if isinstance(v, str):
+            v = v.encode(encoding, "surrogateescape")
+        sql += v
+        rows = 0
+        for arg in args:
+            v = values % escape(arg, conn)
+            if isinstance(v, str):
+                v = v.encode(encoding, "surrogateescape")
+            if len(sql) + len(v) + len(postfix) + 1 > max_stmt_length:
+                rows += self.execute(sql + postfix)
+                sql = prefix
+            else:
+                sql += ",".encode(encoding)
+            sql += v
+        rows += self.execute(sql + postfix)
+        self.rowcount = rows
+        return rows
+
+    def _check_executed(self):
+        if not self._executed:
+            raise err.ProgrammingError("execute() first")
+
+    def fetchone(self):
+        """Fetch the next row"""
+        if not self._executed:
+            raise err.ProgrammingError("execute() first")
+        return self._cursor.fetchone()
+
+    def fetchmany(self, size=1):
+        """Fetch several rows"""
+        if not self._executed:
+            raise err.ProgrammingError("execute() first")
+        return self._cursor.fetchmany(size)
+
+    def fetchall(self):
+        """Fetch all the rows"""
+        if not self._executed:
+            raise err.ProgrammingError("execute() first")
+        return self._cursor.fetchall()
+
+    def nextset(self):
+        """Get the next query set"""
+        # Not support for now
+        return None
+
+    def setinputsizes(self, *args):
+        """Does nothing, required by DB API."""
+
+    def setoutputsizes(self, *args):
+        """Does nothing, required by DB API."""