Fixes to move deployment to stage env

.circleci/integration-tests/script.sh

@@ -12,8 +12,7 @@ set -e
 #         - DOCKER_REGISTRY: Docker registry domain. Script will push the docker image here.
 #         - ORGANIZATION_ID: Astro cloud deployment organization Id. Get it from UI.
 #         - DEPLOYMENT_ID: Astro cloud deployment Id. Get it from UI.
-#         - ASTRONOMER_KEY_ID: Astro cloud deployment service account API key Id.
-#         - ASTRONOMER_KEY_SECRET: Astro cloud deployment service account API key secret.
+#         - TOKEN: Astro workspace token.
 
 SCRIPT_PATH="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 PROJECT_PATH=${SCRIPT_PATH}/../../
@@ -25,9 +24,8 @@ function echo_help() {
     echo "DOCKER_REGISTRY:        Docker registry"
     echo "ORGANIZATION_ID         Astro cloud organization Id"
     echo "DEPLOYMENT_ID           Astro cloud Deployment id"
-    echo "ASTRONOMER_KEY_ID       Astro cloud service account API key id"
-    echo "ASTRONOMER_KEY_SECRET   Astro cloud service account API key secret"
-    echo "bash script.sh astro-cloud <DOCKER_REGISTRY> <ORGANIZATION_ID>  <DEPLOYMENT_ID> <ASTRONOMER_KEY_ID> <ASTRONOMER_KEY_SECRET>"
+    echo "TOKEN     Astro workspace token"
+    echo "bash script.sh astro-cloud <DOCKER_REGISTRY> <ORGANIZATION_ID>  <DEPLOYMENT_ID> <TOKEN>"
 }
 
 # Delete if source old source files exist
@@ -49,15 +47,13 @@ DEPLOYMENT_INSTANCE=$1
 DOCKER_REGISTRY=""
 ORGANIZATION_ID=""
 DEPLOYMENT_ID=""
-ASTRONOMER_KEY_ID=""
-ASTRONOMER_KEY_SECRET=""
+TOKEN=""
 
 if [[ ${DEPLOYMENT_INSTANCE} == "astro-cloud" ]]; then
   DOCKER_REGISTRY=$2
   ORGANIZATION_ID=$3
   DEPLOYMENT_ID=$4
-  ASTRONOMER_KEY_ID=$5
-  ASTRONOMER_KEY_SECRET=$6
+  TOKEN=$5
 else
   echo "Valid value for DEPLOYMENT_INSTANCE can only be astro-cloud"
   echo_help
@@ -80,45 +76,43 @@ BUILD_NUMBER=$(awk 'BEGIN {srand(); print srand()}')
 if [[ ${DEPLOYMENT_INSTANCE} == "astro-cloud" ]]; then
   IMAGE_NAME=${DOCKER_REGISTRY}/${ORGANIZATION_ID}/${DEPLOYMENT_ID}:ci-${BUILD_NUMBER}
   docker build --platform=linux/amd64 -t "${IMAGE_NAME}" -f "${SCRIPT_PATH}"/Dockerfile.astro_cloud "${SCRIPT_PATH}"
-  docker login "${DOCKER_REGISTRY}" -u "${ASTRONOMER_KEY_ID}" -p "${ASTRONOMER_KEY_SECRET}"
+  docker login "${DOCKER_REGISTRY}" -u "cli" -p "${TOKEN}"
   docker push "${IMAGE_NAME}"
 
-  TOKEN=$( curl --location --request POST "https://auth.astronomer.io/oauth/token" \
-        --header "content-type: application/json" \
-        --data-raw "{
-            \"client_id\": \"$ASTRONOMER_KEY_ID\",
-            \"client_secret\": \"$ASTRONOMER_KEY_SECRET\",
-            \"audience\": \"astronomer-ee\",
-            \"grant_type\": \"client_credentials\"}" | jq -r '.access_token' )
+
   # Step 5. Create the Image
   echo "get image id"
-  IMAGE=$( curl --location --request POST "https://api.astronomer.io/hub/v1" \
-        --header "Authorization: Bearer $TOKEN" \
+  IMAGE=$( curl --location --request POST "https://api.astronomer-stage.io/hub/graphql" \
+        --header "Authorization: Bearer "${TOKEN}"" \
         --header "Content-Type: application/json" \
         --data-raw "{
-            \"query\" : \"mutation imageCreate(\n    \$input: ImageCreateInput!\n) {\n    imageCreate (\n    input: \$input\n) {\n    id\n    tag\n    repository\n    digest\n    env\n    labels\n    deploymentId\n  }\n}\",
+            \"query\" : \"mutation CreateImage(\n    \$input: CreateImageInput!\n) {\n    createImage (\n    input: \$input\n) {\n    id\n    tag\n    repository\n    digest\n    env\n    labels\n    deploymentId\n  }\n}\",
             \"variables\" : {
                 \"input\" : {
                     \"deploymentId\" : \"$DEPLOYMENT_ID\",
                     \"tag\" : \"ci-$BUILD_NUMBER\"
                     }
                 }
-            }" | jq -r '.data.imageCreate.id')
+            }" | jq -r '.data.createImage.id')
   # Step 6. Deploy the Image
   echo "deploy image"
-  curl --location --request POST "https://api.astronomer.io/hub/v1" \
+
+  curl --location --request POST "https://api.astronomer-stage.io/hub/graphql" \
           --header "Authorization: Bearer $TOKEN" \
           --header "Content-Type: application/json" \
           --data-raw "{
-              \"query\" : \"mutation imageDeploy(\n    \$input: ImageDeployInput!\n  ) {\n    imageDeploy(\n      input: \$input\n    ) {\n      id\n      deploymentId\n      digest\n      env\n      labels\n      name\n      tag\n      repository\n    }\n}\",
+              \"query\" : \"mutation DeployImage(\n    \$input: DeployImageInput!\n  ) {\n    deployImage(\n      input: \$input\n    ) {\n      id\n      deploymentId\n      digest\n      env\n      labels\n      name\n      tag\n      repository\n    }\n}\",
               \"variables\" : {
                   \"input\" : {
-                      \"id\" : \"$IMAGE\",
+                      \"deploymentId\" : \"$DEPLOYMENT_ID\",
+                      \"imageId\" : \"$IMAGE\",
                       \"tag\" : \"ci-$BUILD_NUMBER\",
-                      \"repository\" : \"images.astronomer.cloud/$ORGANIZATION_ID/$DEPLOYMENT_ID\"
+                      \"repository\" : \"images.astronomer-stage.cloud/$ORGANIZATION_ID/$DEPLOYMENT_ID\",
+                      \"dagDeployEnabled\":false
                       }
                   }
             }"
+
 fi
 
 clean

.github/workflows/deploy-integration-tests.yaml

@@ -39,8 +39,7 @@ jobs:
       docker_registry: ${{ secrets.DOCKER_REGISTRY }}
       organization_id: ${{ secrets.ORGANIZATION_ID }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_SECRET }}
+      bearer_token: ${{ secrets.BEARER_TOKEN }}
       SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }}
 
   wait-for-deployment-to-be-ready-and-trigger-dags-for-providers-integration-tests:
@@ -55,8 +54,6 @@ jobs:
     secrets:
       astro_subdomain: ${{ secrets.ASTRO_SUBDOMAIN }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_SECRET }}
       organization_id: ${{ secrets.ORGANIZATION_ID }}
       bearer_token: ${{ secrets.BEARER_TOKEN }}
 
@@ -72,8 +69,7 @@ jobs:
       docker_registry: ${{ secrets.DOCKER_REGISTRY }}
       organization_id: ${{ secrets.ORGANIZATION_ID }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ON_KE_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ON_KE_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ON_KE_ASTRONOMER_KEY_SECRET }}
+      bearer_token: ${{ secrets.BEARER_TOKEN }}
       SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }}
 
   wait-for-deployment-to-be-ready-and-trigger-dags-for-providers-integration-tests-on-KE:
@@ -88,7 +84,5 @@ jobs:
     secrets:
       astro_subdomain: ${{ secrets.ASTRO_SUBDOMAIN }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ON_KE_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ON_KE_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets. PROVIDER_INTEGRATION_TESTS_ON_KE_ASTRONOMER_KEY_SECRET }}
       organization_id: ${{ secrets.ORGANIZATION_ID }}
       bearer_token: ${{ secrets.BEARER_TOKEN }}

.github/workflows/reuse-wf-deploy-to-astro-cloud.yaml

@@ -23,11 +23,8 @@ on:  # yamllint disable-line rule:truthy
       deployment_id:
         description: 'astro cloud deployment_id'
         required: true
-      astronomer_key_id:
-        description: 'astro cloud astronomer_key_id'
-        required: true
-      astronomer_key_secret:
-        description: 'astro cloud astronomer_key_secret'
+      bearer_token:
+        description: 'workspace bearer token'
         required: true
       SLACK_WEBHOOK_URL:
         description: 'slack webhook url for sending notification'
@@ -54,8 +51,7 @@ jobs:
             ${{ secrets.docker_registry }} \
             ${{ secrets.organization_id }} \
             ${{ secrets.deployment_id }} \
-            ${{ secrets.astronomer_key_id }} \
-            ${{ secrets.astronomer_key_secret }}
+            ${{ secrets.bearer_token }} \
 
       - name: send succeeded notification to Slack
         if: success() && github.event_name == 'workflow_dispatch'

.github/workflows/reuse-wf-trigger-dag.yaml

@@ -23,12 +23,6 @@ on:  # yamllint disable-line rule:truthy
       deployment_id:
         description: 'astro cloud deployment_id'
         required: true
-      astronomer_key_id:
-        description: 'astro cloud astronomer_key_id'
-        required: true
-      astronomer_key_secret:
-        description: 'astro cloud astronomer_key_secret'
-        required: true
       organization_id:
         description: 'astro cloud organization_id'
         required: true
@@ -43,8 +37,8 @@ jobs:
 
       - name: Wait for deployment to be healthy
         run: |
-          astro_core_api="https://api.astronomer.io/v1alpha1/organizations/${{secrets.organization_id }}/\
-          deployments"
+          astro_core_api="https://api.astronomer-stage.io/v1alpha1/organizations/\
+          ${{secrets.organization_id }}/deployments"
           tries=15
           health_flag=false
 
@@ -74,8 +68,7 @@ jobs:
       - name: Trigger DAG(s)
         run: |
           python3 dev/integration_test_scripts/trigger_dag.py \
-            ${{ secrets.astro_subdomain }} \
+            ${{secrets.organization_id }} \
             ${{ secrets.deployment_id }} \
-            ${{ secrets.astronomer_key_id }} \
-            ${{ secrets.astronomer_key_secret }} \
+            ${{ secrets.bearer_token }} \
             --dag-ids "${{ inputs.dags_to_trigger_after_deployment }}"

.github/workflows/test-rc-release.yaml

@@ -51,8 +51,7 @@ jobs:
       docker_registry: ${{ secrets.DOCKER_REGISTRY }}
       organization_id: ${{ secrets.ORGANIZATION_ID }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_SECRET }}
+      bearer_token: ${{ secrets.BEARER_TOKEN }}
       SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }}
 
   wait-for-deployment-to-be-ready-and-trigger-master-dag:
@@ -67,10 +66,8 @@ jobs:
     secrets:
       astro_subdomain: ${{ secrets.ASTRO_SUBDOMAIN }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.PROVIDER_INTEGRATION_TESTS_ASTRONOMER_KEY_SECRET }}
-      organization_id: ${{ secrets.ORGANIZATION_ID }}
       bearer_token: ${{ secrets.BEARER_TOKEN }}
+      organization_id: ${{ secrets.ORGANIZATION_ID }}
 
   deploy-rc-testing-branch-to-astro-cloud-on-GCP:
     needs: check-airflow-provider-rc-release
@@ -85,8 +82,7 @@ jobs:
       docker_registry: ${{ secrets.DOCKER_REGISTRY }}
       organization_id: ${{ secrets.ORGANIZATION_ID }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_GCP_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_GCP_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.PROVIDER_INTEGRATION_TESTS_GCP_ASTRONOMER_KEY_SECRET }}
+      bearer_token: ${{ secrets.BEARER_TOKEN }}
       SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }}
 
   wait-for-deployment-to-be-ready-and-trigger-master-dag-on-GCP:
@@ -101,7 +97,5 @@ jobs:
     secrets:
       astro_subdomain: ${{ secrets.ASTRO_SUBDOMAIN }}
       deployment_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_GCP_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.PROVIDER_INTEGRATION_TESTS_GCP_ASTRONOMER_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.PROVIDER_INTEGRATION_TESTS_GCP_ASTRONOMER_KEY_SECRET }}
       organization_id: ${{ secrets.ORGANIZATION_ID }}
       bearer_token: ${{ secrets.BEARER_TOKEN }}

dev/integration_test_scripts/trigger_dag.py

@@ -9,24 +9,6 @@
 logging.basicConfig(stream=sys.stdout, level=logging.DEBUG)
 
 
-def get_access_token(api_key_id: str, api_key_secret: str) -> str:
-    """
-    Gets bearer access token for the Astro Cloud deployment needed for REST API authentication.
-
-    :param api_key_id: API key ID of the Astro Cloud deployment
-    :param api_key_secret: API key secret of the Astro Cloud deployment
-    """
-    request_json = {
-        "client_id": api_key_id,
-        "client_secret": api_key_secret,
-        "audience": "astronomer-ee",
-        "grant_type": "client_credentials",
-    }
-    response = requests.post("https://auth.astronomer.io/oauth/token", json=request_json)
-    response_json = response.json()
-    return response_json["access_token"]
-
-
 def trigger_dag_runs(
     *, dag_ids: list[str], astro_subdomain: str, deployment_id: str, bearer_token: str
 ) -> None:
@@ -40,7 +22,7 @@ def trigger_dag_runs(
     :param bearer_token: bearer token to be used for authentication with the Airflow REST API
     """
     short_deployment_id = f"d{deployment_id[-7:]}"
-    integration_tests_deployment_url = f"https://{astro_subdomain}.astronomer.run/{short_deployment_id}"
+    integration_tests_deployment_url = f"https://{astro_subdomain}.astronomer-stage.run/{short_deployment_id}"
     headers = {
         "Content-Type": "application/json",
         "Cache-Control": "no-cache",
@@ -64,10 +46,9 @@ def trigger_dag_runs(
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("astro_subdomain", help="subdomain of the Astro Cloud", type=str)
+    parser.add_argument("organization_id", help="organization id", type=str)
     parser.add_argument("deployment_id", help="ID of the deployment in Astro Cloud", type=str)
-    parser.add_argument("astronomer_key_id", help="Key ID of the Astro Cloud deployment", type=str)
-    parser.add_argument("astronomer_key_secret", help="Key secret of the Astro Cloud deployment", type=str)
+    parser.add_argument("token", help="astro workspace token", type=str)
     parser.add_argument(
         "--dag-ids",
         help=(
@@ -80,14 +61,13 @@ def trigger_dag_runs(
     )
 
     args = parser.parse_args()
-    token = get_access_token(args.astronomer_key_id.strip(), args.astronomer_key_secret.strip())
 
     input_dag_ids = args.dag_ids
     dag_ids = [dag_id.strip() for dag_id in input_dag_ids.split(",")]
 
     trigger_dag_runs(
         dag_ids=dag_ids,
-        astro_subdomain=args.astro_subdomain,
+        astro_subdomain=args.organization_id,
         deployment_id=args.deployment_id,
-        bearer_token=token,
+        bearer_token=args.token.strip(),
     )