Your IP : 216.73.217.117


Current Path : /lib/google-cloud-sdk/lib/surface/orchestration_pipelines/
Upload File :
Current File : //lib/google-cloud-sdk/lib/surface/orchestration_pipelines/init_demo.py

# -*- coding: utf-8 -*- #
# Copyright 2026 Google LLC. All Rights Reserved.
#
# 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.
"""Init-demo command for orchestration pipelines."""


import pathlib

from googlecloudsdk.calliope import base as calliope_base
from googlecloudsdk.calliope import exceptions
from googlecloudsdk.core import log
from googlecloudsdk.core.util import files


_PIPELINE_YAML = """\
pipelineId: orchestration-pipeline
description: 'Test deployment pipeline for Dataproc Serverless PySpark job'
runner: 'airflow'
modelVersion: '1.0'
owner: 'data-eng-team'

defaults:
  projectId: {{ project }}
  location: {{ region }}
  executionConfig:
    retries: 1

triggers:
  - schedule:
      interval: '0 2 * * *'  # 2 AM daily
      startTime: '2025-11-01T00:00:00'
      endTime: '2026-12-01T00:00:00'
      catchup: false
      timezone: 'UTC'

actions:
  - pyspark:
      name: test_job
      mainFilePath: 'jobs/job.py'
      environment:
        requirements:
          path: 'jobs/requirements.txt'
      engine:
        dataprocServerless:
          location: {{ region }}
          resourceProfile:
            path: './resources/profiles/serverless-standard.yaml'
"""

_DEPLOYMENT_YAML = """\
environments:
  dev:
    # TODO: Replace with your GCP project
    project: YOUR_PROJECT_ID
    # TODO: Replace with your region
    region: us-central1
    artifact_storage:
      # TODO: Replace with your artifacts bucket
      bucket: YOUR_ARTIFACTS_BUCKET
      path_prefix: pipelines
    # TODO: Replace with your Composer environment
    composer_environment: YOUR_COMPOSER_ENVIRONMENT
    variables:
      # TODO: Replace with your service account
      service_account: YOUR_SERVICE_ACCOUNT
      # TODO: Replace with your network URI
      network_uri: projects/YOUR_PROJECT_ID/global/networks/default
    pipelines:
      - source: orchestration-pipeline.yaml
  production:
    # TODO: Replace with your GCP project
    project: YOUR_PROJECT_ID
    # TODO: Replace with your region
    region: us-central1
    artifact_storage:
      # TODO: Replace with your artifacts bucket
      bucket: YOUR_ARTIFACTS_BUCKET
      path_prefix: pipelines
    # TODO: Replace with your Composer environment
    composer_environment: YOUR_COMPOSER_ENVIRONMENT
    variables:
      # TODO: Replace with your service account
      service_account: YOUR_SERVICE_ACCOUNT
      # TODO: Replace with your network URI
      network_uri: projects/YOUR_PROJECT_ID/global/networks/default
    pipelines:
      - source: orchestration-pipeline.yaml
"""

_JOB_SCRIPT = """\
\"\"\"
Sensor Temperature Pipeline
Converts raw Celsius readings to Fahrenheit using Pint,
then writes the enriched dataset to GCS.
\"\"\"

import pint
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

# TODO: Replace with your output bucket
GCS_OUTPUT_PATH = "gs://YOUR_OUTPUT_BUCKET/deploy-auto/output"

spark = SparkSession.builder.appName("sensor-temperature-pipeline").getOrCreate()

# Raw sensor readings in Celsius
data = [
    (1, "2025-01-01 00:00", 22.5),
    (1, "2025-01-01 01:00", 23.1),
    (2, "2025-01-01 00:00", 30.0),
    (2, "2025-01-01 01:00", 31.2),
    (3, "2025-01-01 00:00", 17.8),
    (3, "2025-01-01 01:00", 16.4),
]
df = spark.createDataFrame(data, ["sensor_id", "timestamp", "temp_celsius"])

# Convert Celsius -> Fahrenheit using Pint
ureg = pint.UnitRegistry()


@F.udf(returnType="double")
def celsius_to_fahrenheit(c):
    return float((c * ureg.degC).to(ureg.degF).magnitude)


df = df.withColumn("temp_fahrenheit", celsius_to_fahrenheit("temp_celsius"))
df = df.withColumn(
    "status",
    F.when(F.col("temp_celsius") > 30, "HIGH")
    .when(F.col("temp_celsius") < 18, "LOW")
    .otherwise("NORMAL"),
)

df.show(truncate=False)

df.write.mode("overwrite").partitionBy("sensor_id").parquet(GCS_OUTPUT_PATH)
print(f"Wrote {df.count()} records to {GCS_OUTPUT_PATH}")

spark.stop()
"""

_REQUIREMENTS_TXT = 'pint==0.24'

_PROFILE_YAML = """\
profileId: serverless-standard
type: dataproc.session
definition:
  environmentConfig:
    execution_config:
      service_account: "{{{{ service_account }}}}"
      network_uri: "{{{{ network_uri }}}}"
"""

_VALIDATE_WORKFLOW_TEMPLATE = """\
name: Validate Pipelines

on:
  pull_request:
    branches: [main]

jobs:
  validate:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3
      - uses: google-github-actions/auth@v2
        with:
          credentials_json: "${{{{ secrets.GCP_SA_KEY }}}}"
      - uses: google-github-actions/setup-gcloud@v1
      - run: gcloud components update --quiet
      - run: gcloud components install beta --quiet
      - uses: astral-sh/setup-uv@v7
      - run: gcloud beta orchestration-pipelines validate --pipeline-paths=orchestration-pipeline.yaml --environment=dev --quiet
      - run: gcloud beta orchestration-pipelines validate --pipeline-paths=orchestration-pipeline.yaml --environment=production --quiet
"""

_DEPLOY_WORKFLOW_TEMPLATE = """\
name: Deploy to {environment_title}

on:
  push:
    branches: [main]

jobs:
  deploy:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v3
      - uses: google-github-actions/auth@v2
        with:
          credentials_json: "${{{{ secrets.GCP_SA_KEY }}}}"
      - uses: google-github-actions/setup-gcloud@v1
      - run: gcloud components update --quiet
      - run: gcloud components install beta --quiet
      - uses: astral-sh/setup-uv@v7
      - run: gcloud beta orchestration-pipelines deploy --environment={environment} --quiet
"""


@calliope_base.Hidden
@calliope_base.DefaultUniverseOnly
@calliope_base.ReleaseTracks(calliope_base.ReleaseTrack.BETA)
class InitDemo(calliope_base.Command):
  """Initialize a demo orchestration pipeline project."""

  @staticmethod
  def Args(parser):
    parser.add_argument(
        'directory',
        nargs='?',
        default='orchestration-pipelines-demo',
        help='The directory to initialize the demo in.',
    )

  def Run(self, args):
    work_dir = pathlib.Path.cwd()
    demo_dir = work_dir / args.directory

    if demo_dir.exists():
      raise exceptions.BadFileException(
          f'Directory {demo_dir} already exists. Please specify a new directory'
          ' or remove the existing one.'
      )

    files.MakeDir(str(demo_dir))

    # Create file structure
    self._WriteFile(demo_dir / 'orchestration-pipeline.yaml', _PIPELINE_YAML)
    self._WriteFile(demo_dir / 'deployment.yaml', _DEPLOYMENT_YAML)

    jobs_dir = demo_dir / 'jobs'
    files.MakeDir(str(jobs_dir))
    self._WriteFile(jobs_dir / 'job.py', _JOB_SCRIPT)

    self._WriteFile(jobs_dir / 'requirements.txt', _REQUIREMENTS_TXT)

    profiles_dir = demo_dir / 'resources' / 'profiles'
    files.MakeDir(str(profiles_dir))
    self._WriteFile(profiles_dir / 'serverless-standard.yaml', _PROFILE_YAML)

    # Create GitHub workflows
    workflows_dir = demo_dir / '.github' / 'workflows'
    files.MakeDir(str(workflows_dir))

    self._WriteFile(
        workflows_dir / 'validate.yaml', _VALIDATE_WORKFLOW_TEMPLATE.format()
    )

    self._WriteFile(
        workflows_dir / 'deploy-dev.yaml',
        _DEPLOY_WORKFLOW_TEMPLATE.format(
            environment='dev', environment_title='Dev'
        ),
    )

    self._WriteFile(
        workflows_dir / 'deploy-prod.yaml',
        _DEPLOY_WORKFLOW_TEMPLATE.format(
            environment='production', environment_title='Production'
        ),
    )

    log.status.Print(f'Initialized demo project in {demo_dir}')

  def _WriteFile(self, path, content):
    files.WriteFileContents(path, content)