OutputBackends.py

#
#	GammaScoutUtil - Tool to communicate with Gamma Scout Geiger counters.
#	Copyright (C) 2011-2011 Johannes Bauer
#
#	This file is part of GammaScoutUtil.
#
#	GammaScoutUtil is free software; you can redistribute it and/or modify
#	it under the terms of the GNU General Public License as published by
#	the Free Software Foundation; this program is ONLY licensed under
#	version 3 of the License, later versions are explicitly excluded.
#
#	GammaScoutUtil is distributed in the hope that it will be useful,
#	but WITHOUT ANY WARRANTY; without even the implied warranty of
#	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#	GNU General Public License for more details.
#
#	You should have received a copy of the GNU General Public License
#	along with GammaScoutUtil; if not, write to the Free Software
#	Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
#	Johannes Bauer <JohannesBauer@gmx.de>
#

import sys
import csv
import struct
import hashlib
import logging
import datetime
import os
try:
	import pymysql
except ImportError:
	pymysql = None

import Globals
from Exceptions import InvalidArgumentException
from DosisConversion import DosisConversion
from SQLite import SQLite

def _strftimeexpand(filename, args):
	if args["localstrftime"]:
		now = datetime.datetime.now()
	else:
		now = datetime.datetime.utcnow()
	filename = now.strftime(filename)
	directory = os.path.dirname(filename)
	try:
		os.makedirs(directory)
	except OSError:
		pass
	return filename

def _parseconnstr(connstr, acceptkeywords):
	"""Acceptkeywords is a dictionary of default items that are returned if the
	argument wasn't specified. Only keys in that dictionary are allowed to be
	specified in the connstr."""
	assert(isinstance(connstr, str))
	assert(isinstance(acceptkeywords, dict))
	resultdict = dict(acceptkeywords)
	for item in connstr.split(","):
		keyvalue = item.split("=")
		if len(keyvalue) != 2:
			raise InvalidArgumentException("Connection string must be of the form key1=value1,key2=value2,... but item '%s' contains an illegal number of values (%d instead of 2)." % (item, len(keyvalue)))
		(key, value) = keyvalue
		if key not in acceptkeywords:
			raise InvalidArgumentException("Connection string recognizes keywords %s, but keyword '%s' was given on command line." % (", ".join(sorted(list(acceptkeywords.keys()))), key))
		resultdict[key] = value
	return resultdict


class OutputBackend():
	def __init__(self, filename, args):
		self._filename = filename
		self._args = args

	def initdata(self, logsize, datablob):
		pass

	def newinterval(self, fromtime, totime, counts):
		pass

	def close(self):
		pass


class FileWrapper():
	def __init__(self, filename, args):
		self._filename = filename
		self._args = args
		if self._filename == "-":
			self._f = sys.stdout
		else:
			self._f = open(filename, "w", encoding = "utf-8")

	def write(self, data, **kwargs):
		try:
			result = self._f.write(data, **kwargs)
		except UnicodeEncodeError:
			data = data.replace("µ", "u")
			data = data.encode("ascii", errors = "replace").decode("ascii")
			result = self._f.write(data, **kwargs)
		if self._args["line_buffered"]:
			self._f.flush()
		return result

	def close(self):
		if self._filename != "-":
			self._f.close()


class OutputBackendBIN(OutputBackend):
	def __init__(self, filename, args):
		OutputBackend.__init__(self, filename, args)
		self._f = open(_strftimeexpand(filename, args), "wb")

	def initdata(self, logsize, datablob):
		data = ("Gamma scout binary data blob " + Globals.VERSION).encode("utf-8")
		data += bytes([ 0 ])
		data += struct.pack("<L", logsize)
		data += datablob
		self._f.write(data)

		# Append SHA-256 hash of data
		self._f.write(hashlib.sha256(data).digest())

	@staticmethod
	def readdata(filename, force):
		log = logging.getLogger("gsu.fileops." + __class__.__name__)

		data = open(filename, "rb").read()
		hashval = data[-32:]
		data = data[:-32]
		calchash = hashlib.sha256(data).digest()
		if calchash != hashval:
			if not force:
				raise InvalidArgumentException("File has incorrect hash value, will not read it without specification of --force")

		loc = data.index(bytes([0]))
		version = data[:loc].decode("utf-8")
		log.debug("Reading file data from file generated by %s" % (version))
		data = data[loc + 1:]

		(length,) = struct.unpack("<L", data[:4])
		data = data[4:]
		return (length, data)


	def close(self):
		self._f.close()


class OutputBackendCSV(OutputBackend):
	def __init__(self, filename, args):
		OutputBackend.__init__(self, filename, args)
		self._f = FileWrapper(_strftimeexpand(filename, args), args)
		self._csv = csv.writer(self._f)
		if not self._args["noheader"]:
			self._csv.writerow(["From", "To", "Counts", "Seconds", "CPM", "CPS", "µSv/h"])

	def newinterval(self, fromtime, totime, counts):
		delta = (totime - fromtime)
		totalseconds = delta.days * 86400 + delta.seconds
		cps = counts / totalseconds
		self._csv.writerow([ fromtime.strftime(self._args["date_format"]), totime.strftime(self._args["date_format"]), counts, totalseconds, 60 * cps, cps, DosisConversion.cts_per_sec_to_usv_per_hr(cps) ])

	def close(self):
		self._f.close()


class OutputBackendTXT(OutputBackend):
	def get_known_args():
		return [ "fromtime", "midtime", "totime", "counts", "intervallen", "cps", "cpm", "usvperhr" ]

	def __init__(self, filename, args):
		OutputBackend.__init__(self, filename, args)
		self._sampleno = 0
		self._f = FileWrapper(_strftimeexpand(filename, args), args)
		if (not self._args["noheader"]) and (not self._args["txt_format"]):
			heading = "%-20s   %-20s   %6s  %6s   %4s   %5s   %6s" % ("From", "To", "Counts", "Seconds", "CPM", "CPS", "µSv/hr")
			print(heading, file = self._f)
			print("-" * len(heading), file = self._f)

	def newinterval(self, fromtime, totime, counts):
		self._sampleno += 1
		delta = (totime - fromtime)
		totalseconds = round(delta.days * 86400 + delta.seconds)
		cps = counts / totalseconds

		midtime = fromtime + datetime.timedelta(0, totalseconds / 2)

		intervallenstr = {
			300:	"5 min",
			600:	"10 min",
		}.get(totalseconds, "?")

		variables = {
			"fromtime":			fromtime.strftime(self._args["date_format"]),
			"totime":			totime.strftime(self._args["date_format"]),
			"midtime":			midtime.strftime(self._args["date_format"]),
			"counts":			counts,
			"intervallen":		totalseconds,
			"cps":				cps,
			"cpm":				60 * cps,
			"usvperhr":			DosisConversion.cts_per_sec_to_usv_per_hr(cps),
			"sampleno":			self._sampleno,
			"intervallenstr":	intervallenstr,
		}
		if self._args["txt_format"] is None:
			if not self._args["gstool_txt_format"]:
				# Regular format
				print("%(fromtime)-20s   %(totime)-20s   %(counts)6d   %(intervallen)6d   %(cpm)6.1f   %(cps)7.3f   %(usvperhr)6.3f" % (variables), file = self._f)
			else:
				# GSTool format
				print("%(sampleno)04d    %(intervallenstr)-7s %(fromtime)-19s     %(totime)-19s     %(counts)010d      %(cps)-6.3f  %(usvperhr)5.3f" % (variables), file = self._f)
		else:
			# Custom format
			print(self._args["txt_format"] % (variables), file = self._f)

	def close(self):
		self._f.close()


class OutputBackendXML(OutputBackend):
	def __init__(self, filename, args):
		OutputBackend.__init__(self, filename, args)
		self._f = FileWrapper(_strftimeexpand(filename, args), args)
		if not self._args["noheader"]:
			print("<?xml version=\"1.0\" encoding=\"utf-8\" ?>", file = self._f)
			print("<gammascout>", file = self._f)

	def newinterval(self, fromtime, totime, counts):
		delta = (totime - fromtime)
		totalseconds = delta.days * 86400 + delta.seconds
		cps = counts / totalseconds
		doserate = DosisConversion.cts_per_sec_to_usv_per_hr(cps)
		print("	<interval from=\"%s\" to=\"%s\" secs=\"%s\" counts=\"%d\" doserate=\"%.5f\" />" % (fromtime.strftime(self._args["date_format"]), totime.strftime(self._args["date_format"]), totalseconds, counts, doserate), file = self._f)

	def close(self):
		print("</gammascout>", file = self._f)
		self._f.close()


class OutputBackendSqlite(OutputBackend):
	def __init__(self, filename, args):
		OutputBackend.__init__(self, filename, args)
		self._db = SQLite(_strftimeexpand(filename, args))
		self._db.exec_mayfail_commit("""CREATE TABLE metadata (
			key character varying PRIMARY KEY,
			value character varying NOT NULL
		);""")
		self._db.exec_mayfail_commit("INSERT INTO metadata (key, value) VALUES ('dbversion', '1');")
		self._db.exec_mayfail_commit("""CREATE TABLE data (
			id integer PRIMARY KEY,
			tfrom timestamp NOT NULL,
			tto timestamp NOT NULL,
			counts integer NOT NULL,
			CHECK(tto > tfrom),
			CHECK(counts >= 0)
		);""")

	def newinterval(self, fromtime, totime, counts):
		self._db.execute("INSERT INTO data (tfrom, tto, counts) VALUES (?, ?, ?);", fromtime, totime, counts)

	def close(self):
		self._db.commit()


class OutputBackendSQL(OutputBackend):
	def __init__(self, connstring, args):
		OutputBackend.__init__(self, connstring, args)
		self._dbdef = _parseconnstr(connstring, {
			"file":			"-",
			"dialect":		"sqlite",
			"dbname":		"gammascout",
			"tablename":	"data",
		})

		knowndialects = [ "sqlite", "mysql" ]
		if self._dbdef["dialect"] not in knowndialects:
			raise InvalidArgumentException("dialect must be one of %s, but %s was given." % (", ".join(sorted(list(knowndialects))), self._dbdef["dialect"]))

		self._f = FileWrapper(_strftimeexpand(self._dbdef["file"], args), args)

		if self._dbdef["dialect"] == "sqlite":
			self._createdb_sqlite()
			self.newinterval = self._newinterval_sqlite
			self.close = self._close_sqlite
		elif self._dbdef["dialect"] == "mysql":
			self._createdb_mysql()
			self.newinterval = self._newinterval_mysql
			self.close = self._close_mysql

	def _createdb_sqlite(self):
		print("BEGIN TRANSACTION;", file = self._f)
		print("""CREATE TABLE IF NOT EXISTS %(dbname)s (
			id integer PRIMARY KEY,
			tfrom timestamp NOT NULL,
			tto timestamp NOT NULL,
			counts integer NOT NULL,
			CHECK(tto > tfrom),
			CHECK(counts >= 0)
		);""" % self._dbdef, file = self._f)

	def _newinterval_sqlite(self, fromtime, totime, counts):
		print("INSERT INTO data (tfrom, tto, counts) VALUES ('%s', '%s', %d);" % (fromtime, totime, counts), file = self._f)

	def _close_sqlite(self):
		print("COMMIT;", file = self._f)

	def _createdb_mysql(self):
		print("CREATE DATABASE /*!32312 IF NOT EXISTS*/ `%(dbname)s` /*!40100 DEFAULT CHARACTER SET utf8 */;" % self._dbdef, file = self._f)
		print("CONNECT `%(dbname)s`;" % self._dbdef, file = self._f)
		print("START TRANSACTION;", file = self._f)
		print("""CREATE TABLE IF NOT EXISTS `%(tablename)s` (
			`id` int(12) NOT NULL AUTO_INCREMENT PRIMARY KEY,
			`tfrom` timestamp NOT NULL,
			`tto` timestamp NOT NULL,
			`counts` int(8) NOT NULL,
			CHECK(tto > tfrom),
			CHECK(counts >= 0)
		);""" % self._dbdef, file = self._f)

	def _newinterval_mysql(self, fromtime, totime, counts):
		print("INSERT INTO data (`tfrom`, `tto`, `counts`) VALUES ('%s', '%s', %d);" % (fromtime, totime, counts), file = self._f)

	def _close_mysql(self):
		print("COMMIT;", file = self._f)


class OutputBackendMySQL(OutputBackend):
	def __init__(self, connstring, args):
		OutputBackend.__init__(self, connstring, args)
		if pymysql is None:
			raise InvalidArgumentException("The MySQL backend is not available from Python. Please install the PyMySQL pacakge available from http://www.pymysql.org")
		self._dbdef = _parseconnstr(connstring, {
			"host":			None,
			"port":			None,
			"unixsocket":	None,
			"user":			None,
			"password":		None,
			"dbname":		"gammascout",
			"tablename":	"data",
		})
		if self._dbdef["port"] is not None:
			self._dbdef["port"] = int(self._dbdef["port"])

		connparams = {
			"host":			self._dbdef["host"],
			"user":			self._dbdef["user"],
			"passwd":		self._dbdef["password"],
			"port":			self._dbdef["port"],
			"unixsocket":	self._dbdef["unixsocket"],
		}
		for key in list(connparams.keys()):
			if connparams[key] is None:
				del connparams[key]
		self._db = pymysql.Connect(**connparams)
		self._cursor = self._db.cursor()
		self._cursor.execute("CREATE DATABASE /*!32312 IF NOT EXISTS*/ %s;" % (self._dbdef["dbname"]))
		self._cursor.execute("USE %s;" % (self._dbdef["dbname"]))
		self._cursor.execute("""CREATE TABLE IF NOT EXISTS metadata (
			`key` varchar(128) PRIMARY KEY,
			`value` varchar(128) NOT NULL
		);""")
		self._exec_mayfail("INSERT INTO metadata (`key`, `value`) VALUES ('dbversion', '1');")

		self._cursor.execute("""CREATE TABLE IF NOT EXISTS `%(tablename)s` (
			`id` int(12) NOT NULL AUTO_INCREMENT PRIMARY KEY,
			`tfrom` timestamp NOT NULL,
			`tto` timestamp NOT NULL,
			`counts` int(8) NOT NULL,
			CHECK(tto > tfrom),
			CHECK(counts >= 0)
		);""" % self._dbdef)

	def _exec_mayfail(self, sql):
		try:
			self._cursor.execute(sql)
		except pymysql.err.IntegrityError:
			pass

	def newinterval(self, fromtime, totime, counts):
		self._cursor.execute("INSERT INTO data (`tfrom`, `tto`, `counts`) VALUES ('%s', '%s', %d);" % (fromtime, totime, counts))

	def close(self):
		self._db.commit()


def getbackendbyname(name):
	return {
		"bin":		OutputBackendBIN,
		"txt":		OutputBackendTXT,
		"csv":		OutputBackendCSV,
		"sqlite":	OutputBackendSqlite,
		"xml":		OutputBackendXML,
		"sql":		OutputBackendSQL,
		"mysql":	OutputBackendMySQL,
	}[name]