feat(dataproc): add dataproc autoscaling policy

Author: james-zhou-inspire11

src/enums/schemasMap.ts

@@ -41,5 +41,6 @@ export default {
   [services.apiKey]: 'gcpApiKey',
   [services.computeProject]: 'gcpComputeProject',
   [services.dataprocCluster]: 'gcpDataprocCluster',
+  [services.dataprocAutoscalingPolicy]: 'gcpDataprocAutoscalingPolicy',
   tag: 'gcpTag',
 }

src/enums/serviceMap.ts

@@ -37,6 +37,7 @@ import GcpServiceAccount from '../services/serviceAccount'
 import GcpApiKey from '../services/apiKey'
 import GcpComputeProject from '../services/computeProject'
 import GcpDataprocCluster from '../services/dataprocCluster'
+import GcpDataprocAutoscalingPolicy from '../services/dataprocAutoscalingPolicy'
 
 /**
  * serviceMap is an object that contains all currently supported services
@@ -80,5 +81,6 @@ export default {
   [services.apiKey]: GcpApiKey,
   [services.computeProject]: GcpComputeProject,
   [services.dataprocCluster]: GcpDataprocCluster,
+  [services.dataprocAutoscalingPolicy]: GcpDataprocAutoscalingPolicy,
   tag: GcpTag,
 }

src/enums/services.ts

@@ -29,6 +29,7 @@ export default {
   // dataCatalog: 'data-catalog',
   // dataflow: 'dataflow',
   dataprocCluster: 'dataprocClusters',
+  dataprocAutoscalingPolicy: 'dataprocAutoscalingPolicies',
   // metastore: 'metastore',
   // pubsub: 'pubsub',
   // bigtable: 'bigtable',

src/services/dataprocAutoscalingPolicy/data.ts

@@ -0,0 +1,62 @@
+import { AutoscalingPolicyServiceClient } from '@google-cloud/dataproc'
+import CloudGraph from '@cloudgraph/sdk'
+import groupBy from 'lodash/groupBy'
+import { google } from '@google-cloud/dataproc/build/protos/protos'
+import gcpLoggerText from '../../properties/logger'
+import { GcpServiceInput } from '../../types'
+import { generateGcpErrorLog } from '../../utils'
+
+const lt = { ...gcpLoggerText }
+const { logger } = CloudGraph
+const serviceName = 'Dataproc Autoscaling Policy'
+
+export interface RawGcpDataprocAutoscalingPolicy extends Omit<google.cloud.dataproc.v1.IAutoscalingPolicy, 'id'|'labels'> {
+  id: string
+  region: string
+  projectId: string
+  Labels: { [key: string]: string }
+}
+
+export default async ({
+  regions,
+  config,
+}: GcpServiceInput): Promise<{
+  [region: string]: RawGcpDataprocAutoscalingPolicy[]
+}> => {
+    const autoscalingPolicyList: RawGcpDataprocAutoscalingPolicy[] = []
+    const { projectId } = config
+
+    for (const region of regions.split(',')) {
+      /**
+       * Get all the Dataproc Autoscaling Policies
+       */
+        
+      try {
+        const dataprocClient = new AutoscalingPolicyServiceClient({ 
+          ...config,
+          apiEndpoint: `${region}-dataproc.googleapis.com`,
+          projectId,
+        })
+        const locationName = dataprocClient.locationPath(projectId, region)
+        const iterable =  dataprocClient.listAutoscalingPoliciesAsync({
+          parent: locationName,
+        })
+        for await (const { id, labels, ...response } of iterable) {
+          if (response) {
+            autoscalingPolicyList.push({
+              ...response,
+              id,
+              projectId,
+              region,
+              Labels: labels,
+            })
+          }
+        }
+      } catch (error) {
+        generateGcpErrorLog(serviceName, 'dataprocAutoscalingPolicy:listAutoscalingPoliciesAsync', error)
+      }
+    }
+
+    logger.debug(lt.foundResources(serviceName, autoscalingPolicyList.length))
+    return groupBy(autoscalingPolicyList, 'region')
+  }

src/services/dataprocAutoscalingPolicy/format.ts

@@ -0,0 +1,43 @@
+import {
+  GcpDataprocAutoscalingPolicy
+} from '../../types/generated'
+import { RawGcpDataprocAutoscalingPolicy } from './data'
+import { formatLabelsFromMap } from '../../utils/format'
+
+export default ({
+  service,
+  region,
+}: {
+  service: RawGcpDataprocAutoscalingPolicy
+  region: string
+}): GcpDataprocAutoscalingPolicy => {
+  const {
+    id,
+    projectId,
+    name,
+    basicAlgorithm = {},
+    workerConfig = {},
+    secondaryWorkerConfig = {},
+    Labels = {},
+  } = service
+
+  return {
+    id,
+    projectId,
+    region,
+    name,
+    basicAlgorithmYarnConfigGracefulDecommissionTimeout: basicAlgorithm?.yarnConfig?.gracefulDecommissionTimeout?.seconds?.toString() || '',
+    basicAlgorithmYarnConfigScaleUpFactor: basicAlgorithm?.yarnConfig?.scaleUpFactor || 0.0,
+    basicAlgorithmYarnConfigScaleDownFactor: basicAlgorithm?.yarnConfig?.scaleDownFactor || 0.0,
+    basicAlgorithmYarnConfigScaleUpMinWorkerFraction: basicAlgorithm?.yarnConfig?.scaleUpMinWorkerFraction || 0.0,
+    basicAlgorithmYarnConfigScaleDownMinWorkerFraction: basicAlgorithm?.yarnConfig?.scaleDownMinWorkerFraction || 0.0,
+    basicAlgorithmCooldownPeriod: basicAlgorithm?.cooldownPeriod?.seconds?.toString() || '',
+    workerConfigMinInstances: workerConfig?.minInstances || 0,
+    workerConfigMaxInstances: workerConfig?.maxInstances || 0,
+    workerConfigWeight: workerConfig?.weight || 0,
+    secondaryWorkerConfigMinInstances: secondaryWorkerConfig?.minInstances || 0,
+    secondaryWorkerConfigMaxInstances: secondaryWorkerConfig?.maxInstances || 0,
+    secondaryWorkerConfigWeight: secondaryWorkerConfig?.weight || 0,
+    labels: formatLabelsFromMap(Labels),
+  }
+}

src/services/dataprocAutoscalingPolicy/index.ts

@@ -0,0 +1,13 @@
+import {Service} from '@cloudgraph/sdk'
+import BaseService from '../base'
+import format from './format'
+import getData from './data'
+import mutation from './mutation'
+
+export default class GcpDataprocAutoscalingPolicy extends BaseService implements Service {
+  format = format.bind(this)
+
+  getData = getData.bind(this)
+
+  mutation = mutation
+}

src/services/dataprocAutoscalingPolicy/mutation.ts

@@ -0,0 +1,5 @@
+export default `mutation($input: [AddgcpDataprocAutoscalingPolicyInput!]!) {
+  addgcpDataprocAutoscalingPolicy(input: $input, upsert: true) {
+    numUids
+  }
+}`;

src/services/dataprocAutoscalingPolicy/schema.graphql

@@ -0,0 +1,22 @@
+type gcpDataprocAutoscalingPolicy implements gcpBaseResource
+  @generate(
+    query: { get: true, query: true, aggregate: true }
+    mutation: { add: true, delete: false }
+    subscription: false
+  )
+  @key(fields: "id") {
+    basicAlgorithmYarnConfigGracefulDecommissionTimeout: String @search(by: [hash, regexp])
+    basicAlgorithmYarnConfigScaleUpFactor: Float @search
+    basicAlgorithmYarnConfigScaleDownFactor: Float @search
+    basicAlgorithmYarnConfigScaleUpMinWorkerFraction: Float @search
+    basicAlgorithmYarnConfigScaleDownMinWorkerFraction: Float @search
+    basicAlgorithmCooldownPeriod: String @search(by: [hash, regexp])
+    workerConfigMinInstances: Int @search
+    workerConfigMaxInstances: Int @search
+    workerConfigWeight: Int @search
+    secondaryWorkerConfigMinInstances: Int @search
+    secondaryWorkerConfigMaxInstances: Int @search
+    secondaryWorkerConfigWeight: Int @search
+    labels: [gcpRawLabel]
+    project: [gcpProject] @hasInverse(field: dataprocAutoscalingPolicies)
+  }

src/services/dataprocCluster/index.ts

@@ -4,7 +4,7 @@ import format from './format'
 import getData from './data'
 import mutation from './mutation'
 
-export default class GcpDataproc extends BaseService implements Service {
+export default class GcpDataprocCluster extends BaseService implements Service {
   format = format.bind(this)
 
   getData = getData.bind(this)

src/services/project/schema.graphql

@@ -47,4 +47,5 @@ type gcpProject @key(fields: "id") {
   kmsCryptoKeys: [gcpKmsCryptoKey] @hasInverse(field: project)
   dataprocCluster: [gcpDataprocCluster] @hasInverse(field: project)
   dataprocClusters: [gcpDataprocCluster] @hasInverse(field: project)
+  dataprocAutoscalingPolicies: [gcpDataprocAutoscalingPolicy] @hasInverse(field: project)
 }