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.
 
 
 

569 lines
21 KiB

  1. # -*- coding: utf-8 -*-
  2. """NilmDB
  3. Object that represents a NILM database file.
  4. Manages both the SQL database and the table storage backend.
  5. """
  6. # Need absolute_import so that "import nilmdb" won't pull in
  7. # nilmdb.py, but will pull the parent nilmdb module instead.
  8. from __future__ import absolute_import
  9. import nilmdb.utils
  10. from nilmdb.utils.printf import *
  11. from nilmdb.server.interval import (Interval, DBInterval,
  12. IntervalSet, IntervalError)
  13. from nilmdb.server import bulkdata
  14. from nilmdb.server.errors import NilmDBError, StreamError, OverlapError
  15. import sqlite3
  16. import os
  17. import errno
  18. import bisect
  19. # Note about performance and transactions:
  20. #
  21. # Committing a transaction in the default sync mode (PRAGMA synchronous=FULL)
  22. # takes about 125msec. sqlite3 will commit transactions at 3 times:
  23. # 1: explicit con.commit()
  24. # 2: between a series of DML commands and non-DML commands, e.g.
  25. # after a series of INSERT, SELECT, but before a CREATE TABLE or PRAGMA.
  26. # 3: at the end of an explicit transaction, e.g. "with self.con as con:"
  27. #
  28. # To speed up testing, or if this transaction speed becomes an issue,
  29. # the sync=False option to NilmDB.__init__ will set PRAGMA synchronous=OFF.
  30. # Don't touch old entries -- just add new ones.
  31. _sql_schema_updates = {
  32. 0: """
  33. -- All streams
  34. CREATE TABLE streams(
  35. id INTEGER PRIMARY KEY, -- stream ID
  36. path TEXT UNIQUE NOT NULL, -- path, e.g. '/newton/prep'
  37. layout TEXT NOT NULL -- layout name, e.g. float32_8
  38. );
  39. -- Individual timestamped ranges in those streams.
  40. -- For a given start_time and end_time, this tells us that the
  41. -- data is stored between start_pos and end_pos.
  42. -- Times are stored as μs since Unix epoch
  43. -- Positions are opaque: PyTables rows, file offsets, etc.
  44. --
  45. -- Note: end_pos points to the row _after_ end_time, so end_pos-1
  46. -- is the last valid row.
  47. CREATE TABLE ranges(
  48. stream_id INTEGER NOT NULL,
  49. start_time INTEGER NOT NULL,
  50. end_time INTEGER NOT NULL,
  51. start_pos INTEGER NOT NULL,
  52. end_pos INTEGER NOT NULL
  53. );
  54. CREATE INDEX _ranges_index ON ranges (stream_id, start_time, end_time);
  55. """,
  56. 1: """
  57. -- Generic dictionary-type metadata that can be associated with a stream
  58. CREATE TABLE metadata(
  59. stream_id INTEGER NOT NULL,
  60. key TEXT NOT NULL,
  61. value TEXT
  62. );
  63. """,
  64. }
  65. @nilmdb.utils.must_close()
  66. class NilmDB(object):
  67. verbose = 0
  68. def __init__(self, basepath, sync=True, max_results=None,
  69. bulkdata_args=None):
  70. if bulkdata_args is None:
  71. bulkdata_args = {}
  72. # set up path
  73. self.basepath = os.path.abspath(basepath)
  74. # Create the database path if it doesn't exist
  75. try:
  76. os.makedirs(self.basepath)
  77. except OSError as e:
  78. if e.errno != errno.EEXIST:
  79. raise IOError("can't create tree " + self.basepath)
  80. # Our data goes inside it
  81. self.data = bulkdata.BulkData(self.basepath, **bulkdata_args)
  82. # SQLite database too
  83. sqlfilename = os.path.join(self.basepath, "data.sql")
  84. self.con = sqlite3.connect(sqlfilename, check_same_thread = True)
  85. self._sql_schema_update()
  86. # See big comment at top about the performance implications of this
  87. if sync:
  88. self.con.execute("PRAGMA synchronous=FULL")
  89. else:
  90. self.con.execute("PRAGMA synchronous=OFF")
  91. # Approximate largest number of elements that we want to send
  92. # in a single reply (for stream_intervals, stream_extract)
  93. if max_results:
  94. self.max_results = max_results
  95. else:
  96. self.max_results = 16384
  97. def get_basepath(self):
  98. return self.basepath
  99. def close(self):
  100. if self.con:
  101. self.con.commit()
  102. self.con.close()
  103. self.data.close()
  104. def _sql_schema_update(self):
  105. cur = self.con.cursor()
  106. version = cur.execute("PRAGMA user_version").fetchone()[0]
  107. oldversion = version
  108. while version in _sql_schema_updates:
  109. cur.executescript(_sql_schema_updates[version])
  110. version = version + 1
  111. if self.verbose: # pragma: no cover
  112. printf("Schema updated to %d\n", version)
  113. if version != oldversion:
  114. with self.con:
  115. cur.execute("PRAGMA user_version = {v:d}".format(v=version))
  116. def _check_user_times(self, start, end):
  117. if start is None:
  118. start = -1e12
  119. if end is None:
  120. end = 1e12
  121. if start >= end:
  122. raise NilmDBError("start must precede end")
  123. return (start, end)
  124. @nilmdb.utils.lru_cache(size = 16)
  125. def _get_intervals(self, stream_id):
  126. """
  127. Return a mutable IntervalSet corresponding to the given stream ID.
  128. """
  129. iset = IntervalSet()
  130. result = self.con.execute("SELECT start_time, end_time, "
  131. "start_pos, end_pos "
  132. "FROM ranges "
  133. "WHERE stream_id=?", (stream_id,))
  134. try:
  135. for (start_time, end_time, start_pos, end_pos) in result:
  136. iset += DBInterval(start_time, end_time,
  137. start_time, end_time,
  138. start_pos, end_pos)
  139. except IntervalError: # pragma: no cover
  140. raise NilmDBError("unexpected overlap in ranges table!")
  141. return iset
  142. def _sql_interval_insert(self, id, start, end, start_pos, end_pos):
  143. """Helper that adds interval to the SQL database only"""
  144. self.con.execute("INSERT INTO ranges "
  145. "(stream_id,start_time,end_time,start_pos,end_pos) "
  146. "VALUES (?,?,?,?,?)",
  147. (id, start, end, start_pos, end_pos))
  148. def _sql_interval_delete(self, id, start, end, start_pos, end_pos):
  149. """Helper that removes interval from the SQL database only"""
  150. self.con.execute("DELETE FROM ranges WHERE "
  151. "stream_id=? AND start_time=? AND "
  152. "end_time=? AND start_pos=? AND end_pos=?",
  153. (id, start, end, start_pos, end_pos))
  154. def _add_interval(self, stream_id, interval, start_pos, end_pos):
  155. """
  156. Add interval to the internal interval cache, and to the database.
  157. Note: arguments must be ints (not numpy.int64, etc)
  158. """
  159. # Load this stream's intervals
  160. iset = self._get_intervals(stream_id)
  161. # Check for overlap
  162. if iset.intersects(interval): # pragma: no cover (gets caught earlier)
  163. raise NilmDBError("new interval overlaps existing data")
  164. # Check for adjacency. If there's a stream in the database
  165. # that ends exactly when this one starts, and the database
  166. # rows match up, we can make one interval that covers the
  167. # time range [adjacent.start -> interval.end)
  168. # and database rows [ adjacent.start_pos -> end_pos ].
  169. # Only do this if the resulting interval isn't too large.
  170. max_merged_rows = 8000 * 60 * 60 * 1.05 # 1.05 hours at 8 KHz
  171. adjacent = iset.find_end(interval.start)
  172. if (adjacent is not None and
  173. start_pos == adjacent.db_endpos and
  174. (end_pos - adjacent.db_startpos) < max_merged_rows):
  175. # First delete the old one, both from our iset and the
  176. # database
  177. iset -= adjacent
  178. self._sql_interval_delete(stream_id,
  179. adjacent.db_start, adjacent.db_end,
  180. adjacent.db_startpos, adjacent.db_endpos)
  181. # Now update our interval so the fallthrough add is
  182. # correct.
  183. interval.start = adjacent.start
  184. start_pos = adjacent.db_startpos
  185. # Add the new interval to the iset
  186. iset.iadd_nocheck(DBInterval(interval.start, interval.end,
  187. interval.start, interval.end,
  188. start_pos, end_pos))
  189. # Insert into the database
  190. self._sql_interval_insert(stream_id, interval.start, interval.end,
  191. int(start_pos), int(end_pos))
  192. self.con.commit()
  193. def _remove_interval(self, stream_id, original, remove):
  194. """
  195. Remove an interval from the internal cache and the database.
  196. stream_id: id of stream
  197. original: original DBInterval; must be already present in DB
  198. to_remove: DBInterval to remove; must be subset of 'original'
  199. """
  200. # Just return if we have nothing to remove
  201. if remove.start == remove.end: # pragma: no cover
  202. return
  203. # Load this stream's intervals
  204. iset = self._get_intervals(stream_id)
  205. # Remove existing interval from the cached set and the database
  206. iset -= original
  207. self._sql_interval_delete(stream_id,
  208. original.db_start, original.db_end,
  209. original.db_startpos, original.db_endpos)
  210. # Add back the intervals that would be left over if the
  211. # requested interval is removed. There may be two of them, if
  212. # the removed piece was in the middle.
  213. def add(iset, start, end, start_pos, end_pos):
  214. iset += DBInterval(start, end, start, end, start_pos, end_pos)
  215. self._sql_interval_insert(stream_id, start, end, start_pos, end_pos)
  216. if original.start != remove.start:
  217. # Interval before the removed region
  218. add(iset, original.start, remove.start,
  219. original.db_startpos, remove.db_startpos)
  220. if original.end != remove.end:
  221. # Interval after the removed region
  222. add(iset, remove.end, original.end,
  223. remove.db_endpos, original.db_endpos)
  224. # Commit SQL changes
  225. self.con.commit()
  226. return
  227. def stream_list(self, path = None, layout = None, extent = False):
  228. """Return list of lists of all streams in the database.
  229. If path is specified, include only streams with a path that
  230. matches the given string.
  231. If layout is specified, include only streams with a layout
  232. that matches the given string.
  233. If extent = False, returns a list of lists containing
  234. the path and layout: [ path, layout ]
  235. If extent = True, returns a list of lists containing the
  236. path, layout, and min/max extent of the data:
  237. [ path, layout, extent_min, extent_max ]
  238. """
  239. params = ()
  240. query = "SELECT streams.path, streams.layout"
  241. if extent:
  242. query += ", min(ranges.start_time), max(ranges.end_time)"
  243. query += " FROM streams"
  244. if extent:
  245. query += " LEFT JOIN ranges ON streams.id = ranges.stream_id"
  246. query += " WHERE 1=1"
  247. if layout is not None:
  248. query += " AND streams.layout=?"
  249. params += (layout,)
  250. if path is not None:
  251. query += " AND streams.path=?"
  252. params += (path,)
  253. query += " GROUP BY streams.id ORDER BY streams.path"
  254. result = self.con.execute(query, params).fetchall()
  255. return [ list(x) for x in result ]
  256. def stream_intervals(self, path, start = None, end = None):
  257. """
  258. Returns (intervals, restart) tuple.
  259. intervals is a list of [start,end] timestamps of all intervals
  260. that exist for path, between start and end.
  261. restart, if nonzero, means that there were too many results to
  262. return in a single request. The data is complete from the
  263. starting timestamp to the point at which it was truncated,
  264. and a new request with a start time of 'restart' will fetch
  265. the next block of data.
  266. """
  267. stream_id = self._stream_id(path)
  268. intervals = self._get_intervals(stream_id)
  269. (start, end) = self._check_user_times(start, end)
  270. requested = Interval(start, end)
  271. result = []
  272. for n, i in enumerate(intervals.intersection(requested)):
  273. if n >= self.max_results:
  274. restart = i.start
  275. break
  276. result.append([i.start, i.end])
  277. else:
  278. restart = 0
  279. return (result, restart)
  280. def stream_create(self, path, layout_name):
  281. """Create a new table in the database.
  282. path: path to the data (e.g. '/newton/prep').
  283. Paths must contain at least two elements, e.g.:
  284. /newton/prep
  285. /newton/raw
  286. /newton/upstairs/prep
  287. /newton/upstairs/raw
  288. layout_name: string for nilmdb.layout.get_named(), e.g. 'float32_8'
  289. """
  290. # Create the bulk storage. Raises ValueError on error, which we
  291. # pass along.
  292. self.data.create(path, layout_name)
  293. # Insert into SQL database once the bulk storage is happy
  294. with self.con as con:
  295. con.execute("INSERT INTO streams (path, layout) VALUES (?,?)",
  296. (path, layout_name))
  297. def _stream_id(self, path):
  298. """Return unique stream ID"""
  299. result = self.con.execute("SELECT id FROM streams WHERE path=?",
  300. (path,)).fetchone()
  301. if result is None:
  302. raise StreamError("No stream at path " + path)
  303. return result[0]
  304. def stream_set_metadata(self, path, data):
  305. """Set stream metadata from a dictionary, e.g.
  306. { description = 'Downstairs lighting',
  307. v_scaling = 123.45 }
  308. This replaces all existing metadata.
  309. """
  310. stream_id = self._stream_id(path)
  311. with self.con as con:
  312. con.execute("DELETE FROM metadata WHERE stream_id=?", (stream_id,))
  313. for key in data:
  314. if data[key] != '':
  315. con.execute("INSERT INTO metadata VALUES (?, ?, ?)",
  316. (stream_id, key, data[key]))
  317. def stream_get_metadata(self, path):
  318. """Return stream metadata as a dictionary."""
  319. stream_id = self._stream_id(path)
  320. result = self.con.execute("SELECT metadata.key, metadata.value "
  321. "FROM metadata "
  322. "WHERE metadata.stream_id=?", (stream_id,))
  323. data = {}
  324. for (key, value) in result:
  325. data[key] = value
  326. return data
  327. def stream_update_metadata(self, path, newdata):
  328. """Update stream metadata from a dictionary"""
  329. data = self.stream_get_metadata(path)
  330. data.update(newdata)
  331. self.stream_set_metadata(path, data)
  332. def stream_destroy(self, path):
  333. """Fully remove a table and all of its data from the database.
  334. No way to undo it! Metadata is removed."""
  335. stream_id = self._stream_id(path)
  336. # Delete the cached interval data (if it was cached)
  337. self._get_intervals.cache_remove(self, stream_id)
  338. # Delete the data
  339. self.data.destroy(path)
  340. # Delete metadata, stream, intervals
  341. with self.con as con:
  342. con.execute("DELETE FROM metadata WHERE stream_id=?", (stream_id,))
  343. con.execute("DELETE FROM ranges WHERE stream_id=?", (stream_id,))
  344. con.execute("DELETE FROM streams WHERE id=?", (stream_id,))
  345. def stream_insert(self, path, start, end, data):
  346. """Insert new data into the database.
  347. path: Path at which to add the data
  348. start: Starting timestamp
  349. end: Ending timestamp
  350. data: Textual data, formatted according to the layout of path
  351. """
  352. # First check for basic overlap using timestamp info given.
  353. stream_id = self._stream_id(path)
  354. iset = self._get_intervals(stream_id)
  355. interval = Interval(start, end)
  356. if iset.intersects(interval):
  357. raise OverlapError("new data overlaps existing data at range: "
  358. + str(iset & interval))
  359. # Tenatively append the data. This will raise a ValueError if
  360. # there are any parse errors.
  361. table = self.data.getnode(path)
  362. row_start = table.nrows
  363. table.append_string(data, start, end)
  364. row_end = table.nrows
  365. # Insert the record into the sql database.
  366. self._add_interval(stream_id, interval, row_start, row_end)
  367. # And that's all
  368. return
  369. def _find_start(self, table, dbinterval):
  370. """
  371. Given a DBInterval, find the row in the database that
  372. corresponds to the start time. Return the first database
  373. position with a timestamp (first element) greater than or
  374. equal to 'start'.
  375. """
  376. # Optimization for the common case where an interval wasn't truncated
  377. if dbinterval.start == dbinterval.db_start:
  378. return dbinterval.db_startpos
  379. return bisect.bisect_left(bulkdata.TimestampOnlyTable(table),
  380. dbinterval.start,
  381. dbinterval.db_startpos,
  382. dbinterval.db_endpos)
  383. def _find_end(self, table, dbinterval):
  384. """
  385. Given a DBInterval, find the row in the database that follows
  386. the end time. Return the first database position after the
  387. row with timestamp (first element) greater than or equal
  388. to 'end'.
  389. """
  390. # Optimization for the common case where an interval wasn't truncated
  391. if dbinterval.end == dbinterval.db_end:
  392. return dbinterval.db_endpos
  393. # Note that we still use bisect_left here, because we don't
  394. # want to include the given timestamp in the results. This is
  395. # so a queries like 1:00 -> 2:00 and 2:00 -> 3:00 return
  396. # non-overlapping data.
  397. return bisect.bisect_left(bulkdata.TimestampOnlyTable(table),
  398. dbinterval.end,
  399. dbinterval.db_startpos,
  400. dbinterval.db_endpos)
  401. def stream_extract(self, path, start = None, end = None, count = False):
  402. """
  403. Returns (data, restart) tuple.
  404. data is ASCII-formatted data from the database, formatted
  405. according to the layout of the stream.
  406. restart, if nonzero, means that there were too many results to
  407. return in a single request. The data is complete from the
  408. starting timestamp to the point at which it was truncated,
  409. and a new request with a start time of 'restart' will fetch
  410. the next block of data.
  411. count, if true, means to not return raw data, but just the count
  412. of rows that would have been returned. This is much faster
  413. than actually fetching the data. It is not limited by
  414. max_results.
  415. """
  416. stream_id = self._stream_id(path)
  417. table = self.data.getnode(path)
  418. intervals = self._get_intervals(stream_id)
  419. (start, end) = self._check_user_times(start, end)
  420. requested = Interval(start, end)
  421. result = []
  422. matched = 0
  423. remaining = self.max_results
  424. restart = 0
  425. for interval in intervals.intersection(requested):
  426. # Reading single rows from the table is too slow, so
  427. # we use two bisections to find both the starting and
  428. # ending row for this particular interval, then
  429. # read the entire range as one slice.
  430. row_start = self._find_start(table, interval)
  431. row_end = self._find_end(table, interval)
  432. if count:
  433. matched += row_end - row_start
  434. continue
  435. # Shorten it if we'll hit the maximum number of results
  436. row_max = row_start + remaining
  437. if row_max < row_end:
  438. row_end = row_max
  439. restart = table[row_max][0]
  440. # Gather these results up
  441. result.append(table.get_as_text(row_start, row_end))
  442. # Count them
  443. remaining -= row_end - row_start
  444. if restart:
  445. break
  446. if count:
  447. return matched
  448. return ("".join(result), restart)
  449. def stream_remove(self, path, start = None, end = None):
  450. """
  451. Remove data from the specified time interval within a stream.
  452. Removes all data in the interval [start, end), and intervals
  453. are truncated or split appropriately. Returns the number of
  454. data points removed.
  455. """
  456. stream_id = self._stream_id(path)
  457. table = self.data.getnode(path)
  458. intervals = self._get_intervals(stream_id)
  459. (start, end) = self._check_user_times(start, end)
  460. to_remove = Interval(start, end)
  461. removed = 0
  462. # Can't remove intervals from within the iterator, so we need to
  463. # remember what's currently in the intersection now.
  464. all_candidates = list(intervals.intersection(to_remove, orig = True))
  465. for (dbint, orig) in all_candidates:
  466. # Find row start and end
  467. row_start = self._find_start(table, dbint)
  468. row_end = self._find_end(table, dbint)
  469. # Adjust the DBInterval to match the newly found ends
  470. dbint.db_start = dbint.start
  471. dbint.db_end = dbint.end
  472. dbint.db_startpos = row_start
  473. dbint.db_endpos = row_end
  474. # Remove interval from the database
  475. self._remove_interval(stream_id, orig, dbint)
  476. # Remove data from the underlying table storage
  477. table.remove(row_start, row_end)
  478. # Count how many were removed
  479. removed += row_end - row_start
  480. return removed