Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 7e07ff8

Browse filesBrowse files
author
Ace Nassri
authored
GCF: Add samples for billing (GoogleCloudPlatform#1725)
* Add billing GCF samples Change-Id: Ie43de503c9600f23b7bb168e62948a01f868457d * Fix spacing Change-Id: Iaa73d0548b5196a9e002b330165fd09c172928da * Add token comment Change-Id: I537817dda32abb03629e8ebc4c9e45d69efb59f7
1 parent 23dfa12 commit 7e07ff8
Copy full SHA for 7e07ff8

File tree

Expand file treeCollapse file tree

3 files changed

+288
-0
lines changed
Filter options
Expand file treeCollapse file tree

3 files changed

+288
-0
lines changed

‎functions/billing/main.py

Copy file name to clipboard
+169Lines changed: 169 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,169 @@
1+
# Copyright 2018 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the 'License');
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an 'AS IS' BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
# [START functions_billing_limit]
16+
# [START functions_billing_stop]
17+
import base64
18+
import json
19+
# [END functions_billing_stop]
20+
import os
21+
# [END functions_billing_limit]
22+
23+
# [START functions_billing_limit]
24+
# [START functions_billing_stop]
25+
from googleapiclient import discovery
26+
from oauth2client.client import GoogleCredentials
27+
28+
# [END functions_billing_stop]
29+
# [END functions_billing_limit]
30+
31+
# [START functions_billing_slack]
32+
from slackclient import SlackClient
33+
# [END functions_billing_slack]
34+
35+
# [START functions_billing_limit]
36+
# [START functions_billing_stop]
37+
PROJECT_ID = os.getenv('GCP_PROJECT')
38+
PROJECT_NAME = f'projects/{PROJECT_ID}'
39+
# [END functions_billing_stop]
40+
# [END functions_billing_limit]
41+
42+
# [START functions_billing_slack]
43+
44+
# See https://api.slack.com/docs/token-types#bot for more info
45+
BOT_ACCESS_TOKEN = 'xxxx-111111111111-abcdefghidklmnopq'
46+
47+
CHANNEL = 'general'
48+
49+
slack_client = SlackClient(BOT_ACCESS_TOKEN)
50+
51+
52+
def notify_slack(data, context):
53+
pubsub_message = data
54+
55+
notification_attrs = json.dumps(pubsub_message['attributes'])
56+
notification_data = base64.b64decode(data['data']).decode('utf-8')
57+
budget_notification_text = f'{notification_attrs}, {notification_data}'
58+
59+
res = slack_client.api_call(
60+
'chat.postMessage',
61+
channel=CHANNEL,
62+
text=budget_notification_text)
63+
print(res)
64+
# [END functions_billing_slack]
65+
66+
67+
# [START functions_billing_limit]
68+
def stop_billing(data, context):
69+
pubsub_data = base64.b64decode(data['data']).decode('utf-8')
70+
pubsub_json = json.loads(pubsub_data)
71+
cost_amount = pubsub_json['costAmount']
72+
budget_amount = pubsub_json['budgetAmount']
73+
if cost_amount <= budget_amount:
74+
print(f'No action necessary. (Current cost: {cost_amount})')
75+
return
76+
77+
billing = discovery.build(
78+
'cloudbilling',
79+
'v1',
80+
cache_discovery=False,
81+
credentials=GoogleCredentials.get_application_default()
82+
)
83+
84+
projects = billing.projects()
85+
86+
if __is_billing_enabled(PROJECT_NAME, projects):
87+
print(__disable_billing_for_project(PROJECT_NAME, projects))
88+
else:
89+
print('Billing already disabled')
90+
91+
92+
def __is_billing_enabled(project_name, projects):
93+
"""
94+
Determine whether billing is enabled for a project
95+
@param {string} project_name Name of project to check if billing is enabled
96+
@return {bool} Whether project has billing enabled or not
97+
"""
98+
res = projects.getBillingInfo(name=project_name).execute()
99+
return res['billingEnabled']
100+
101+
102+
def __disable_billing_for_project(project_name, projects):
103+
"""
104+
Disable billing for a project by removing its billing account
105+
@param {string} project_name Name of project disable billing on
106+
@return {string} Text containing response from disabling billing
107+
"""
108+
body = {'billingAccountName': ''} # Disable billing
109+
res = projects.updateBillingInfo(name=project_name, body=body).execute()
110+
print(f'Billing disabled: {json.dumps(res)}')
111+
# [END functions_billing_stop]
112+
113+
114+
# [START functions_billing_limit]
115+
ZONE = 'us-west1-b'
116+
117+
118+
def limit_use(data, context):
119+
pubsub_data = base64.b64decode(data['data']).decode('utf-8')
120+
pubsub_json = json.loads(pubsub_data)
121+
cost_amount = pubsub_json['costAmount']
122+
budget_amount = pubsub_json['budgetAmount']
123+
if cost_amount <= budget_amount:
124+
print(f'No action necessary. (Current cost: {cost_amount})')
125+
return
126+
127+
compute = discovery.build(
128+
'compute',
129+
'v1',
130+
cache_discovery=False,
131+
credentials=GoogleCredentials.get_application_default()
132+
)
133+
instances = compute.instances()
134+
135+
instance_names = __list_running_instances(PROJECT_ID, ZONE, instances)
136+
__stop_instances(PROJECT_ID, ZONE, instance_names, instances)
137+
138+
139+
def __list_running_instances(project_id, zone, instances):
140+
"""
141+
@param {string} project_id ID of project that contains instances to stop
142+
@param {string} zone Zone that contains instances to stop
143+
@return {Promise} Array of names of running instances
144+
"""
145+
res = instances.list(project=project_id, zone=zone).execute()
146+
147+
items = res['items']
148+
running_names = [i['name'] for i in items if i['status'] == 'RUNNING']
149+
return running_names
150+
151+
152+
def __stop_instances(project_id, zone, instance_names, instances):
153+
"""
154+
@param {string} project_id ID of project that contains instances to stop
155+
@param {string} zone Zone that contains instances to stop
156+
@param {Array} instance_names Names of instance to stop
157+
@return {Promise} Response from stopping instances
158+
"""
159+
if not len(instance_names):
160+
print('No running instances were found.')
161+
return
162+
163+
for name in instance_names:
164+
instances.stop(
165+
project=project_id,
166+
zone=zone,
167+
instance=name).execute()
168+
print(f'Instance stopped successfully: {name}')
169+
# [END functions_billing_limit]

‎functions/billing/main_test.py

Copy file name to clipboard
+116Lines changed: 116 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,116 @@
1+
# Copyright 2018, Google, LLC.
2+
# Licensed under the Apache License, Version 2.0 (the "License");
3+
# you may not use this file except in compliance with the License.
4+
# You may obtain a copy of the License at
5+
#
6+
# http://www.apache.org/licenses/LICENSE-2.0
7+
#
8+
# Unless required by applicable law or agreed to in writing, software
9+
# distributed under the License is distributed on an "AS IS" BASIS,
10+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
# See the License for the specific language governing permissions and
12+
# limitations under the License.
13+
14+
import base64
15+
import json
16+
17+
from mock import MagicMock, patch
18+
19+
import main
20+
21+
22+
@patch('main.slack_client')
23+
def test_notify_slack(slack_client):
24+
slack_client.api_call = MagicMock()
25+
26+
data = {"budgetAmount": 400, "costAmount": 500}
27+
attrs = {"foo": "bar"}
28+
29+
pubsub_message = {
30+
"data": base64.b64encode(bytes(json.dumps(data), 'utf-8')),
31+
"attributes": attrs
32+
}
33+
34+
main.notify_slack(pubsub_message, None)
35+
36+
assert slack_client.api_call.called
37+
38+
39+
@patch('main.PROJECT_ID')
40+
@patch('main.discovery')
41+
def test_disable_billing(discovery_mock, PROJECT_ID):
42+
PROJECT_ID = 'my-project'
43+
PROJECT_NAME = f'projects/{PROJECT_ID}'
44+
45+
data = {"budgetAmount": 400, "costAmount": 500}
46+
47+
pubsub_message = {
48+
"data": base64.b64encode(bytes(json.dumps(data), 'utf-8')),
49+
"attributes": {}
50+
}
51+
52+
projects_mock = MagicMock()
53+
projects_mock.projects = MagicMock(return_value=projects_mock)
54+
projects_mock.getBillingInfo = MagicMock(return_value=projects_mock)
55+
projects_mock.updateBillingInfo = MagicMock(return_value=projects_mock)
56+
projects_mock.execute = MagicMock(return_value={'billingEnabled': True})
57+
58+
discovery_mock.build = MagicMock(return_value=projects_mock)
59+
60+
main.stop_billing(pubsub_message, None)
61+
62+
assert projects_mock.getBillingInfo.called_with(name=PROJECT_NAME)
63+
assert projects_mock.updateBillingInfo.called_with(
64+
name=PROJECT_NAME,
65+
body={'billingAccountName': ''}
66+
)
67+
assert projects_mock.execute.call_count == 2
68+
69+
70+
@patch('main.PROJECT_ID')
71+
@patch('main.ZONE')
72+
@patch('main.discovery')
73+
def test_limit_use(discovery_mock, ZONE, PROJECT_ID):
74+
PROJECT_ID = 'my-project'
75+
PROJECT_NAME = f'projects/{PROJECT_ID}'
76+
ZONE = 'my-zone'
77+
78+
data = {"budgetAmount": 400, "costAmount": 500}
79+
80+
pubsub_message = {
81+
"data": base64.b64encode(bytes(json.dumps(data), 'utf-8')),
82+
"attributes": {}
83+
}
84+
85+
instances_list = {
86+
"items": [
87+
{"name": "instance-1", "status": "RUNNING"},
88+
{"name": "instance-2", "status": "TERMINATED"}
89+
]
90+
}
91+
92+
instances_mock = MagicMock()
93+
instances_mock.instances = MagicMock(return_value=instances_mock)
94+
instances_mock.list = MagicMock(return_value=instances_mock)
95+
instances_mock.stop = MagicMock(return_value=instances_mock)
96+
instances_mock.execute = MagicMock(return_value=instances_list)
97+
98+
projects_mock = MagicMock()
99+
projects_mock.projects = MagicMock(return_value=projects_mock)
100+
projects_mock.getBillingInfo = MagicMock(return_value=projects_mock)
101+
projects_mock.execute = MagicMock(return_value={'billingEnabled': True})
102+
103+
def discovery_mocker(x, *args, **kwargs):
104+
if x == 'compute':
105+
return instances_mock
106+
else:
107+
return projects_mock
108+
109+
discovery_mock.build = MagicMock(side_effect=discovery_mocker)
110+
111+
main.limit_use(pubsub_message, None)
112+
113+
assert projects_mock.getBillingInfo.called_with(name=PROJECT_NAME)
114+
assert instances_mock.list.calledWith(project=PROJECT_ID, zone=ZONE)
115+
assert instances_mock.stop.call_count == 1
116+
assert instances_mock.execute.call_count == 2

‎functions/billing/requirements.txt

Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
slackclient==1.3.0
2+
oauth2client==4.1.3
3+
google-api-python-client==1.7.4

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.