1# Licensed to the Apache Software Foundation (ASF) under one or more
2# contributor license agreements.  See the NOTICE file distributed with
3# this work for additional information regarding copyright ownership.
4# The ASF licenses this file to You under the Apache License, Version 2.0
5# (the "License"); you may not use this file except in compliance with
6# the License.  You may obtain a copy of the License at
7#
8#     http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15
16"""
17A class which handles loading the pricing files.
18"""
19
20from __future__ import with_statement
21
22from typing import Dict
23from typing import Optional
24from typing import Union
25
26import os.path
27from os.path import join as pjoin
28
29try:
30    import simplejson as json
31    try:
32        JSONDecodeError = json.JSONDecodeError
33    except AttributeError:
34        # simplejson < 2.1.0 does not have the JSONDecodeError exception class
35        JSONDecodeError = ValueError  # type: ignore
36except ImportError:
37    import json  # type: ignore
38    JSONDecodeError = ValueError  # type: ignore
39
40__all__ = [
41    'get_pricing',
42    'get_size_price',
43    'set_pricing',
44    'clear_pricing_data',
45    'download_pricing_file'
46]
47
48# Default URL to the pricing file in a git repo
49DEFAULT_FILE_URL_GIT = 'https://git-wip-us.apache.org/repos/asf?p=libcloud.git;a=blob_plain;f=libcloud/data/pricing.json'  # NOQA
50
51DEFAULT_FILE_URL_S3_BUCKET = 'https://libcloud-pricing-data.s3.amazonaws.com/pricing.json'  # NOQA
52
53CURRENT_DIRECTORY = os.path.dirname(os.path.abspath(__file__))
54DEFAULT_PRICING_FILE_PATH = pjoin(CURRENT_DIRECTORY, 'data/pricing.json')
55CUSTOM_PRICING_FILE_PATH = os.path.expanduser('~/.libcloud/pricing.json')
56
57# Pricing data cache
58PRICING_DATA = {
59    'compute': {},
60    'storage': {}
61}  # type: Dict[str, Dict]
62
63VALID_PRICING_DRIVER_TYPES = ['compute', 'storage']
64
65# Set this to True to cache all the pricing data in memory instead of just the
66# one for the drivers which are used
67CACHE_ALL_PRICING_DATA = False
68
69
70def get_pricing_file_path(file_path=None):
71    # type: (Optional[str]) -> str
72    if os.path.exists(CUSTOM_PRICING_FILE_PATH) and \
73       os.path.isfile(CUSTOM_PRICING_FILE_PATH):
74        # Custom pricing file is available, use it
75        return CUSTOM_PRICING_FILE_PATH
76
77    return DEFAULT_PRICING_FILE_PATH
78
79
80def get_pricing(driver_type, driver_name, pricing_file_path=None,
81                cache_all=False):
82    # type: (str, str, Optional[str], bool) -> Optional[dict]
83    """
84    Return pricing for the provided driver.
85
86    NOTE: This method will also cache data for the requested driver
87    memory.
88
89    We intentionally only cache data for the requested driver and not all the
90    pricing data since the whole pricing data is quite large (~2 MB). This
91    way we avoid unncessary memory overhead.
92
93    :type driver_type: ``str``
94    :param driver_type: Driver type ('compute' or 'storage')
95
96    :type driver_name: ``str``
97    :param driver_name: Driver name
98
99    :type pricing_file_path: ``str``
100    :param pricing_file_path: Custom path to a price file. If not provided
101                              it uses a default path.
102
103    :type cache_all: ``bool``
104    :param cache_all: True to cache pricing data in memory for all the drivers
105                      and not just for the requested one.
106
107    :rtype: ``dict``
108    :return: Dictionary with pricing where a key name is size ID and
109             the value is a price.
110    """
111    cache_all = cache_all or CACHE_ALL_PRICING_DATA
112
113    if driver_type not in VALID_PRICING_DRIVER_TYPES:
114        raise AttributeError('Invalid driver type: %s', driver_type)
115
116    if driver_name in PRICING_DATA[driver_type]:
117        return PRICING_DATA[driver_type][driver_name]
118
119    if not pricing_file_path:
120        pricing_file_path = get_pricing_file_path(file_path=pricing_file_path)
121
122    with open(pricing_file_path, "r") as fp:
123        content = fp.read()
124
125    pricing_data = json.loads(content)
126    driver_pricing = pricing_data[driver_type][driver_name]
127
128    # NOTE: We only cache prices in memory for the the requested drivers.
129    # This way we avoid storing massive pricing data for all the drivers in
130    # memory
131
132    if cache_all:
133        for driver_type in VALID_PRICING_DRIVER_TYPES:
134            # pylint: disable=maybe-no-member
135            pricing = pricing_data.get(driver_type, None)
136
137            if not pricing:
138                continue
139
140            PRICING_DATA[driver_type] = pricing
141    else:
142        set_pricing(driver_type=driver_type, driver_name=driver_name,
143                    pricing=driver_pricing)
144
145    return driver_pricing
146
147
148def set_pricing(driver_type, driver_name, pricing):
149    # type: (str, str, dict) -> None
150    """
151    Populate the driver pricing dictionary.
152
153    :type driver_type: ``str``
154    :param driver_type: Driver type ('compute' or 'storage')
155
156    :type driver_name: ``str``
157    :param driver_name: Driver name
158
159    :type pricing: ``dict``
160    :param pricing: Dictionary where a key is a size ID and a value is a price.
161    """
162
163    PRICING_DATA[driver_type][driver_name] = pricing
164
165
166def get_size_price(driver_type, driver_name, size_id, region=None):
167    # type: (str, str, Union[str,int], Optional[str]) -> Optional[float]
168    """
169    Return price for the provided size.
170
171    :type driver_type: ``str``
172    :param driver_type: Driver type ('compute' or 'storage')
173
174    :type driver_name: ``str``
175    :param driver_name: Driver name
176
177    :type size_id: ``str`` or ``int``
178    :param size_id: Unique size ID (can be an integer or a string - depends on
179                    the driver)
180
181    :rtype: ``float``
182    :return: Size price.
183    """
184    pricing = get_pricing(driver_type=driver_type, driver_name=driver_name)
185    assert pricing is not None
186
187    price = None  # Type: Optional[float]
188
189    try:
190        if region is None:
191            price = float(pricing[size_id])
192        else:
193            price = float(pricing[size_id][region])
194    except KeyError:
195        # Price not available
196        price = None
197
198    return price
199
200
201def invalidate_pricing_cache():
202    # type: () -> None
203    """
204    Invalidate pricing cache for all the drivers.
205    """
206    PRICING_DATA['compute'] = {}
207    PRICING_DATA['storage'] = {}
208
209
210def clear_pricing_data():
211    # type: () -> None
212    """
213    Invalidate pricing cache for all the drivers.
214
215    Note: This method does the same thing as invalidate_pricing_cache and is
216    here for backward compatibility reasons.
217    """
218    invalidate_pricing_cache()
219
220
221def invalidate_module_pricing_cache(driver_type, driver_name):
222    # type: (str, str) -> None
223    """
224    Invalidate the cache for the specified driver.
225
226    :type driver_type: ``str``
227    :param driver_type: Driver type ('compute' or 'storage')
228
229    :type driver_name: ``str``
230    :param driver_name: Driver name
231    """
232    if driver_name in PRICING_DATA[driver_type]:
233        del PRICING_DATA[driver_type][driver_name]
234
235
236def download_pricing_file(file_url=DEFAULT_FILE_URL_S3_BUCKET,
237                          file_path=CUSTOM_PRICING_FILE_PATH):
238    # type: (str, str) -> None
239    """
240    Download pricing file from the file_url and save it to file_path.
241
242    :type file_url: ``str``
243    :param file_url: URL pointing to the pricing file.
244
245    :type file_path: ``str``
246    :param file_path: Path where a download pricing file will be saved.
247    """
248    from libcloud.utils.connection import get_response_object
249
250    dir_name = os.path.dirname(file_path)
251
252    if not os.path.exists(dir_name):
253        # Verify a valid path is provided
254        msg = ('Can\'t write to %s, directory %s, doesn\'t exist' %
255               (file_path, dir_name))
256        raise ValueError(msg)
257
258    if os.path.exists(file_path) and os.path.isdir(file_path):
259        msg = ('Can\'t write to %s file path because it\'s a'
260               ' directory' % (file_path))
261        raise ValueError(msg)
262
263    response = get_response_object(file_url)
264    body = response.body
265
266    # Verify pricing file is valid
267    try:
268        data = json.loads(body)
269    except JSONDecodeError:
270        msg = 'Provided URL doesn\'t contain valid pricing data'
271        raise Exception(msg)
272
273    # pylint: disable=maybe-no-member
274    if not data.get('updated', None):
275        msg = 'Provided URL doesn\'t contain valid pricing data'
276        raise Exception(msg)
277
278    # No need to stream it since file is small
279    with open(file_path, 'w') as file_handle:
280        file_handle.write(body)
281