|
| 1 | +import logging |
| 2 | +import platform |
| 3 | +from ctypes import ( |
| 4 | + CDLL, POINTER, CFUNCTYPE, |
| 5 | + pointer, byref, string_at, cast, |
| 6 | + c_void_p, c_char_p, |
| 7 | + c_int, c_int64, c_uint64, c_double, c_char, |
| 8 | +) |
| 9 | + |
| 10 | +from sqlite3_error_map import sqlite_error_code_to_name |
| 11 | + |
| 12 | +logger = logging.getLogger("server") |
| 13 | + |
| 14 | +c_sqlite3_p = c_void_p |
| 15 | +c_sqlite3_stmt_p = c_void_p |
| 16 | +c_exec_callback_fn = CFUNCTYPE(c_int, c_void_p, c_int, POINTER(c_char_p), POINTER(c_char_p)) |
| 17 | +c_destructor_fn = CFUNCTYPE(None, c_void_p) |
| 18 | + |
| 19 | +libfile_platform = { |
| 20 | + "Linux": "libsqlite3.so", |
| 21 | + "Darwin": "libsqlite3.dylib", |
| 22 | +} |
| 23 | + |
| 24 | +platform_name = platform.system() |
| 25 | +libfile = libfile_platform[platform_name] |
| 26 | +lib = CDLL(libfile) |
| 27 | +lib.sqlite3_open_v2.argtypes = (c_char_p, POINTER(c_sqlite3_p), c_int, c_char_p,) |
| 28 | +lib.sqlite3_open_v2.restype = c_int |
| 29 | +lib.sqlite3_close_v2.argtypes = (c_sqlite3_p,) |
| 30 | +lib.sqlite3_close_v2.restype = c_int |
| 31 | +lib.sqlite3_extended_result_codes.argtypes = (c_sqlite3_p, c_int,) |
| 32 | +lib.sqlite3_extended_result_codes.restype = c_int |
| 33 | +lib.sqlite3_errmsg.argtypes = (c_sqlite3_p,) |
| 34 | +lib.sqlite3_errmsg.restype = c_char_p |
| 35 | +lib.sqlite3_errstr.argtypes = (c_int,) |
| 36 | +lib.sqlite3_errstr.restype = c_char_p |
| 37 | +lib.sqlite3_exec.argtypes = (c_sqlite3_p, c_char_p, c_exec_callback_fn, c_void_p, POINTER(c_char_p),) |
| 38 | +lib.sqlite3_exec.restype = c_int |
| 39 | +lib.sqlite3_txn_state.argtypes = (c_sqlite3_p, c_char_p,) |
| 40 | +lib.sqlite3_txn_state.restype = c_int |
| 41 | +lib.sqlite3_changes64.argtypes = (c_sqlite3_p,) |
| 42 | +lib.sqlite3_changes64.restype = c_int64 |
| 43 | +lib.sqlite3_total_changes64.argtypes = (c_sqlite3_p,) |
| 44 | +lib.sqlite3_total_changes64.restype = c_int64 |
| 45 | +lib.sqlite3_last_insert_rowid.argtypes = (c_sqlite3_p,) |
| 46 | +lib.sqlite3_last_insert_rowid.restype = c_int64 |
| 47 | +lib.sqlite3_limit.argtypes = (c_sqlite3_p, c_int, c_int,) |
| 48 | +lib.sqlite3_limit.restype = c_int |
| 49 | +lib.sqlite3_busy_timeout.argtypes = (c_sqlite3_p, c_int,) |
| 50 | +lib.sqlite3_busy_timeout.restype = c_int |
| 51 | +lib.sqlite3_get_autocommit.argtypes = (c_sqlite3_p,) |
| 52 | +lib.sqlite3_get_autocommit.restype = c_int |
| 53 | + |
| 54 | +lib.sqlite3_prepare_v2.argtypes = ( |
| 55 | + c_sqlite3_p, c_void_p, c_int, POINTER(c_sqlite3_stmt_p), POINTER(c_void_p),) |
| 56 | +lib.sqlite3_prepare_v2.restype = c_int |
| 57 | +lib.sqlite3_finalize.argtypes = (c_sqlite3_stmt_p,) |
| 58 | +lib.sqlite3_finalize.restype = c_int |
| 59 | +lib.sqlite3_step.argtypes = (c_sqlite3_stmt_p,) |
| 60 | +lib.sqlite3_step.restype = c_int |
| 61 | +lib.sqlite3_bind_parameter_count.argtypes = (c_sqlite3_stmt_p,) |
| 62 | +lib.sqlite3_bind_parameter_count.restype = c_int |
| 63 | +lib.sqlite3_bind_parameter_index.argtypes = (c_sqlite3_stmt_p, c_char_p,) |
| 64 | +lib.sqlite3_bind_parameter_index.restype = c_int |
| 65 | +lib.sqlite3_bind_parameter_name.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 66 | +lib.sqlite3_bind_parameter_name.restype = c_char_p |
| 67 | +lib.sqlite3_bind_blob64.argtypes = (c_sqlite3_stmt_p, c_int, c_void_p, c_uint64, c_destructor_fn,) |
| 68 | +lib.sqlite3_bind_blob64.restype = c_int |
| 69 | +lib.sqlite3_bind_text.argtypes = (c_sqlite3_stmt_p, c_int, POINTER(c_char), c_int, c_destructor_fn,) |
| 70 | +lib.sqlite3_bind_text.restype = c_int |
| 71 | +lib.sqlite3_bind_double.argtypes = (c_sqlite3_stmt_p, c_int, c_double,) |
| 72 | +lib.sqlite3_bind_double.restype = c_int |
| 73 | +lib.sqlite3_bind_int64.argtypes = (c_sqlite3_stmt_p, c_int, c_int64,) |
| 74 | +lib.sqlite3_bind_int64.restype = c_int |
| 75 | +lib.sqlite3_bind_null.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 76 | +lib.sqlite3_bind_null.restype = c_int |
| 77 | +lib.sqlite3_column_count.argtypes = (c_sqlite3_stmt_p,) |
| 78 | +lib.sqlite3_column_count.restype = c_int |
| 79 | +lib.sqlite3_column_name.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 80 | +lib.sqlite3_column_name.restype = c_char_p |
| 81 | +lib.sqlite3_column_decltype.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 82 | +lib.sqlite3_column_decltype.restype = c_char_p |
| 83 | +lib.sqlite3_column_type.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 84 | +lib.sqlite3_column_type.restype = c_int |
| 85 | +lib.sqlite3_column_blob.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 86 | +lib.sqlite3_column_blob.restype = c_void_p |
| 87 | +lib.sqlite3_column_text.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 88 | +lib.sqlite3_column_text.restype = c_void_p |
| 89 | +lib.sqlite3_column_bytes.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 90 | +lib.sqlite3_column_bytes.restype = c_int |
| 91 | +lib.sqlite3_column_double.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 92 | +lib.sqlite3_column_double.restype = c_double |
| 93 | +lib.sqlite3_column_int64.argtypes = (c_sqlite3_stmt_p, c_int,) |
| 94 | +lib.sqlite3_column_int64.restype = c_int64 |
| 95 | +lib.sqlite3_stmt_readonly.argtypes = (c_sqlite3_stmt_p,) |
| 96 | +lib.sqlite3_stmt_readonly.restype = c_int |
| 97 | +lib.sqlite3_stmt_isexplain.argtypes = (c_sqlite3_stmt_p,) |
| 98 | +lib.sqlite3_stmt_isexplain.restype = c_int |
| 99 | + |
| 100 | +SQLITE_OPEN_READWRITE = 0x00000002 |
| 101 | +SQLITE_OPEN_CREATE = 0x00000004 |
| 102 | +SQLITE_TRANSIENT = c_destructor_fn(-1) |
| 103 | +SQLITE_ROW = 100 |
| 104 | +SQLITE_DONE = 101 |
| 105 | + |
| 106 | +SQLITE_INTEGER = 1 |
| 107 | +SQLITE_FLOAT = 2 |
| 108 | +SQLITE_BLOB = 4 |
| 109 | +SQLITE_NULL = 5 |
| 110 | +SQLITE_TEXT = 3 |
| 111 | + |
| 112 | +SQLITE_LIMIT_LENGTH = 0 |
| 113 | +SQLITE_LIMIT_SQL_LENGTH = 1 |
| 114 | +SQLITE_LIMIT_COLUMN = 2 |
| 115 | +SQLITE_LIMIT_EXPR_DEPTH = 3 |
| 116 | +SQLITE_LIMIT_COMPOUND_SELECT = 4 |
| 117 | +SQLITE_LIMIT_VDBE_OP = 5 |
| 118 | +SQLITE_LIMIT_FUNCTION_ARG = 6 |
| 119 | +SQLITE_LIMIT_ATTACHED = 7 |
| 120 | +SQLITE_LIMIT_LIKE_PATTERN_LENGTH = 8 |
| 121 | +SQLITE_LIMIT_VARIABLE_NUMBER = 9 |
| 122 | +SQLITE_LIMIT_TRIGGER_DEPTH = 10 |
| 123 | +SQLITE_LIMIT_WORKER_THREADS = 11 |
| 124 | + |
| 125 | + |
| 126 | +class Conn: |
| 127 | + def __init__(self, db_ptr): |
| 128 | + self.db_ptr = db_ptr |
| 129 | + |
| 130 | + @classmethod |
| 131 | + def open(cls, filename): |
| 132 | + filename_ptr = c_char_p(filename.encode()) |
| 133 | + db_ptr = c_sqlite3_p() |
| 134 | + flags = SQLITE_OPEN_READWRITE|SQLITE_OPEN_CREATE |
| 135 | + vfs_ptr = c_char_p() |
| 136 | + _try(lib.sqlite3_open_v2(filename_ptr, byref(db_ptr), flags, vfs_ptr)) |
| 137 | + return cls(db_ptr) |
| 138 | + |
| 139 | + def close(self): |
| 140 | + if self.db_ptr is not None: |
| 141 | + lib.sqlite3_close_v2(self.db_ptr) |
| 142 | + self.db_ptr = None |
| 143 | + |
| 144 | + def __del__(self): |
| 145 | + self.close() |
| 146 | + |
| 147 | + def extended_result_codes(self, onoff): |
| 148 | + assert self.db_ptr is not None |
| 149 | + lib.sqlite3_extended_result_codes(self.db_ptr, onoff) |
| 150 | + |
| 151 | + def errmsg(self): |
| 152 | + assert self.db_ptr is not None |
| 153 | + return str(lib.sqlite3_errmsg(self.db_ptr).decode()) |
| 154 | + |
| 155 | + @classmethod |
| 156 | + def errstr(cls, code): |
| 157 | + return str(lib.sqlite3_errstr(code).decode()) |
| 158 | + |
| 159 | + def exec(self, sql): |
| 160 | + assert self.db_ptr is not None |
| 161 | + sql_ptr = c_char_p(sql.encode()) |
| 162 | + callback_ptr = c_exec_callback_fn() |
| 163 | + arg_ptr = c_void_p() |
| 164 | + errmsg_ptr_ptr = pointer(c_char_p()) |
| 165 | + _try(lib.sqlite3_exec(self.db_ptr, sql_ptr, callback_ptr, arg_ptr, errmsg_ptr_ptr), self) |
| 166 | + |
| 167 | + def txn_state(self): |
| 168 | + assert self.db_ptr is not None |
| 169 | + schema_ptr = c_char_p() |
| 170 | + return lib.sqlite3_txn_state(self.db_ptr, schema_ptr) |
| 171 | + |
| 172 | + def prepare(self, sql): |
| 173 | + assert self.db_ptr is not None |
| 174 | + sql = sql.encode() |
| 175 | + sql_data = c_char_p(sql) |
| 176 | + sql_ptr = cast(sql_data, c_void_p) |
| 177 | + sql_len = c_int(len(sql) + 1) |
| 178 | + stmt_ptr = c_sqlite3_stmt_p() |
| 179 | + tail_ptr = c_void_p() |
| 180 | + _try(lib.sqlite3_prepare_v2(self.db_ptr, sql_ptr, sql_len, byref(stmt_ptr), byref(tail_ptr)), self) |
| 181 | + if stmt_ptr.value is None: |
| 182 | + return None, b"" |
| 183 | + tail = sql[tail_ptr.value - sql_ptr.value:] |
| 184 | + return Stmt(self, stmt_ptr), tail.decode() |
| 185 | + |
| 186 | + def changes(self): |
| 187 | + assert self.db_ptr is not None |
| 188 | + return lib.sqlite3_changes64(self.db_ptr) |
| 189 | + |
| 190 | + def total_changes(self): |
| 191 | + assert self.db_ptr is not None |
| 192 | + return lib.sqlite3_total_changes64(self.db_ptr) |
| 193 | + |
| 194 | + def last_insert_rowid(self): |
| 195 | + assert self.db_ptr is not None |
| 196 | + return lib.sqlite3_last_insert_rowid(self.db_ptr) |
| 197 | + |
| 198 | + def limit(self, id, new_val): |
| 199 | + assert self.db_ptr is not None |
| 200 | + return lib.sqlite3_limit(self.db_ptr, id, new_val) |
| 201 | + |
| 202 | + def busy_timeout(self, ms): |
| 203 | + assert self.db_ptr is not None |
| 204 | + lib.sqlite3_busy_timeout(self.db_ptr, ms) |
| 205 | + |
| 206 | + def get_autocommit(self): |
| 207 | + assert self.db_ptr is not None |
| 208 | + return lib.sqlite3_get_autocommit(self.db_ptr) != 0 |
| 209 | + |
| 210 | +class Stmt: |
| 211 | + def __init__(self, conn, stmt_ptr): |
| 212 | + self.conn = conn |
| 213 | + self.stmt_ptr = stmt_ptr |
| 214 | + |
| 215 | + def close(self): |
| 216 | + if self.stmt_ptr is not None: |
| 217 | + lib.sqlite3_finalize(self.stmt_ptr) |
| 218 | + self.stmt_ptr = None |
| 219 | + |
| 220 | + def __del__(self): |
| 221 | + self.close() |
| 222 | + |
| 223 | + def param_count(self): |
| 224 | + assert self.stmt_ptr is not None |
| 225 | + return lib.sqlite3_bind_parameter_count(self.stmt_ptr) |
| 226 | + |
| 227 | + def param_index(self, name): |
| 228 | + assert self.stmt_ptr is not None |
| 229 | + name_ptr = c_char_p(name.encode()) |
| 230 | + return lib.sqlite3_bind_parameter_index(self.stmt_ptr, name_ptr) |
| 231 | + |
| 232 | + def param_name(self, param_i): |
| 233 | + assert self.stmt_ptr is not None |
| 234 | + name = lib.sqlite3_bind_parameter_name(self.stmt_ptr, param_i) |
| 235 | + return name.decode() if name is not None else None |
| 236 | + |
| 237 | + def bind(self, param_i, value): |
| 238 | + assert self.stmt_ptr is not None |
| 239 | + if isinstance(value, str): |
| 240 | + value = value.encode() |
| 241 | + value_ptr, value_len = c_char_p(value), c_int(len(value)) |
| 242 | + _try(lib.sqlite3_bind_text(self.stmt_ptr, param_i, value_ptr, value_len, SQLITE_TRANSIENT), self.conn) |
| 243 | + elif isinstance(value, bytes): |
| 244 | + value_ptr, value_len = c_char_p(value), c_uint64(len(value)) |
| 245 | + _try(lib.sqlite3_bind_blob64(self.stmt_ptr, param_i, value_ptr, value_len, SQLITE_TRANSIENT), self.conn) |
| 246 | + elif isinstance(value, int): |
| 247 | + _try(lib.sqlite3_bind_int64(self.stmt_ptr, param_i, c_int64(value)), self.conn) |
| 248 | + elif isinstance(value, float): |
| 249 | + _try(lib.sqlite3_bind_double(self.stmt_ptr, param_i, c_double(value)), self.conn) |
| 250 | + elif value is None: |
| 251 | + _try(lib.sqlite3_bind_null(self.stmt_ptr, param_i), self.conn) |
| 252 | + else: |
| 253 | + raise ValueError(f"Cannot bind {type(value)!r}") |
| 254 | + |
| 255 | + def step(self): |
| 256 | + assert self.stmt_ptr is not None |
| 257 | + res = lib.sqlite3_step(self.stmt_ptr) |
| 258 | + if res == SQLITE_DONE: |
| 259 | + return False |
| 260 | + elif res == SQLITE_ROW: |
| 261 | + return True |
| 262 | + _try(res, self.conn) |
| 263 | + |
| 264 | + def column_count(self): |
| 265 | + assert self.stmt_ptr is not None |
| 266 | + return lib.sqlite3_column_count(self.stmt_ptr) |
| 267 | + |
| 268 | + def column_name(self, column_i): |
| 269 | + assert self.stmt_ptr is not None |
| 270 | + return lib.sqlite3_column_name(self.stmt_ptr, column_i).decode() |
| 271 | + |
| 272 | + def column_decltype(self, column_i): |
| 273 | + assert self.stmt_ptr is not None |
| 274 | + name = lib.sqlite3_column_decltype(self.stmt_ptr, column_i) |
| 275 | + return name.decode() if name is not None else name |
| 276 | + |
| 277 | + def column(self, column_i): |
| 278 | + assert self.stmt_ptr is not None |
| 279 | + typ = lib.sqlite3_column_type(self.stmt_ptr, column_i) |
| 280 | + if typ == SQLITE_INTEGER: |
| 281 | + return lib.sqlite3_column_int64(self.stmt_ptr, column_i) |
| 282 | + elif typ == SQLITE_FLOAT: |
| 283 | + return lib.sqlite3_column_double(self.stmt_ptr, column_i) |
| 284 | + elif typ == SQLITE_BLOB: |
| 285 | + data_ptr = lib.sqlite3_column_blob(self.stmt_ptr, column_i) |
| 286 | + data_len = lib.sqlite3_column_bytes(self.stmt_ptr, column_i) |
| 287 | + return bytes(string_at(data_ptr, data_len)) |
| 288 | + elif typ == SQLITE_TEXT: |
| 289 | + data_ptr = lib.sqlite3_column_text(self.stmt_ptr, column_i) |
| 290 | + data_len = lib.sqlite3_column_bytes(self.stmt_ptr, column_i) |
| 291 | + b = bytes(string_at(data_ptr, data_len)) |
| 292 | + try: |
| 293 | + return b.decode() |
| 294 | + except UnicodeDecodeError: |
| 295 | + logger.debug("Could not decode column %s, bytes %s", column_i, b, exc_info=True) |
| 296 | + raise |
| 297 | + elif typ == SQLITE_NULL: |
| 298 | + return None |
| 299 | + else: |
| 300 | + raise ValueError(f"Unknown SQLite type {typ}") |
| 301 | + |
| 302 | + def readonly(self): |
| 303 | + assert self.stmt_ptr is not None |
| 304 | + return lib.sqlite3_stmt_readonly(self.stmt_ptr) != 0 |
| 305 | + |
| 306 | + def isexplain(self): |
| 307 | + assert self.stmt_ptr is not None |
| 308 | + return lib.sqlite3_stmt_isexplain(self.stmt_ptr) |
| 309 | + |
| 310 | + |
| 311 | +class SqliteError(RuntimeError): |
| 312 | + def __init__(self, message, error_code=None) -> None: |
| 313 | + super().__init__(message) |
| 314 | + self.error_code = error_code |
| 315 | + self.error_name = sqlite_error_code_to_name.get(error_code) |
| 316 | + |
| 317 | + |
| 318 | +def _try(error_code, conn=None): |
| 319 | + if error_code == 0: |
| 320 | + return |
| 321 | + |
| 322 | + error_str = Conn.errstr(error_code) |
| 323 | + if conn is not None: |
| 324 | + details = f": {conn.errmsg()}" |
| 325 | + |
| 326 | + message = f"SQLite function returned error code {error_code} ({error_str}){details}" |
| 327 | + raise SqliteError(message, error_code) |
0 commit comments