Ejemplo n.º 1
0
class VedisBackend(DBPlugin):
    def connect(self):
        path = os.path.join(self.project.outdir, DB_NAME)
        self._db = Vedis(path)

    def get_sample(self, sha256):
        s = self._samples[sha256]
        j = json.loads(s)
        return Sample(self.project, j['file'], j['tags'], j['mutations'],
                      j.get('base_sample', None))

    def put_sample(self, sample):
        self._samples[sample.sha256] = sample.to_json()

        # keep reference of the sample for each tag :)
        for t in sample.tags:
            self.put_tag(t, sample.sha256)

        self._db.commit()

    def put_scan(self, scan, sample=None, sha256=None):
        if sha256:
            h = sha256
        else:
            h = sample.sha256
        self._scans[h] = json.dumps(scan)
        self._db.commit()

    def get_scan(self, sample=None, sha256=None):
        if sha256:
            h = sha256
        else:
            h = sample.sha256
        res = self._scans[h]
        if res:
            return json.loads(res)
        return None

    def put_tag(self, tag, sha256):
        t = self._db.Set('tag_' + tag)
        t.add(sha256)

    def get_tagged_samples(self, tag):
        return self._db.Set('tag_' + tag)

    @property
    def _samples(self):
        return self._db.Hash('samples')

    @property
    def _scans(self):
        return self._db.Hash('scans')

    @property
    def all_samples(self):
        # no way to use a generator with vedis
        # still ok for small analyses, we'll need to
        # support other DBs to reduce memory usage
        for s in self._samples.keys():
            yield self.get_sample(s)
Ejemplo n.º 2
0
class Base(object):
    __slots__ = ["db"]

    def __init__(self, **kwargs):
        self.db = Vedis(kwargs["dataFile"])

    def get(self, user_id):
        try:
            data = self.db[user_id]

            return ast.literal_eval(data)
        except KeyError:
            return None

    def set(self, user_id, shop, value):
        try:
            data = self.get(user_id)

            if not data:
                data = dict()

            data[shop] = value

            self.db[user_id] = data
            self.db.commit()

            return True
        except:
            return False