diff --git a/.kokoro/functions/billing.cfg b/.kokoro/functions/billing.cfg new file mode 100644 index 0000000000..9668c9d2cb --- /dev/null +++ b/.kokoro/functions/billing.cfg @@ -0,0 +1,13 @@ +# Format: //devtools/kokoro/config/proto/build.proto + +# Set the folder in which the tests are run +env_vars: { + key: "PROJECT" + value: "functions/billing" +} + +# Tell the trampoline which build file to use. +env_vars: { + key: "TRAMPOLINE_BUILD_FILE" + value: "github/nodejs-docs-samples/.kokoro/build.sh" +} diff --git a/functions/billing/.gcloudignore b/functions/billing/.gcloudignore new file mode 100644 index 0000000000..ccc4eb240e --- /dev/null +++ b/functions/billing/.gcloudignore @@ -0,0 +1,16 @@ +# This file specifies files that are *not* uploaded to Google Cloud Platform +# using gcloud. It follows the same syntax as .gitignore, with the addition of +# "#!include" directives (which insert the entries of the given .gitignore-style +# file at that point). +# +# For more information, run: +# $ gcloud topic gcloudignore +# +.gcloudignore +# If you would like to upload your .git directory, .gitignore file or files +# from your .gitignore file, remove the corresponding line +# below: +.git +.gitignore + +node_modules diff --git a/functions/billing/index.js b/functions/billing/index.js new file mode 100644 index 0000000000..ffe433516d --- /dev/null +++ b/functions/billing/index.js @@ -0,0 +1,156 @@ +/** + * Copyright 2018, Google LLC. + * Licensed 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. + */ + +// [START functions_billing_limit] +// [START functions_billing_stop] +const {google} = require('googleapis'); +const {auth} = require('google-auth-library'); + +const PROJECT_ID = process.env.GCP_PROJECT; +const PROJECT_NAME = `projects/${PROJECT_ID}`; +// [END functions_billing_stop] +// [END functions_billing_limit] + +// [START functions_billing_slack] +const slack = require('slack'); + +const BOT_ACCESS_TOKEN = 'xxxx-111111111111-abcdefghidklmnopq'; +const CHANNEL = 'general'; + +exports.notifySlack = async (data, context) => { + const pubsubMessage = data; + const pubsubAttrs = JSON.stringify(pubsubMessage.attributes); + const pubsubData = Buffer.from(pubsubMessage.data, 'base64').toString(); + const budgetNotificationText = `${pubsubAttrs}, ${pubsubData}`; + + const res = await slack.chat.postMessage({ + token: BOT_ACCESS_TOKEN, + channel: CHANNEL, + text: budgetNotificationText + }); + console.log(res); +}; +// [END functions_billing_slack] + +// [START functions_billing_stop] +const billing = google.cloudbilling('v1').projects; + +exports.stopBilling = async (data, context) => { + const pubsubData = JSON.parse(Buffer.from(data.data, 'base64').toString()); + if (pubsubData.costAmount <= pubsubData.budgetAmount) { + return `No action necessary. (Current cost: ${pubsubData.costAmount})`; + } + + await _setAuthCredential(); + if (await _isBillingEnabled(PROJECT_NAME)) { + return _disableBillingForProject(PROJECT_NAME); + } else { + return 'Billing already disabled'; + } +}; + +/** + * @return {Promise} Credentials set globally + */ +const _setAuthCredential = async () => { + const res = await auth.getApplicationDefault(); + + let client = res.credential; + if (client.hasScopes && !client.hasScopes()) { + client = client.createScoped([ + 'https://www.googleapis.com/auth/cloud-billing', + 'https://www.googleapis.com/auth/cloud-platform' + ]); + } + + // Set credential globally for all requests + google.options({ + auth: client + }); +}; + +/** + * Determine whether billing is enabled for a project + * @param {string} projectName Name of project to check if billing is enabled + * @return {bool} Whether project has billing enabled or not + */ +const _isBillingEnabled = async (projectName) => { + const res = await billing.getBillingInfo({ name: projectName }); + return res.data.billingEnabled; +}; + +/** + * Disable billing for a project by removing its billing account + * @param {string} projectName Name of project disable billing on + * @return {string} Text containing response from disabling billing + */ +const _disableBillingForProject = async (projectName) => { + const res = await billing.updateBillingInfo({ + name: projectName, + resource: { 'billingAccountName': '' } // Disable billing + }); + return `Billing disabled: ${JSON.stringify(res.data)}`; +}; +// [END functions_billing_stop] + +// [START functions_billing_limit] +const compute = google.compute('v1'); +const ZONE = 'us-west1-b'; + +exports.limitUse = async (data, context) => { + const pubsubData = JSON.parse(Buffer.from(data.data, 'base64').toString()); + if (pubsubData.costAmount <= pubsubData.budgetAmount) { + return `No action necessary. (Current cost: ${pubsubData.costAmount})`; + } + + await _setAuthCredential(); + const instanceNames = await _listRunningInstances(PROJECT_ID, ZONE); + await _stopInstances(PROJECT_ID, ZONE, instanceNames); +}; + +/** + * @return {Promise} Array of names of running instances + */ +const _listRunningInstances = async (projectId, zone) => { + const res = await compute.instances.list({ + project: projectId, + zone: zone + }); + + const instances = res.data.items || []; + const ranInstances = instances.filter(item => item.status === 'RUNNING'); + return ranInstances.map(item => item.name); +}; + +/** + * @param {Array} instanceNames Names of instance to stop + * @return {Promise} Response from stopping instances + */ +const _stopInstances = async (projectId, zone, instanceNames) => { + if (!instanceNames.length) { + return 'No running instances were found.'; + } + await Promise.all(instanceNames.map(instanceName => { + return compute.instances.stop({ + project: projectId, + zone: zone, + instance: instanceName + }).then((res) => { + console.log('Instance stopped successfully: ' + instanceName); + return res.data; + }); + })); +}; +// [END functions_billing_limit] diff --git a/functions/billing/package.json b/functions/billing/package.json new file mode 100644 index 0000000000..465dd1066b --- /dev/null +++ b/functions/billing/package.json @@ -0,0 +1,24 @@ +{ + "name": "cloud-functions-billing", + "version": "0.0.1", + "description": "Examples of integrating Cloud Functions with billing", + "main": "index.js", + "scripts": { + "lint": "repo-tools lint", + "pretest": "npm run lint", + "test": "ava test/*" + }, + "author": "Ace Nassri ", + "license": "Apache-2.0", + "dependencies": { + "google-auth-library": "^2.0.0", + "googleapis": "^33.0.0", + "slack": "^11.0.1" + }, + "devDependencies": { + "@google-cloud/nodejs-repo-tools": "^2.2.5", + "ava": "^0.25.0", + "proxyquire": "^2.1.0", + "sinon": "^6.3.4" + } +} diff --git a/functions/billing/test/index.test.js b/functions/billing/test/index.test.js new file mode 100644 index 0000000000..ef94be98de --- /dev/null +++ b/functions/billing/test/index.test.js @@ -0,0 +1,144 @@ +/** + * Copyright 2018, Google LLC. + * Licensed 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. + */ + +const proxyquire = require(`proxyquire`).noCallThru(); +const sinon = require(`sinon`); +const test = require(`ava`); + +function getSample () { + const instanceListMock = [ + { name: 'test-instance-1', status: 'RUNNING' }, + { name: 'test-instance-2', status: 'RUNNING' } + ]; + + const listInstancesResponseMock = { + data: { + items: instanceListMock + } + }; + + const computeMock = { + instances: { + list: sinon.stub().returns(listInstancesResponseMock), + stop: sinon.stub().resolves({ data: {} }) + } + }; + + const cloudbillingMock = { + projects: { + getBillingInfo: sinon.stub().resolves({ + data: { + billingEnabled: true + } + }), + updateBillingInfo: sinon.stub().returns({ + data: {} + }) + } + }; + + const googleMock = { + cloudbilling: sinon.stub().returns(cloudbillingMock), + compute: sinon.stub().returns(computeMock), + options: sinon.stub() + }; + + const googleapisMock = { + google: googleMock + }; + + const slackMock = { + chat: { + postMessage: sinon.stub().returns({ data: {} }) + } + }; + + const credentialMock = { + hasScopes: sinon.stub().returns(false) + }; + credentialMock.createScoped = sinon.stub().returns(credentialMock); + + const googleAuthMock = { + auth: { + getApplicationDefault: sinon.stub().resolves({ + credential: credentialMock + }) + } + }; + + return { + program: proxyquire(`../`, { + 'google-auth-library': googleAuthMock, + 'googleapis': googleapisMock, + 'slack': slackMock + }), + mocks: { + google: googleMock, + googleAuth: googleAuthMock, + googleapis: googleapisMock, + compute: computeMock, + cloudbilling: cloudbillingMock, + credential: credentialMock, + slack: slackMock, + instanceList: instanceListMock + } + }; +} + +test(`should notify Slack when budget is exceeded`, async t => { + const {program, mocks} = getSample(); + + const jsonData = { cost: 500, budget: 400 }; + const pubsubData = { + data: Buffer.from(JSON.stringify(jsonData)).toString('base64'), + attributes: {} + }; + + await program.notifySlack(pubsubData, null); + + t.true(mocks.slack.chat.postMessage.calledOnce); +}); + +test(`should disable billing when budget is exceeded`, async t => { + const {program, mocks} = getSample(); + + const jsonData = { cost: 500, budget: 400 }; + const pubsubData = { + data: Buffer.from(JSON.stringify(jsonData)).toString('base64'), + attributes: {} + }; + + await program.stopBilling(pubsubData, null); + + t.true(mocks.credential.createScoped.calledOnce); + t.true(mocks.cloudbilling.projects.getBillingInfo.calledOnce); + t.true(mocks.cloudbilling.projects.updateBillingInfo.calledOnce); +}); + +test(`should shut down GCE instances when budget is exceeded`, async t => { + const {program, mocks} = getSample(); + + const jsonData = { cost: 500, budget: 400 }; + const pubsubData = { + data: Buffer.from(JSON.stringify(jsonData)).toString('base64'), + attributes: {} + }; + + await program.limitUse(pubsubData, null); + + t.true(mocks.credential.createScoped.calledOnce); + t.true(mocks.compute.instances.list.calledOnce); + t.is(mocks.compute.instances.stop.callCount, mocks.instanceList.length); +});