/utils/db.py

import os
import re
import sys
import couchdb
import simplejson
from os.path import dirname, basename
from urlparse import urlsplit, urlunsplit

from dotcloud.utils import shellCommand

class ViewDict(dict):
    """ Results of a CouchDB view mapped to a dict """

    duplicates = property(lambda self: self._duplicates)

    def __init__(self, view):
        dict.__init__(self)
        self._duplicates = []
        for (key, value) in ((r.key, r.value) for r in view.rows):
            if key not in self:
                dict.__setitem__(self, key, value)
            else:
                self._duplicates.append((key, value))

    def __setitem__(self, key, value):
        raise AttributeError()

    def __delitem__(self, key):
        raise AttributeError()


class TypedDocs(ViewDict):

    def __init__(self, typename, db):
        self._db = db
        self._typename = typename
        super(self.__class__, self).__init__(self._view())
        # FIXME: delete duplicates

    def sync(self, src, pretend=False):
        for (key, doc) in self.items():
            if key not in src:
                if not pretend:
                    del self._db[doc["_id"]]
        for (key, src_doc) in src.items():
            if key in self:
                doc = self[key]
                doc.update(src_doc)
                if not pretend:
                    self._db[doc["_id"]] = doc
            else:
                doc = dict(src_doc, type=self._typename)
                if not pretend:
                    self._db.create(doc)

    def _view(self):
        return self._db.view(
                "types/%s" % self._typename,
                js_map="""if (doc.type == "%s") { emit (doc.name ? doc.name : doc._id, doc); }""" % self._typename
            )

class DB(couchdb.Database):
    """ A customized CouchDB database """

    def __init__(self, url, debug=False):
        self._debug = debug
        (url_scheme, url_host, url_path, url_query, url_fragment) = urlsplit(url)
        server = couchdb.Server(urlunsplit((url_scheme, url_host, dirname(url_path), url_query, url_fragment)))
        name = basename(url_path)
        couchdb.Database.__init__(
                self,
                couchdb.client.uri(server.resource.uri, name),
                couchdb.client.validate_dbname(name),
                http=server.resource.http)
    
    def typed_docs(self, typename):
        return TypedDocs(typename, db=self)

    def view(self, name, js_map=None, js_reduce=None, *args, **kw):
        if js_map:
            self.log("UPDATING VIEW %s" % name)
            self.set_view(name, js_map=js_map, js_reduce=js_reduce)
        return couchdb.Database.view(self, name=name, group=True, *args, **kw)

    def log(self, msg, **args):
        if self._debug:
            sys.stderr.write("[%s] %s%s\n" % (self.resource.uri, msg, args and " "+args or ""))

    def __delitem__(self, key):
        self.log("DELETING %s" % (key))
        return couchdb.Database.__delitem__(self, key)

    def __setitem__(self, key, value):
        self.log("UPDATING %s" % (key))
        return couchdb.Database.__setitem__(self, key, value)

    def create(self, value):
        result = couchdb.Database.create(self, value)
        self.log("CREATED %s" % result)
        return result

    def get_view(self, path):
        (doc_name, view_name) = os.path.split(path)
        return self["_design/%s" % doc_name]["views"][view_name]

    def set_view(self, path, js_map, js_reduce=None):
        (doc_name, view_name) = os.path.split(path)
        try:
            doc = self["_design/%s" % doc_name]
        except:
            doc = {"language" : "javascript", "views" : {}}
        doc["views"][view_name] = {
                "map"       : "function(doc) { %s }" % js_map,
                "reduce"    : "function(key, values, rereduce) { %s }" % js_reduce if js_reduce else None
            }
        self["_design/%s" % doc_name] = doc
        self.log("Uploaded view %s" % path)
        return path

    def delete_view(self, path):
        (doc_name, view_name) = os.path.split(path)
        full_doc_name = "_design/%s" % doc_name
        if full_doc_name:
            del self[full_doc_name]
            self.log("Deleted view %s" % path)

class DotCloudDB(DB):
    """ Code for the dotCloud runtime mapped to a CouchDB database """

    # DOWNLOAD
    ###########

    def get_config(self):
        self.typed_docs("config").sync({"config" : {"type" : "config", "name" : "config", "content" : Config()}})
        return ViewDict(self.view("dotcloud/config"))
    config = property(get_config)

    def get_instances(self):
        self.typed_docs("ovz_container").sync(OVZContainers())
        return ViewDict(self.view("dotcloud/instances"))
    instances = property(get_instances)

    def get_volumes(self):
        self.typed_docs("volume").sync(Volumes(self.config["path"]))
        return ViewDict(self.view("dotcloud/volumes"))
    volumes = property(get_volumes)

    def view(self, name, *args, **kw):
        if name in self.views:
            kw["js_map"] = self.views[name]["map"]
            kw["js_reduce"] = self.views[name].get("reduce")
        return super(self.__class__, self).view(name, *args, **kw)

    views = {
        "dotcloud/volumes"   : {
            "map":
            """
            if (doc.type == "volume") { emit (doc.name, doc); }
            """
        },
        "dotcloud/instances" : {
            "map":
                """
                if (doc.type == "ovz_container") {
                    emit (doc.name, doc);
                }
                else if (doc.type == "instance") {
                    emit (doc.name, doc);
                }
                """,
            "reduce":
                """
                var result = {"managed": false};
                for (var idx in values) {
                    var object = values[idx];
                    if (object["type"] == "instance") { result["managed"] = true; }
                    for (var key in object) {
                        if (["_id", "_rev", "type"].indexOf(key) != -1) continue;
                        result[key] = object[key];
                    }
                }
                return result;
                """
        },
        "dotcloud/config" : {
            "map":
                """
                 if (doc.type == "config" && doc.content) {
                    for (var key in doc.content) {
                       emit (key, doc.content[key]);
                    }
                 }
                """,
            "reduce":
                """
                return values[0];
                """
        }
    }


# System classes
################

class Config(dict):

    def __init__(self, **override):
        super(self.__class__, self).__init__()
        config_search = os.environ.get("DOTCLOUD_CONFIG_PATH", "/etc/dotcloud").split(":")
        for (path, content) in self._load_files(*config_search):
            for (key, value) in content.items():
                self[key] = value
        for (key, value) in override.items():
            self[key] = value

    def _load_files(self, *search):

        def config_paths():
            for config_path in map(os.path.abspath, search):
                if os.path.isdir(config_path):
                    for (path, dirnames, filenames) in os.walk(config_path):
                        for filename in filenames:
                            yield os.path.normpath(os.path.join(path, filename))
                else:
                    yield os.path.normpath(config_path)
        
        for path in config_paths():
            try:
                yield (path, simplejson.loads(file(path).read()))
            except Exception, e:
                sys.stderr.write("Skipping %s: %s\n" % (path, e))
                continue

class IPAddresses(dict):
    
    _re = re.compile("^[^:]*: ([^ :]+).*inet ([^ ]*) .*$")

    def __init__(self):
        dict.__init__(self)
        lines = shellCommand("ip -o addr").strip().split("\n")
        for line_match in filter(None, map(self._re.match, lines)):
            doc = {
                "iface"     : line_match.groups()[0],
                "address"   : line_match.groups()[1]
                }
            self[doc["address"]] = doc

class Volumes(dict):

    def __init__(self, dotcloud_path):
        dict.__init__(self)
        volumes_path = os.path.join(dotcloud_path, "volumes")
        for dirname in os.listdir(volumes_path):
            path = os.path.normpath(os.path.join(volumes_path, dirname))
            self[dirname] = {"name": dirname, "path":path}

class AUFSMounts(dict):

    re_mount = re.compile("^aufs on ([^ ]*) type aufs \((.*)\)$")

    def __init__(self):
        lines = shellCommand("mount").strip().split("\n")
        for line_match in filter(None, map(cls.re_mount.match, lines)):
            path = os.path.normpath(line_match.groups()[0])
            self[path] = {"name" : path, "options" : line_match.groups()[1].split(",")}


class BindFSMounts(dict):

    re_mount    = re.compile("^bindfs on ([^ ]*) type fuse\.bindfs \((.*)\)$")

    def __init__(self):
        lines = shellCommand("mount").strip().split("\n")
        for line_match in filter(None, map(cls.re_mount.match, lines)):
            path = os.path.normpath(line_match.groups()[0])
            self[path] = {"name" : path, "option" : line_match.groups()[1].split(",")}

class OVZContainers(dict):

    def __init__(self):
        dict.__init__(self)
        for line in shellCommand("vzlist -a -o name,veid,ip,status -H").strip().split("\n"):
            if not line:
                continue
            container = dict(zip(
                ("name", "veid", "ip", "running"),
                line.split()
                ))
            container["running"] = (True if container["running"] == "running" else False)
            if (container["ip"] == "-"):
                container["ip"] = None
            if (container["name"] == "-"):
                container["name"] = container["veid"]
            else:
                container["name"] = cls.get_long_name(container["name"])
            self[container["name"]] = container

    re_short_name = re.compile("^(........)(....)(....)(....)(............)$")
    @classmethod
    def get_long_name(cls, short_name):
        return "-".join(cls.re_short_name.match(short_name).groups())

    @classmethod
    def get_short_name(cls, long_name):
        return re.sub("-", "", long_name)