You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
454 lines
12 KiB
454 lines
12 KiB
|
|
import os |
|
import shutil |
|
import sqlite3 |
|
import contextlib |
|
from collections import OrderedDict |
|
|
|
from ..base import WriteConcern |
|
from ..types import unicode_, bson |
|
from . import ( |
|
AbstractStorage, |
|
AbstractDatabase, |
|
AbstractCollection, |
|
AbstractCursor, |
|
|
|
StorageDuplicateKeyError, |
|
) |
|
|
|
|
|
sqlite_324 = sqlite3.sqlite_version_info >= (3, 24, 0) |
|
|
|
|
|
""" |
|
(NOTE) SQLite3 pragmas DEFAULT value: |
|
|
|
* Applies to the database |
|
- journal_mode=delete |
|
|
|
* Applies to the connection |
|
- busy_timeout=5000 |
|
- foreign_keys=OFF |
|
- automatic_index=ON |
|
""" |
|
|
|
|
|
SQLITE_DB_EXT = ".collection" |
|
SQLITE_RECORD_TABLE = "documents" |
|
|
|
|
|
"""SQL""" |
|
|
|
CREATE_TABLE = """ |
|
CREATE TABLE [{}]( |
|
k text NOT NULL, |
|
v text NOT NULL, |
|
PRIMARY KEY(k) |
|
); |
|
""" |
|
|
|
INSERT_RECORD = """ |
|
INSERT INTO [{}](k, v) VALUES (?, ?); |
|
""" |
|
|
|
UPSERT_RECORD = """ |
|
INSERT INTO [{}] (v, k) VALUES(?, ?) |
|
ON CONFLICT(k) |
|
DO UPDATE SET v=excluded.v; |
|
""" # need sqlite_version >= 3.24.0 |
|
|
|
UPDATE_RECORD = """ |
|
UPDATE [{}] SET v = (?) WHERE k = (?); |
|
""" |
|
|
|
INSORE_RECORD = """ |
|
INSERT OR IGNORE INTO [{}](v, k) VALUES (?, ?); |
|
""" # for sqlite_version < 3.24, UPDATE + INSERT_IGNORE = UPSERT SCRIPT |
|
|
|
DELETE_RECORD = """ |
|
DELETE FROM [{}] WHERE k = (?); |
|
""" |
|
|
|
SELECT_ALL_RECORD = """ |
|
SELECT v FROM [{}]; |
|
""" |
|
|
|
SELECT_LIMIT_RECORD = """ |
|
SELECT v FROM [{0}] LIMIT {1}; |
|
""" |
|
|
|
SELECT_ALL_KEYS = """ |
|
SELECT k FROM [{}]; |
|
""" |
|
|
|
|
|
class SQLiteKVEngine(object): |
|
|
|
def __init__(self, config): |
|
self.__conn = None |
|
|
|
self.__db_pragmas = { |
|
key: config[key] |
|
for key in [ |
|
"journal_mode", |
|
] |
|
if key in config |
|
} |
|
|
|
self.__conn_kwargs = { |
|
key: cast(config[key]) |
|
for key, cast in [ |
|
("check_same_thread", bool) |
|
] |
|
if key in config |
|
} |
|
|
|
@property |
|
def db_pragmas(self): |
|
return self._assemble_pragmas(self.__db_pragmas) |
|
|
|
def _connect(self, db_file, wconcern=None): |
|
self.__conn = sqlite3.connect(db_file, **self.__conn_kwargs) |
|
self.__conn.text_factory = str |
|
|
|
wcon_pragmas = "" |
|
if wconcern: |
|
wcon_doc = wconcern.document |
|
# wtimeout (milliseconds) -> busy_timeout (milliseconds) |
|
timeout = wcon_doc.get("wtimeout") |
|
if timeout: |
|
del wcon_doc["wtimeout"] |
|
wcon_doc["busy_timeout"] = timeout |
|
|
|
wcon_pragmas = self._assemble_pragmas(wcon_doc) |
|
|
|
# update connection pragmas and write_concern pragmas |
|
self.__conn.executescript(self.db_pragmas + ";" + wcon_pragmas) |
|
|
|
return contextlib.closing(self.__conn) |
|
|
|
def _assemble_pragmas(self, pragma_dict): |
|
return ";".join(["PRAGMA {0}={1}".format(k, v) |
|
for k, v in pragma_dict.items()]) |
|
|
|
def create_table(self, db_file): |
|
with self._connect(db_file) as conn: |
|
with conn: |
|
conn.execute(CREATE_TABLE.format(SQLITE_RECORD_TABLE)) |
|
|
|
def write_one(self, db_file, params, wconcern=None): |
|
with self._connect(db_file, wconcern) as conn: |
|
with conn: |
|
sql = INSERT_RECORD.format(SQLITE_RECORD_TABLE) |
|
conn.execute(sql, params) |
|
|
|
def write_many(self, db_file, seq_params, wconcern=None): |
|
with self._connect(db_file, wconcern) as conn: |
|
with conn: |
|
sql = INSERT_RECORD.format(SQLITE_RECORD_TABLE) |
|
conn.executemany(sql, seq_params) |
|
|
|
def update_one(self, db_file, params, wconcern=None): |
|
with self._connect(db_file, wconcern) as conn: |
|
with conn: |
|
sql = UPDATE_RECORD.format(SQLITE_RECORD_TABLE) |
|
conn.execute(sql, params) |
|
|
|
def update_many(self, db_file, seq_params, wconcern=None): |
|
with self._connect(db_file, wconcern) as conn: |
|
with conn: |
|
sql = UPDATE_RECORD.format(SQLITE_RECORD_TABLE) |
|
conn.executemany(sql, seq_params) |
|
|
|
def delete_one(self, db_file, params, wconcern=None): |
|
with self._connect(db_file, wconcern) as conn: |
|
with conn: |
|
sql = DELETE_RECORD.format(SQLITE_RECORD_TABLE) |
|
conn.execute(sql, params) |
|
|
|
def delete_many(self, db_file, seq_params, wconcern=None): |
|
with self._connect(db_file, wconcern) as conn: |
|
with conn: |
|
sql = DELETE_RECORD.format(SQLITE_RECORD_TABLE) |
|
conn.executemany(sql, seq_params) |
|
|
|
def read_all(self, db_file, limit): |
|
if not os.path.isfile(db_file): |
|
return [] |
|
with self._connect(db_file) as conn: |
|
with conn: |
|
if limit: |
|
sql = SELECT_LIMIT_RECORD.format( |
|
SQLITE_RECORD_TABLE, limit) |
|
else: |
|
sql = SELECT_ALL_RECORD.format(SQLITE_RECORD_TABLE) |
|
|
|
return conn.execute(sql).fetchall() |
|
|
|
def read_all_keys(self, db_file): |
|
if not os.path.isfile(db_file): |
|
return [] |
|
with self._connect(db_file) as conn: |
|
with conn: |
|
sql = SELECT_ALL_KEYS.format(SQLITE_RECORD_TABLE) |
|
return conn.execute(sql).fetchall() |
|
|
|
|
|
class SQLiteWriteConcern(WriteConcern): |
|
""" |
|
|
|
Args: |
|
busy_timeout (int): Default 5000 |
|
|
|
synchronous (int, str): Default "NORMAL" |
|
- type: integer |
|
enum: [0, 1, 2, 3] |
|
- type: string |
|
enum: ["OFF", NORMAL, FULL, EXTRA, "0", "1", "2", "3"] |
|
|
|
automatic_index (bool, str): Default False |
|
- type: boolean |
|
- type: string |
|
enum: ["ON", "OFF"] |
|
|
|
""" |
|
|
|
def __init__(self, |
|
busy_timeout=5000, |
|
synchronous="NORMAL", |
|
automatic_index=False): |
|
|
|
super(SQLiteWriteConcern, self).__init__(busy_timeout) |
|
|
|
if synchronous is not None: |
|
self._document["synchronous"] = synchronous |
|
if automatic_index is not None: |
|
self._document["automatic_index"] = automatic_index |
|
|
|
def __repr__(self): |
|
return ("SQLiteWriteConcern({})".format( |
|
", ".join("%s=%s" % kvt for kvt in self.document.items()),)) |
|
|
|
|
|
class SQLiteStorage(AbstractStorage): |
|
""" |
|
""" |
|
|
|
def __init__(self, repository, storage_config): |
|
super(SQLiteStorage, self).__init__(repository, storage_config) |
|
self._conn = SQLiteKVEngine(self._config) |
|
|
|
def _db_path(self, db_name): |
|
""" |
|
Get Monty database dir path. |
|
""" |
|
return os.path.join(self._repository, db_name) |
|
|
|
@classmethod |
|
def nice_name(cls): |
|
return "sqlite" |
|
|
|
@classmethod |
|
def config(cls, journal_mode="WAL", check_same_thread=True, **kwargs): |
|
""" |
|
|
|
Args: |
|
journal_mode (str): Default "WAL" |
|
type: string |
|
enum: [DELETE, TRUNCATE, PERSIST, MEMORY, WAL, "OFF"] |
|
check_same_thread (bool): Default True |
|
See `sqlite3.connect` |
|
|
|
""" |
|
return { |
|
"journal_mode": journal_mode, |
|
"check_same_thread": check_same_thread, |
|
} |
|
|
|
def wconcern_parser(self, |
|
wtimeout=None, |
|
busy_timeout=None, |
|
synchronous=None, |
|
automatic_index=None, |
|
**kwargs): |
|
return SQLiteWriteConcern(wtimeout or busy_timeout, |
|
synchronous, |
|
automatic_index) |
|
|
|
def database_create(self, db_name): |
|
if not os.path.isdir(self._db_path(db_name)): |
|
os.makedirs(self._db_path(db_name)) |
|
|
|
def database_drop(self, db_name): |
|
db_path = self._db_path(db_name) |
|
if os.path.isdir(db_path): |
|
shutil.rmtree(db_path) |
|
|
|
def database_list(self): |
|
return [ |
|
name for name in os.listdir(unicode_(self._repository)) |
|
if os.path.isdir(self._db_path(name)) |
|
] |
|
|
|
|
|
class SQLiteDatabase(AbstractDatabase): |
|
|
|
def __init__(self, storage, subject): |
|
super(SQLiteDatabase, self).__init__(storage, subject) |
|
self._db_path = storage._db_path(self._name) |
|
|
|
def _col_path(self, col_name): |
|
""" |
|
Get SQLite database file path, which is Monty collection. |
|
""" |
|
return os.path.join(self._db_path, col_name) + SQLITE_DB_EXT |
|
|
|
@property |
|
def _conn(self): |
|
return self._storage._conn |
|
|
|
def database_exists(self): |
|
return os.path.isdir(self._db_path) |
|
|
|
def collection_exists(self, col_name): |
|
return os.path.isfile(self._col_path(col_name)) |
|
|
|
def collection_create(self, col_name): |
|
if not self.database_exists(): |
|
self._storage.database_create(self._name) |
|
self._conn.create_table(self._col_path(col_name)) |
|
|
|
def collection_drop(self, col_name): |
|
if self.collection_exists(col_name): |
|
os.remove(self._col_path(col_name)) |
|
|
|
def collection_list(self): |
|
if not self.database_exists(): |
|
return [] |
|
return [os.path.splitext(name)[0] |
|
for name in os.listdir(unicode_(self._db_path))] |
|
|
|
|
|
SQLiteStorage.contractor_cls = SQLiteDatabase |
|
|
|
|
|
class SQLiteCollection(AbstractCollection): |
|
|
|
def __init__(self, database, subject): |
|
super(SQLiteCollection, self).__init__(database, subject) |
|
|
|
self._col_path = self._database._col_path(self._name) |
|
|
|
def _ensure_table(func): |
|
def make_table(self, *args, **kwargs): |
|
if not self._database.collection_exists(self._name): |
|
self._database.collection_create(self._name) |
|
return func(self, *args, **kwargs) |
|
return make_table |
|
|
|
@property |
|
def _conn(self): |
|
return self._database._conn |
|
|
|
@_ensure_table |
|
def write_one(self, doc, check_keys=True): |
|
""" |
|
""" |
|
_id = doc["_id"] |
|
try: |
|
self._conn.write_one( |
|
self._col_path, |
|
(bson.id_encode(_id), |
|
self._encode_doc(doc, check_keys),), |
|
self.wconcern |
|
) |
|
except sqlite3.IntegrityError: |
|
raise StorageDuplicateKeyError() |
|
|
|
return _id |
|
|
|
@_ensure_table |
|
def write_many(self, docs, check_keys=True, ordered=True): |
|
""" |
|
""" |
|
_docs = OrderedDict() |
|
ids = list() |
|
keys = self._conn.read_all_keys(self._col_path) |
|
has_duplicated_key = False |
|
for doc in docs: |
|
_id = doc["_id"] |
|
b_id = bson.id_encode(_id) |
|
if b_id in keys or b_id in _docs: |
|
has_duplicated_key = True |
|
break |
|
|
|
_docs[b_id] = self._encode_doc(doc, check_keys) |
|
ids.append(_id) |
|
|
|
self._conn.write_many( |
|
self._col_path, |
|
_docs.items(), |
|
self.wconcern |
|
) |
|
|
|
if has_duplicated_key: |
|
raise StorageDuplicateKeyError() |
|
|
|
return ids |
|
|
|
def update_one(self, doc): |
|
""" |
|
""" |
|
self._conn.update_one( |
|
self._col_path, |
|
(self._encode_doc(doc), bson.id_encode(doc["_id"])), |
|
self.wconcern |
|
) |
|
|
|
def update_many(self, docs): |
|
""" |
|
""" |
|
self._conn.update_many( |
|
self._col_path, |
|
[(self._encode_doc(doc), bson.id_encode(doc["_id"])) |
|
for doc in docs], |
|
self.wconcern |
|
) |
|
|
|
def delete_one(self, id): |
|
self._conn.delete_one( |
|
self._col_path, |
|
(bson.id_encode(id),), |
|
self.wconcern |
|
) |
|
|
|
def delete_many(self, ids): |
|
self._conn.delete_many( |
|
self._col_path, |
|
[(bson.id_encode(id),) for id in ids], |
|
self.wconcern |
|
) |
|
|
|
|
|
SQLiteDatabase.contractor_cls = SQLiteCollection |
|
|
|
|
|
class SQLiteCursor(AbstractCursor): |
|
|
|
def __init__(self, collection, subject): |
|
super(SQLiteCursor, self).__init__(collection, subject) |
|
|
|
@property |
|
def _conn(self): |
|
return self._collection._conn |
|
|
|
@property |
|
def _col_path(self): |
|
return self._collection._col_path |
|
|
|
def query(self, max_scan): |
|
docs = self._conn.read_all(self._col_path, max_scan) |
|
return (self._decode_doc(doc[0]) for doc in docs) |
|
|
|
|
|
SQLiteCollection.contractor_cls = SQLiteCursor
|
|
|