summaryrefslogtreecommitdiffstats
path: root/azure/aria/aria-extension-cloudify/src/aria/aria/storage/core.py
blob: 7e9b20123cf7908172d42c16de9202b810ea6fb8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Storage API management.
"""

import copy
from contextlib import contextmanager

from aria.logger import LoggerMixin
from . import sql_mapi

__all__ = (
    'Storage',
    'ModelStorage',
    'ResourceStorage'
)


class Storage(LoggerMixin):
    """
    Base class for storage managers.
    """
    def __init__(self,
                 api_cls,
                 api_kwargs=None,
                 items=(),
                 initiator=None,
                 initiator_kwargs=None,
                 **kwargs):
        """
        :param api_cls: API class for each entry
        :param api_kwargs:
        :param items: items to register
        :param initiator: function which initializes the storage before the first use; this function
         should return a dict, this dict would be passed in addition to the API kwargs; this enables
         the creation of non-serializable objects
        :param initiator_kwargs:
        :param kwargs:
        """
        super(Storage, self).__init__(**kwargs)
        self.api = api_cls
        self.registered = {}
        self._initiator = initiator
        self._initiator_kwargs = initiator_kwargs or {}
        self._api_kwargs = api_kwargs or {}
        self._additional_api_kwargs = {}
        if self._initiator:
            self._additional_api_kwargs = self._initiator(**self._initiator_kwargs)
        for item in items:
            self.register(item)
        self.logger.debug('{name} object is ready: {0!r}'.format(
            self, name=self.__class__.__name__))

    @property
    def _all_api_kwargs(self):
        kwargs = self._api_kwargs.copy()
        kwargs.update(self._additional_api_kwargs)
        return kwargs

    def __repr__(self):
        return '{name}(api={self.api})'.format(name=self.__class__.__name__, self=self)

    def __getattr__(self, item):
        try:
            return self.registered[item]
        except KeyError:
            return super(Storage, self).__getattribute__(item)

    @property
    def serialization_dict(self):
        return {
            'api': self.api,
            'api_kwargs': self._api_kwargs,
            'initiator': self._initiator,
            'initiator_kwargs': self._initiator_kwargs
        }

    def register(self, entry):
        """
        Register an API.

        :param entry:
        """
        raise NotImplementedError('Subclass must implement abstract register method')


class ResourceStorage(Storage):
    """
    Manages storage resource APIs ("RAPIs").
    """
    def register(self, name):
        """
        Register a storage resource API ("RAPI").

        :param name: name
        """
        self.registered[name] = self.api(name=name, **self._all_api_kwargs)
        self.registered[name].create()
        self.logger.debug('setup {name} in storage {self!r}'.format(name=name, self=self))


class ModelStorage(Storage):
    """
    Manages storage model APIs ("MAPIs").
    """
    def __init__(self, *args, **kwargs):
        if kwargs.get('initiator', None) is None:
            kwargs['initiator'] = sql_mapi.init_storage
        super(ModelStorage, self).__init__(*args, **kwargs)

    def register(self, model_cls):
        """
        Register a storage model API ("MAPI").

        :param model_cls: model API to register
        """
        model_name = model_cls.__modelname__
        if model_name in self.registered:
            self.logger.debug('{name} already in storage {self!r}'.format(name=model_name,
                                                                          self=self))
            return
        self.registered[model_name] = self.api(name=model_name,
                                               model_cls=model_cls,
                                               **self._all_api_kwargs)
        self.registered[model_name].create()
        self.logger.debug('setup {name} in storage {self!r}'.format(name=model_name, self=self))

    def drop(self):
        """
        Drop all the tables.
        """
        for mapi in self.registered.itervalues():
            mapi.drop()

    @contextmanager
    def instrument(self, *instrumentation):
        original_instrumentation = {}

        try:
            for mapi in self.registered.itervalues():
                original_instrumentation[mapi] = copy.copy(mapi._instrumentation)
                mapi._instrumentation.extend(instrumentation)
            yield self
        finally:
            for mapi in self.registered.itervalues():
                mapi._instrumentation[:] = original_instrumentation[mapi]