1"""
2Utilities for working with memcache
3
4:depends:  - python-memcached
5
6This library sets up a connection object for memcache, using the configuration
7passed into the get_conn() function. Normally, this is __opts__. Optionally,
8a profile or specific host and port may be passed in. If neither profile nor
9host and port are provided, the defaults of '`127.0.0.`` and ``11211`` are
10used. The following configurations are both valid:
11
12.. code-block:: yaml
13
14    # No profile name
15    memcached.host: 127.0.0.1
16    memcached.port: 11211
17
18    # One or more profiles defined
19    my_memcached_config:
20      memcached.host: 127.0.0.1
21      memcached.port: 11211
22
23Once configured, the get_conn() function is passed a set of opts, and,
24optionally, the name of a profile to be used.
25
26.. code-block:: python
27
28    import salt.utils.memcached_utils.py
29    conn = salt.utils.memcached_utils.get_conn(__opts__,
30                                              profile='my_memcached_config')
31
32It should be noted that some usages of memcached may require a profile to be
33specified, rather than top-level configurations. This being the case, it is
34better to always use a named configuration profile, as shown above.
35"""
36
37
38import logging
39
40from salt.exceptions import CommandExecutionError, SaltInvocationError
41
42try:
43    import memcache
44
45    HAS_LIBS = True
46except ImportError:
47    HAS_LIBS = False
48
49DEFAULT_HOST = "127.0.0.1"
50DEFAULT_PORT = 11211
51DEFAULT_TIME = 0
52DEFAULT_MIN_COMPRESS_LEN = 0
53
54# Set up logging
55log = logging.getLogger(__name__)
56
57# Don't shadow built-ins
58__func_alias__ = {"set_": "set"}
59
60
61# Although utils are often directly imported, it is also possible
62# to use the loader.
63def __virtual__():
64    """
65    Only load if python-memcached is installed
66    """
67    return True if HAS_LIBS else False
68
69
70def get_conn(opts, profile=None, host=None, port=None):
71    """
72    Return a conn object for accessing memcached
73    """
74    if not (host and port):
75        opts_pillar = opts.get("pillar", {})
76        opts_master = opts_pillar.get("master", {})
77
78        opts_merged = {}
79        opts_merged.update(opts_master)
80        opts_merged.update(opts_pillar)
81        opts_merged.update(opts)
82
83        if profile:
84            conf = opts_merged.get(profile, {})
85        else:
86            conf = opts_merged
87
88        host = conf.get("memcached.host", DEFAULT_HOST)
89        port = conf.get("memcached.port", DEFAULT_PORT)
90
91    if not str(port).isdigit():
92        raise SaltInvocationError("port must be an integer")
93
94    if HAS_LIBS:
95        return memcache.Client(["{}:{}".format(host, port)])
96    else:
97        raise CommandExecutionError(
98            "(unable to import memcache, module most likely not installed)"
99        )
100
101
102def _check_stats(conn):
103    """
104    Helper function to check the stats data passed into it, and raise an
105    exception if none are returned. Otherwise, the stats are returned.
106    """
107    stats = conn.get_stats()
108    if not stats:
109        raise CommandExecutionError("memcached server is down or does not exist")
110    return stats
111
112
113def set_(
114    conn, key, value, time=DEFAULT_TIME, min_compress_len=DEFAULT_MIN_COMPRESS_LEN
115):
116    """
117    Set a key on the memcached server, overwriting the value if it exists.
118    """
119    if not isinstance(time, int):
120        raise SaltInvocationError("'time' must be an integer")
121    if not isinstance(min_compress_len, int):
122        raise SaltInvocationError("'min_compress_len' must be an integer")
123    _check_stats(conn)
124    return conn.set(key, value, time, min_compress_len)
125
126
127def get(conn, key):
128    """
129    Retrieve value for a key
130    """
131    _check_stats(conn)
132    return conn.get(key)
133