blob: 2c0a26b3b0e599c43243c195a7bba828f01e1feb [file] [log] [blame]
koder aka kdanilov22d134e2016-11-08 11:33:19 +02001"""
2This module contains interfaces for storage classes
3"""
4
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +02005import os
koder aka kdanilov22d134e2016-11-08 11:33:19 +02006import abc
koder aka kdanilov39e449e2016-12-17 15:15:26 +02007import array
koder aka kdanilov23e6bdf2016-12-24 02:18:54 +02008import shutil
koder aka kdanilov7f59d562016-12-26 01:34:23 +02009import sqlite3
10from typing import Any, TypeVar, Type, IO, Tuple, cast, List, Dict, Iterable
koder aka kdanilov39e449e2016-12-17 15:15:26 +020011
12import yaml
13try:
14 from yaml import CLoader as Loader, CDumper as Dumper # type: ignore
15except ImportError:
16 from yaml import Loader, Dumper # type: ignore
koder aka kdanilov22d134e2016-11-08 11:33:19 +020017
18
koder aka kdanilov7f59d562016-12-26 01:34:23 +020019from .result_classes import Storable, IStorable
koder aka kdanilov22d134e2016-11-08 11:33:19 +020020
21
koder aka kdanilov22d134e2016-11-08 11:33:19 +020022class ISimpleStorage(metaclass=abc.ABCMeta):
23 """interface for low-level storage, which doesn't support serialization
24 and can operate only on bytes"""
25
26 @abc.abstractmethod
koder aka kdanilov7f59d562016-12-26 01:34:23 +020027 def put(self, value: bytes, path: str) -> None:
koder aka kdanilov22d134e2016-11-08 11:33:19 +020028 pass
29
30 @abc.abstractmethod
koder aka kdanilov7f59d562016-12-26 01:34:23 +020031 def get(self, path: str) -> bytes:
koder aka kdanilov22d134e2016-11-08 11:33:19 +020032 pass
33
34 @abc.abstractmethod
koder aka kdanilov7f59d562016-12-26 01:34:23 +020035 def rm(self, path: str) -> None:
36 pass
37
38 @abc.abstractmethod
39 def sync(self) -> None:
koder aka kdanilov73084622016-11-16 21:51:08 +020040 pass
41
42 @abc.abstractmethod
koder aka kdanilov22d134e2016-11-08 11:33:19 +020043 def __contains__(self, path: str) -> bool:
44 pass
45
46 @abc.abstractmethod
koder aka kdanilov7f59d562016-12-26 01:34:23 +020047 def get_fd(self, path: str, mode: str = "rb+") -> IO:
koder aka kdanilov39e449e2016-12-17 15:15:26 +020048 pass
49
50 @abc.abstractmethod
51 def sub_storage(self, path: str) -> 'ISimpleStorage':
koder aka kdanilov22d134e2016-11-08 11:33:19 +020052 pass
53
54
55class ISerializer(metaclass=abc.ABCMeta):
56 """Interface for serialization class"""
57 @abc.abstractmethod
koder aka kdanilov7f59d562016-12-26 01:34:23 +020058 def pack(self, value: Storable) -> bytes:
koder aka kdanilov22d134e2016-11-08 11:33:19 +020059 pass
60
61 @abc.abstractmethod
koder aka kdanilov7f59d562016-12-26 01:34:23 +020062 def unpack(self, data: bytes) -> Any:
koder aka kdanilov22d134e2016-11-08 11:33:19 +020063 pass
64
65
koder aka kdanilov7f59d562016-12-26 01:34:23 +020066class DBStorage(ISimpleStorage):
67
68 create_tb_sql = "CREATE TABLE IF NOT EXISTS wally_storage (key text, data blob, type text)"
69 insert_sql = "INSERT INTO wally_storage VALUES (?, ?, ?)"
70 update_sql = "UPDATE wally_storage SET data=?, type=? WHERE key=?"
71 select_sql = "SELECT data, type FROM wally_storage WHERE key=?"
72 contains_sql = "SELECT 1 FROM wally_storage WHERE key=?"
73 rm_sql = "DELETE FROM wally_storage WHERE key LIKE '{}%'"
74 list2_sql = "SELECT key, length(data), type FROM wally_storage"
75
76 def __init__(self, db_path: str = None, existing: bool = False,
77 prefix: str = None, db: sqlite3.Connection = None) -> None:
78
79 assert not prefix or "'" not in prefix, "Broken sql prefix {!r}".format(prefix)
80
81 if db_path:
82 self.existing = existing
83 if existing:
84 if not os.path.isfile(db_path):
85 raise IOError("No storage found at {!r}".format(db_path))
86
87 os.makedirs(os.path.dirname(db_path), exist_ok=True)
88 try:
89 self.db = sqlite3.connect(db_path)
90 except sqlite3.OperationalError as exc:
91 raise IOError("Can't open database at {!r}".format(db_path)) from exc
92
93 self.db.execute(self.create_tb_sql)
94 else:
95 if db is None:
96 raise ValueError("Either db or db_path parameter must be passed")
97 self.db = db
98
99 if prefix is None:
100 self.prefix = ""
101 elif not prefix.endswith('/'):
102 self.prefix = prefix + '/'
103 else:
104 self.prefix = prefix
105
106 def put(self, value: bytes, path: str) -> None:
107 c = self.db.cursor()
108 fpath = self.prefix + path
109 c.execute(self.contains_sql, (fpath,))
110 if len(c.fetchall()) == 0:
111 c.execute(self.insert_sql, (fpath, value, 'yaml'))
112 else:
113 c.execute(self.update_sql, (value, 'yaml', fpath))
114
115 def get(self, path: str) -> bytes:
116 c = self.db.cursor()
117 c.execute(self.select_sql, (self.prefix + path,))
118 res = cast(List[Tuple[bytes, str]], c.fetchall()) # type: List[Tuple[bytes, str]]
119 if not res:
120 raise KeyError(path)
121 assert len(res) == 1
122 val, tp = res[0]
123 assert tp == 'yaml'
124 return val
125
126 def rm(self, path: str) -> None:
127 c = self.db.cursor()
128 path = self.prefix + path
129 assert "'" not in path, "Broken sql path {!r}".format(path)
130 c.execute(self.rm_sql.format(path))
131
132 def __contains__(self, path: str) -> bool:
133 c = self.db.cursor()
134 path = self.prefix + path
135 c.execute(self.contains_sql, (self.prefix + path,))
136 return len(c.fetchall()) != 0
137
138 def print_tree(self):
139 c = self.db.cursor()
140 c.execute(self.list2_sql)
141 data = list(c.fetchall())
142 data.sort()
143 print("------------------ DB ---------------------")
144 for key, data_ln, type in data:
145 print(key, data_ln, type)
146 print("------------------ END --------------------")
147
148 def get_fd(self, path: str, mode: str = "rb+") -> IO[bytes]:
149 raise NotImplementedError("SQLITE3 doesn't provide fd-like interface")
150
151 def sub_storage(self, path: str) -> 'DBStorage':
152 return self.__class__(prefix=self.prefix + path, db=self.db)
153
154 def sync(self):
155 self.db.commit()
156
157
158DB_REL_PATH = "__db__.db"
159
160
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200161class FSStorage(ISimpleStorage):
koder aka kdanilov22d134e2016-11-08 11:33:19 +0200162 """Store all data in files on FS"""
163
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200164 def __init__(self, root_path: str, existing: bool) -> None:
165 self.root_path = root_path
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200166 self.existing = existing
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200167
168 def j(self, path: str) -> str:
169 return os.path.join(self.root_path, path)
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200170
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200171 def put(self, value: bytes, path: str) -> None:
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200172 jpath = self.j(path)
173 os.makedirs(os.path.dirname(jpath), exist_ok=True)
174 with open(jpath, "wb") as fd:
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200175 fd.write(value)
176
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200177 def get(self, path: str) -> bytes:
koder aka kdanilov73084622016-11-16 21:51:08 +0200178 try:
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200179 with open(self.j(path), "rb") as fd:
180 return fd.read()
181 except FileNotFoundError as exc:
182 raise KeyError(path) from exc
koder aka kdanilov73084622016-11-16 21:51:08 +0200183
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200184 def rm(self, path: str) -> None:
185 if os.path.isdir(path):
186 shutil.rmtree(path, ignore_errors=True)
187 elif os.path.exists(path):
188 os.unlink(path)
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200189
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200190 def __contains__(self, path: str) -> bool:
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200191 return os.path.exists(self.j(path))
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200192
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200193 def get_fd(self, path: str, mode: str = "rb+") -> IO[bytes]:
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200194 jpath = self.j(path)
195
196 if "cb" == mode:
197 create_on_fail = True
198 mode = "rb+"
199 else:
200 create_on_fail = False
201
koder aka kdanilov962ee5f2016-12-19 02:40:08 +0200202 os.makedirs(os.path.dirname(jpath), exist_ok=True)
203
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200204 try:
205 fd = open(jpath, mode)
206 except IOError:
207 if not create_on_fail:
208 raise
209 fd = open(jpath, "wb")
210
211 return cast(IO[bytes], fd)
212
213 def sub_storage(self, path: str) -> 'FSStorage':
214 return self.__class__(self.j(path), self.existing)
koder aka kdanilov22d134e2016-11-08 11:33:19 +0200215
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200216 def sync(self):
217 pass
koder aka kdanilov23e6bdf2016-12-24 02:18:54 +0200218
koder aka kdanilov22d134e2016-11-08 11:33:19 +0200219
220class YAMLSerializer(ISerializer):
221 """Serialize data to yaml"""
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200222 def pack(self, value: Storable) -> bytes:
223 try:
224 return yaml.dump(value, Dumper=Dumper, encoding="utf8")
225 except Exception as exc:
226 raise ValueError("Can't pickle object {!r} to yaml".format(type(value))) from exc
koder aka kdanilov22d134e2016-11-08 11:33:19 +0200227
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200228 def unpack(self, data: bytes) -> Any:
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200229 return yaml.load(data, Loader=Loader)
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200230
231
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200232class SAFEYAMLSerializer(ISerializer):
233 """Serialize data to yaml"""
234 def pack(self, value: Storable) -> bytes:
235 try:
236 return yaml.safe_dump(value, encoding="utf8")
237 except Exception as exc:
238 raise ValueError("Can't pickle object {!r} to yaml".format(type(value))) from exc
239
240 def unpack(self, data: bytes) -> Any:
241 return yaml.safe_load(data)
242
243
244ObjClass = TypeVar('ObjClass', bound=IStorable)
245
246
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200247class Storage:
248 """interface for storage"""
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200249 def __init__(self, fs_storage: ISimpleStorage, db_storage: ISimpleStorage, serializer: ISerializer) -> None:
250 self.fs = fs_storage
251 self.db = db_storage
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200252 self.serializer = serializer
253
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200254 def sub_storage(self, *path: str) -> 'Storage':
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200255 fpath = "/".join(path)
256 return self.__class__(self.fs.sub_storage(fpath), self.db.sub_storage(fpath), self.serializer)
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200257
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200258 def put(self, value: Storable, *path: str) -> None:
259 dct_value = value.raw() if isinstance(value, IStorable) else value
260 serialized = self.serializer.pack(dct_value)
261 fpath = "/".join(path)
262 self.db.put(serialized, fpath)
263 self.fs.put(serialized, fpath)
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200264
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200265 def put_list(self, value: Iterable[IStorable], *path: str) -> None:
266 serialized = self.serializer.pack([obj.raw() for obj in value])
267 fpath = "/".join(path)
268 self.db.put(serialized, fpath)
269 self.fs.put(serialized, fpath)
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200270
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200271 def get(self, *path: str) -> Any:
272 return self.serializer.unpack(self.db.get("/".join(path)))
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200273
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200274 def rm(self, *path: str) -> None:
275 fpath = "/".join(path)
276 self.fs.rm(fpath)
277 self.db.rm(fpath)
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200278
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200279 def __contains__(self, path: str) -> bool:
280 return path in self.fs or path in self.db
koder aka kdanilov73084622016-11-16 21:51:08 +0200281
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200282 def put_raw(self, val: bytes, *path: str) -> None:
283 self.fs.put(val, "/".join(path))
koder aka kdanilov3af3c332016-12-19 17:12:34 +0200284
285 def get_raw(self, *path: str) -> bytes:
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200286 return self.fs.get("/".join(path))
koder aka kdanilov3af3c332016-12-19 17:12:34 +0200287
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200288 def get_fd(self, path: str, mode: str = "r") -> IO:
289 return self.fs.get_fd(path, mode)
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200290
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200291 def put_array(self, value: array.array, *path: str) -> None:
292 with self.get_fd("/".join(path), "wb") as fd:
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200293 value.tofile(fd) # type: ignore
294
295 def get_array(self, typecode: str, *path: str) -> array.array:
296 res = array.array(typecode)
297 path_s = "/".join(path)
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200298 with self.get_fd(path_s, "rb") as fd:
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200299 fd.seek(0, os.SEEK_END)
300 size = fd.tell()
301 fd.seek(0, os.SEEK_SET)
302 assert size % res.itemsize == 0, "Storage object at path {} contains no array of {} or corrupted."\
303 .format(path_s, typecode)
304 res.fromfile(fd, size // res.itemsize) # type: ignore
305 return res
306
307 def append(self, value: array.array, *path: str) -> None:
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200308 with self.get_fd("/".join(path), "cb") as fd:
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200309 fd.seek(0, os.SEEK_END)
310 value.tofile(fd) # type: ignore
311
koder aka kdanilov70227062016-11-26 23:23:21 +0200312 def load_list(self, obj_class: Type[ObjClass], *path: str) -> List[ObjClass]:
313 path_s = "/".join(path)
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200314 raw_val = cast(List[Dict[str, Any]], self.get(path_s))
koder aka kdanilov73084622016-11-16 21:51:08 +0200315 assert isinstance(raw_val, list)
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200316 return [obj_class.fromraw(val) for val in raw_val]
koder aka kdanilov73084622016-11-16 21:51:08 +0200317
koder aka kdanilov70227062016-11-26 23:23:21 +0200318 def load(self, obj_class: Type[ObjClass], *path: str) -> ObjClass:
319 path_s = "/".join(path)
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200320 return obj_class.fromraw(self.get(path_s))
koder aka kdanilov73084622016-11-16 21:51:08 +0200321
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200322 def sync(self) -> None:
323 self.db.sync()
324 self.fs.sync()
koder aka kdanilov73084622016-11-16 21:51:08 +0200325
koder aka kdanilov39e449e2016-12-17 15:15:26 +0200326 def __enter__(self) -> 'Storage':
327 return self
328
329 def __exit__(self, x: Any, y: Any, z: Any) -> None:
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200330 self.sync()
koder aka kdanilov70227062016-11-26 23:23:21 +0200331
koder aka kdanilov3d2bc4f2016-11-12 18:31:18 +0200332
333def make_storage(url: str, existing: bool = False) -> Storage:
koder aka kdanilov7f59d562016-12-26 01:34:23 +0200334 return Storage(FSStorage(url, existing),
335 DBStorage(os.path.join(url, DB_REL_PATH)),
336 SAFEYAMLSerializer())
koder aka kdanilov22d134e2016-11-08 11:33:19 +0200337