machbaseapi 2.0.0__py3-none-any.whl

machbaseAPI/machbaseAPI.py

@@ -0,0 +1,318 @@
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Callable, Dict, List, Optional, Sequence
+
+from .connector import MachbaseConnection
+from .errors import Error, ProgrammingError
+
+gCharset = "UTF-8"
+_SELECT_PREFIX = re.compile(r"^\s*(SELECT|WITH|DESC|DESCRIBE|SHOW)\b", re.IGNORECASE)
+
+
+def _stringify(value: Any) -> Any:
+    if value is None:
+        return None
+    if isinstance(value, (bytes, bytearray, memoryview)):
+        value = bytes(value)
+        return value.hex().upper()
+    if isinstance(value, float):
+        return str(int(value)) if value.is_integer() else str(value)
+    if isinstance(value, bool):
+        return "1" if value else "0"
+    return str(value)
+
+
+def _encode_row(row: Dict[str, Any]) -> str:
+    encoded = {str(key).lower(): _stringify(value) for key, value in row.items()}
+    if "_arrival_time" in encoded and "_ARRIVAL_TIME" not in encoded:
+        encoded["_ARRIVAL_TIME"] = encoded["_arrival_time"]
+    return json.dumps(encoded, separators=(",", ":"))
+
+
+def _encode_rows(rows: Sequence[Dict[str, Any]]) -> str:
+    return ",".join(_encode_row(row) for row in rows)
+
+
+class machbaseAPI(object):
+    pass
+
+
+class machbase(object):
+    def __init__(self):
+        self._conn: Optional[MachbaseConnection] = None
+        self._opened = False
+        self._append_session = None
+        self._rows: List[Dict[str, Any]] = []
+        self._last_result: str = ""
+        self._append_types: Optional[List[int]] = None
+
+    def _set_ok(self, text: str) -> int:
+        self._last_result = json.dumps({"EXECUTE RESULT": text}, separators=(",", ":"))
+        return 1
+
+    def _set_error(self, text: str, detail: str = "") -> int:
+        payload = {"EXECUTE ERROR": text}
+        if detail:
+            payload["MACHBASE_ERROR"] = detail
+        self._last_result = json.dumps(payload, separators=(",", ":"))
+        return 0
+
+    def _ensure_conn(self) -> MachbaseConnection:
+        if self._conn is None:
+            raise RuntimeError("Connection is not open")
+        return self._conn
+
+    @staticmethod
+    def _normalize_append_payload(
+        types_or_values: Any,
+        explicit_values: Any = None,
+    ) -> tuple[Optional[List[int]], List[Sequence[Any]]]:
+        if explicit_values is not None:
+            if types_or_values is None:
+                types = None
+            else:
+                types = [int(value) for value in types_or_values]
+            raw_rows = explicit_values
+        else:
+            types = None
+            raw_rows = types_or_values
+
+        if raw_rows is None:
+            return types, []
+
+        if isinstance(raw_rows, (list, tuple)):
+            if len(raw_rows) == 0:
+                return types, []
+            if isinstance(raw_rows[0], (list, tuple)):
+                rows = [list(row) for row in raw_rows]
+            else:
+                rows = [list(raw_rows)]
+        else:
+            raise ProgrammingError("Append values must be sequence of rows.")
+
+        return types, rows
+
+    def isConnected(self):
+        return 1 if self._conn and self._conn.is_connected() else 0
+
+    def isOpened(self):
+        return 1 if self._opened else 0
+
+    def open(self, aHost="127.0.0.1", aUser="SYS", aPw="MANAGER", aPort=5656):
+        try:
+            self._conn = MachbaseConnection(host=aHost, user=aUser.upper(), password=aPw, port=int(aPort))
+            self._opened = True
+            return 1
+        except Exception as exc:
+            self._conn = None
+            self._opened = False
+            return self._set_error("Connection error.", str(exc))
+
+    def openEx(self, aHost="127.0.0.1", aUser="SYS", aPw="MANAGER", aPort=5656, aConnStr=""):
+        try:
+            self._conn = MachbaseConnection(
+                host=aHost,
+                user=aUser.upper(),
+                password=aPw,
+                port=int(aPort),
+                connection_string=aConnStr,
+            )
+            self._opened = True
+            return 1
+        except Exception as exc:
+            self._conn = None
+            self._opened = False
+            return self._set_error("Connection error.", str(exc))
+
+    def close(self):
+        if self._conn is None:
+            self._opened = False
+            return 1
+        try:
+            self._conn.close()
+            self._opened = False
+            self._conn = None
+            self._append_session = None
+            self._rows = []
+            return 1
+        except Exception as exc:
+            return self._set_error("Connection close error.", str(exc))
+
+    def execute(self, aSql):
+        try:
+            if _SELECT_PREFIX.match(aSql):
+                if self.select(aSql) == 0:
+                    return 0
+                return 1
+            self._ensure_conn()._client.exec_direct(aSql)
+            self._rows = []
+            return self._set_ok("Execute Success")
+        except Exception as exc:
+            return self._set_error("SQLPrepare error.", str(exc))
+
+    def select(self, aSql):
+        try:
+            result = self._ensure_conn()._client.query(aSql)
+            self._rows = list(result.rows)
+            self._last_result = _encode_rows(self._rows)
+            return 1
+        except Exception as exc:
+            self._rows = []
+            return self._set_error("SQLPrepare error.", str(exc))
+
+    def fetch(self):
+        if not self._rows:
+            return 0, None
+        return 1, _encode_row(self._rows.pop(0))
+
+    def selectClose(self):
+        self._rows = []
+        return 1
+
+    def column(self, aTableName):
+        return self.columns(aTableName)
+
+    def statistics(self, aTableName, aUser="SYS"):
+        sql = (
+            "select name, type, colcount from m$tables "
+            "where user_id = (select user_id from m$users where name = '%s') and name = '%s'"
+            % (aUser.upper(), aTableName.upper())
+        )
+        return self.select(sql)
+
+    def schema(self, aSql):
+        return self.execute(aSql)
+
+    def tables(self):
+        sql = (
+            "select t.NAME as name, u.NAME as username, t.COLCOUNT as colcount "
+            "from m$sys_tables t, m$sys_users u where u.user_id=t.user_id"
+        )
+        return self.select(sql)
+
+    def columns(self, aTableName):
+        if aTableName is None:
+            sql = (
+                "select c.name name, c.type type, c.length length "
+                "from m$sys_columns c, m$sys_tables t where c.table_id = t.id "
+                "and c.id not in(0,65534) order by c.id"
+            )
+        else:
+            name = aTableName.upper()
+            if name.startswith("M$"):
+                sql = (
+                    "select name,type,length from m$columns where table_id = "
+                    "(select id from m$tables where name = '%s') and id not in(0,65534) order by id" % name
+                )
+            elif name.startswith("V$"):
+                sql = (
+                    "select name,type,length from v$columns where table_id = "
+                    "(select id from v$tables where name = '%s') and id not in(0,65534) order by id" % name
+                )
+            else:
+                sql = (
+                    "select name,type,length from m$sys_columns where table_id = "
+                    "(select id from m$sys_tables where name = '%s') and name not in('_RID','_ARRIVAL_TIME') order by id"
+                    % name
+                )
+        return self.select(sql)
+
+    def appendOpen(self, aTableName, aTypes=None):
+        try:
+            self._append_session = self._ensure_conn()._client.append_open(aTableName)
+            self._append_types = [int(v) for v in aTypes] if aTypes else None
+            return 1
+        except Exception as exc:
+            return self._set_error("Append open error.", str(exc))
+
+    def appendClose(self):
+        if self._append_session is None:
+            return 0
+        try:
+            self._ensure_conn()._client.append_close(self._append_session)
+            self._append_session = None
+            self._append_types = None
+            return self._set_ok("Append success")
+        except Exception as exc:
+            return self._set_error("Append close error.", str(exc))
+
+    def appendData(
+        self,
+        aTableName,
+        aTypes,
+        aValues=None,
+        aFormat="YYYY-MM-DD HH24:MI:SS",
+        on_ack: Optional[Callable[[dict], None]] = None,
+    ):
+        _ = aTableName, aFormat
+        if self._append_session is None:
+            return 0
+        try:
+            types, rows = self._normalize_append_payload(aTypes, aValues)
+            self._ensure_conn()._client.append_data(
+                self._append_session,
+                rows,
+                types=types or self._append_types,
+                on_ack=on_ack,
+            )
+            return 1
+        except Exception as exc:
+            return self._set_error("Append data error.", str(exc))
+
+    def appendDataByTime(
+        self,
+        aTableName,
+        aTypes,
+        aValues,
+        aFormat="YYYY-MM-DD HH24:MI:SS",
+        aTimes=None,
+        on_ack: Optional[Callable[[dict], None]] = None,
+    ):
+        _ = aTableName, aFormat
+        if self._append_session is None:
+            return 0
+        try:
+            types, rows = self._normalize_append_payload(aTypes, aValues)
+            self._ensure_conn()._client.append_data(
+                self._append_session,
+                rows,
+                types=types or self._append_types,
+                times=aTimes,
+                on_ack=on_ack,
+            )
+            return 1
+        except Exception as exc:
+            return self._set_error("Append data error.", str(exc))
+
+    def appendFlush(self):
+        return 1 if self._append_session is not None else 0
+
+    def append(self, aTableName, aTypes, aValues=None, aFormat="YYYY-MM-DD HH24:MI:SS"):
+        _ = aFormat
+        types, rows = self._normalize_append_payload(aTypes, aValues)
+        if not rows:
+            return 0
+        try:
+            count = self._ensure_conn()._client.append(aTableName, rows, types=types)
+            self._last_result = json.dumps({"EXECUTE RESULT": "Append success"}, separators=(",", ":"))
+            return 1
+        except Exception as exc:
+            return self._set_error("Append error.", str(exc))
+
+    def appendByTime(self, aTableName, aTypes, aValues=None, aFormat="YYYY-MM-DD HH24:MI:SS", aTimes=None):
+        _ = aFormat
+        if aValues is None and isinstance(aTypes, (list, tuple)) and aTimes is not None and aTypes and isinstance(aTypes[0], (list, tuple)):
+            aValues = aTypes
+            aTypes = None
+        try:
+            types, rows = self._normalize_append_payload(aTypes, aValues)
+            count = self._ensure_conn()._client.append(aTableName, rows, types=types, times=aTimes)
+            self._last_result = json.dumps({"EXECUTE RESULT": "Append success"}, separators=(",", ":"))
+            return 1
+        except Exception as exc:
+            return self._set_error("Append error.", str(exc))
+
+    def result(self):
+        return self._last_result

machbaseAPI/marshal.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, Iterable, Iterator, List, Optional, Union
+
+from .constants import CMI_BINARY_TYPE, CMI_PACKET_MAX_BODY, CMI_ROWS_TYPE, CMI_STRING_TYPE
+
+
+MarshalPrimitive = Union[int, float, bool, bytes, bytearray, str, None]
+
+
+@dataclass
+class MarshalUnit:
+    id: int
+    type: int
+    length: int
+    data: bytes
+
+
+def _align8(length: int) -> int:
+    return (length + 7) & ~7
+
+
+def _read_u32le(payload: bytes) -> int:
+    if len(payload) < 4:
+        return 0
+    return int.from_bytes(payload[:4], "big")
+
+
+def _read_fixed_size(type_id: int) -> Optional[int]:
+    if type_id in {
+        CMI_STRING_TYPE,
+        CMI_BINARY_TYPE,
+        0x0000000C,
+        0x000000F1,
+        0x000000F2,
+        0x0000000D,
+    }:
+        return None
+    if type_id in {0x00000004, 0x00000005}:
+        return 1
+    if type_id in {0x00000006, 0x00000007}:  # 16-bit
+        return 2
+    if type_id in {0x00000008, 0x00000009, 0x00000018, 0x00000019}:  # 32-bit
+        return 4
+    return 8
+
+
+def decode_scalar(unit: MarshalUnit) -> MarshalPrimitive:
+    if unit.type in {
+        CMI_STRING_TYPE,
+        0x0000000C,
+        0x000000F1,
+        0x000000F2,
+        0x0000000D,
+    }:
+        try:
+            return unit.data.decode("utf-8")
+        except Exception:
+            return unit.data.decode("utf-8", errors="replace")
+    if unit.type == 0x00000004:
+        return unit.data[0] if unit.data else 0
+    if unit.type == 0x00000005:
+        return unit.data[0] if unit.data else 0
+    if unit.type == 0x00000006:
+        return int.from_bytes(unit.data[:2], "little", signed=True) if len(unit.data) >= 2 else 0
+    if unit.type == 0x00000007:
+        return int.from_bytes(unit.data[:2], "little", signed=False) if len(unit.data) >= 2 else 0
+    if unit.type == 0x00000008:
+        return int.from_bytes(unit.data[:4], "little", signed=True) if len(unit.data) >= 4 else 0
+    if unit.type == 0x00000009:
+        return int.from_bytes(unit.data[:4], "little", signed=False) if len(unit.data) >= 4 else 0
+    if unit.type == 0x0000000A:
+        return int.from_bytes(unit.data[:8], "little", signed=True) if len(unit.data) >= 8 else 0
+    if unit.type == 0x0000000B:
+        return int.from_bytes(unit.data[:8], "little", signed=False) if len(unit.data) >= 8 else 0
+    return unit.data
+
+
+class MarshalWriter:
+    """Build marshal units and split them into protocol packets."""
+
+    HEADER_SIZE = 16
+
+    def __init__(self, protocol_id: int, stmt_id: int, adds: int = 0):
+        self.protocol_id = protocol_id & 0xFF
+        self.stmt_id = stmt_id & 0xFFFFFFFF
+        self.adds = adds & 0xFFFF
+        self._units: List[bytes] = []
+        self._current_length = 0
+        self._bodies: List[bytes] = []
+
+    def add_string(self, unit_id: int, value: str) -> None:
+        payload = value.encode("utf-8")
+        self._add_variable(unit_id, CMI_STRING_TYPE, payload)
+
+    def add_binary(self, unit_id: int, value: Union[bytes, bytearray]) -> None:
+        self._add_variable(unit_id, CMI_BINARY_TYPE, bytes(value))
+
+    def add_uint32(self, unit_id: int, value: int) -> None:
+        payload = value.to_bytes(4, "little", signed=False)
+        self._add_fixed(unit_id, 0x00000009, payload)
+
+    def add_sint32(self, unit_id: int, value: int) -> None:
+        payload = value.to_bytes(4, "little", signed=True)
+        self._add_fixed(unit_id, 0x00000008, payload)
+
+    def add_uint64(self, unit_id: int, value: int) -> None:
+        payload = value.to_bytes(8, "little", signed=False)
+        self._add_fixed(unit_id, 0x0000000B, payload)
+
+    def add_sint64(self, unit_id: int, value: int) -> None:
+        payload = value.to_bytes(8, "little", signed=True)
+        self._add_fixed(unit_id, 0x0000000A, payload)
+
+    def add_boolean(self, unit_id: int, value: bool) -> None:
+        payload = b"\x01" if value else b"\x00"
+        self._add_fixed(unit_id, 0x00000026, payload)
+
+    def finalize(self) -> List[bytes]:
+        self._flush_current()
+        if not self._bodies:
+            self._bodies = [b""]
+        count = len(self._bodies)
+        result = []
+        for index, body in enumerate(self._bodies):
+            if count == 1:
+                flag = 0
+            elif index == 0:
+                flag = 1
+            elif index + 1 == count:
+                flag = 3
+            else:
+                flag = 2
+            result.append(self._build_packet(body, flag))
+        return result
+
+    def _add_variable(self, unit_id: int, type_id: int, payload: bytes) -> None:
+        length = _align8(len(payload))
+        unit = bytearray(self.HEADER_SIZE + length)
+        unit[0:4] = int.to_bytes(unit_id & 0xFFFFFFFF, 4, "little")
+        unit[4:8] = int.to_bytes(type_id & 0xFFFFFFFF, 4, "little")
+        unit[8:16] = int.to_bytes(len(payload), 8, "little")
+        unit[16 : 16 + len(payload)] = payload
+        self._enqueue(bytes(unit))
+
+    def _add_fixed(self, unit_id: int, type_id: int, payload: bytes) -> None:
+        unit = bytearray(self.HEADER_SIZE)
+        unit[0:4] = int.to_bytes(unit_id & 0xFFFFFFFF, 4, "little")
+        unit[4:8] = int.to_bytes(type_id & 0xFFFFFFFF, 4, "little")
+        fixed = bytes(payload)
+        unit[8 : 8 + len(fixed)] = fixed
+        self._enqueue(bytes(unit))
+
+    def _enqueue(self, unit: bytes) -> None:
+        if len(unit) > CMI_PACKET_MAX_BODY:
+            self._flush_current()
+            self._bodies.append(unit)
+            return
+        if self._current_length + len(unit) > CMI_PACKET_MAX_BODY:
+            self._flush_current()
+        self._units.append(unit)
+        self._current_length += len(unit)
+
+    def _flush_current(self) -> None:
+        if not self._units:
+            self._current_length = 0
+            return
+        body = b"".join(self._units)
+        self._bodies.append(body)
+        self._units.clear()
+        self._current_length = 0
+
+    def _build_packet(self, body: bytes, flag: int) -> bytes:
+        length_with_flag = (flag << 30) | (len(body) & 0x3FFFFFFF)
+        header = bytearray(16)
+        header[0] = 0
+        header[1:3] = self.adds.to_bytes(2, "big")
+        header[3] = self.protocol_id & 0xFF
+        header[4:8] = length_with_flag.to_bytes(4, "big")
+        header[8:12] = self.stmt_id.to_bytes(4, "big")
+        return bytes(header) + body
+
+
+class MarshalReader:
+    def __init__(self, payload: bytes):
+        self._payload = payload
+
+    def __iter__(self) -> Iterator[MarshalUnit]:
+        return self._iter_units()
+
+    def _iter_units(self) -> Iterator[MarshalUnit]:
+        offset = 0
+        while offset < len(self._payload):
+            if offset + 16 > len(self._payload):
+                raise ValueError("Incomplete marshal unit header")
+            unit_id = int.from_bytes(self._payload[offset : offset + 4], "little")
+            type_id = int.from_bytes(self._payload[offset + 4 : offset + 8], "little")
+            raw_length = int.from_bytes(self._payload[offset + 8 : offset + 16], "little")
+            body_offset = offset + 8
+
+            if type_id in {
+                CMI_STRING_TYPE,
+                CMI_BINARY_TYPE,
+                0x0000000C,
+                0x000000F1,
+                0x000000F2,
+                CMI_ROWS_TYPE,
+            }:
+                aligned = _align8(raw_length)
+                body_offset = offset + 16
+                if body_offset + aligned > len(self._payload):
+                    raise ValueError("Incomplete variable marshal payload")
+                data = self._payload[body_offset : body_offset + raw_length]
+                yield MarshalUnit(unit_id, type_id, raw_length, data)
+                offset = body_offset + aligned
+            else:
+                size = _read_fixed_size(type_id)
+                if size is None:
+                    size = 8
+                if body_offset + size > offset + 16:
+                    raise ValueError("Incomplete fixed marshal payload")
+                data = self._payload[body_offset : body_offset + size]
+                yield MarshalUnit(unit_id, type_id, size, data)
+                offset = offset + 16

machbaseAPI/packet.py

@@ -0,0 +1,87 @@
+"""Packet utilities for Machbase binary protocol."""
+
+from __future__ import annotations
+
+import socket
+import time
+from dataclasses import dataclass
+from typing import Deque, List, Optional
+
+from collections import deque
+
+
+@dataclass
+class Packet:
+    protocol: int
+    flag: int
+    adds: int
+    stmt_id: int
+    body: bytes
+
+
+class PacketReader:
+    HEADER_SIZE = 16
+
+    def __init__(self, sock: socket.socket):
+        self._socket = sock
+        self._buffer = b""
+        self._queue: Deque[Packet] = deque()
+        self._closed = False
+        self._error: Optional[BaseException] = None
+
+    def next(self, timeout_ms: Optional[int] = None) -> Packet:
+        if self._error:
+            raise self._error
+        if self._queue:
+            return self._queue.popleft()
+
+        deadline = None
+        if timeout_ms is not None:
+            deadline = time.monotonic() + timeout_ms / 1000
+
+        while True:
+            if self._queue:
+                return self._queue.popleft()
+            if self._closed:
+                raise ConnectionError("socket closed")
+            now = time.monotonic()
+            if deadline is not None and now >= deadline:
+                raise TimeoutError("timed out waiting packet")
+
+            timeout = None
+            if deadline is not None:
+                timeout = max(0.0, deadline - now)
+
+            try:
+                self._socket.settimeout(timeout)
+                chunk = self._socket.recv(65536)
+                if not chunk:
+                    self._closed = True
+                    if not self._queue:
+                        raise ConnectionError("connection closed by server")
+                    return self._queue.popleft()
+                self._buffer += chunk
+                self._parse_buffer()
+            except TimeoutError as exc:
+                raise exc
+            except Exception as exc:
+                self._error = exc
+                raise
+
+    def _parse_buffer(self) -> None:
+        while len(self._buffer) >= self.HEADER_SIZE:
+            length_field = int.from_bytes(self._buffer[4:8], "big")
+            flag = (length_field >> 30) & 0x03
+            body_length = length_field & 0x3FFFFFFF
+            total_length = self.HEADER_SIZE + body_length
+            if len(self._buffer) < total_length:
+                return
+            header = self._buffer[:self.HEADER_SIZE]
+            body = self._buffer[self.HEADER_SIZE : total_length]
+            self._buffer = self._buffer[total_length:]
+
+            protocol = header[3]
+            adds = int.from_bytes(header[1:3], "big")
+            stmt_id = int.from_bytes(header[8:12], "big")
+            self._queue.append(Packet(protocol=protocol, flag=flag, adds=adds, stmt_id=stmt_id, body=body))
+