Fixes to move deployment to stage env (#2053)

As part of moving the deployments to the stage, we had to make some fixes in the script. sh and GitHub Actions files.

.github/workflows/ci-astro-deploy.yml

@@ -57,8 +57,7 @@ jobs:
       docker_registry:  ${{ secrets.ASTRO_DOCKER_REGISTRY }}
       organization_id:  ${{ secrets.ORGANIZATION_ID }}
       deployment_id:  ${{ secrets.ASTRO_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.ASTRO_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.ASTRO_KEY_SECRET }}
+      bearer_token: ${{ secrets.BEARER_TOKEN }}
       SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }}
 
   wait-for-deployment-to-be-ready-and-trigger-dags-for-astro-sdk-integration-tests:
@@ -73,10 +72,8 @@ jobs:
     secrets:
       astro_subdomain: ${{ secrets.ASTRO_SUBDOMAIN }}
       deployment_id: ${{ secrets.ASTRO_DEPLOYMENT_ID }}
-      astronomer_key_id: ${{ secrets.ASTRO_KEY_ID }}
-      astronomer_key_secret: ${{ secrets.ASTRO_KEY_SECRET }}
-      organization_id: ${{ secrets.ORGANIZATION_ID }}
       bearer_token: ${{ secrets.BEARER_TOKEN }}
+      organization_id: ${{ secrets.ORGANIZATION_ID }}
 
   deploy-to-astro-sdk-integration-tests-on-KE:
     if: |
@@ -90,9 +87,8 @@ jobs:
       docker_registry:  ${{ secrets.ASTRO_DOCKER_REGISTRY }}
       organization_id:  ${{ secrets.ORGANIZATION_ID }}
       deployment_id:  ${{ secrets.ASTRO_DEPLOYMENT_ID_KE }}
-      astronomer_key_id: ${{ secrets.ASTRO_KEY_ID_KE }}
-      astronomer_key_secret: ${{ secrets.ASTRO_KEY_SECRET_KE }}
       SLACK_WEBHOOK_URL: ${{ secrets.SLACK_WEBHOOK_URL }}
+      bearer_token: ${{ secrets.BEARER_TOKEN }}
 
   wait-for-deployment-to-be-ready-and-trigger-dags-for-astro-sdk-integration-tests-on-KE:
     if: |
@@ -106,7 +102,5 @@ jobs:
     secrets:
       astro_subdomain: ${{ secrets.ASTRO_SUBDOMAIN }}
       deployment_id: ${{ secrets.ASTRO_DEPLOYMENT_ID_KE }}
-      astronomer_key_id: ${{ secrets.ASTRO_KEY_ID_KE }}
-      astronomer_key_secret: ${{ secrets.ASTRO_KEY_SECRET_KE }}
       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'
@@ -53,8 +50,7 @@ jobs:
           bash deploy.sh  ${{ 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
 
@@ -75,8 +69,7 @@ jobs:
       - name: Trigger DAG(s)
         run: |
           python3 python-sdk/dev/integration_test_scripts/trigger_dag.py \
-            "org-airflow-team-org" \
+            ${{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 }}"

python-sdk/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_name: 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",
@@ -66,8 +48,7 @@ def trigger_dag_runs(
     parser = argparse.ArgumentParser()
     parser.add_argument("astro_subdomain", help="subdomain of the Astro Cloud", 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="Key ID of the Astro Cloud deployment", type=str)
     parser.add_argument(
         "--dag-ids",
         help=(
@@ -80,7 +61,7 @@ def trigger_dag_runs(
     )
 
     args = parser.parse_args()
-    token = get_access_token(args.astronomer_key_id.strip(), args.astronomer_key_secret.strip())
+    token = args.token.strip()
 
     input_dag_ids = args.dag_ids
     dag_ids = [dag_id.strip() for dag_id in input_dag_ids.split(",")]

python-sdk/tests_integration/astro_deploy/deploy.sh

@@ -12,8 +12,7 @@ set -e
 #         - ASTRO_DOCKER_REGISTRY: Docker registry domain. Script will push the docker image here.
 #         - ASTRO_ORGANIZATION_ID: Astro cloud deployment organization Id. Get it from UI.
 #         - ASTRO_DEPLOYMENT_ID: Astro cloud deployment Id. Get it from UI.
-#         - ASTRO_KEY_ID: Astro cloud deployment service account API key Id.
-#         - ASTRO_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}/../.."
@@ -23,9 +22,8 @@ function echo_help() {
     echo "ASTRO_DOCKER_REGISTRY:        Docker registry"
     echo "ASTRO_ORGANIZATION_ID         Astro cloud organization Id"
     echo "ASTRO_DEPLOYMENT_ID           Astro cloud Deployment id"
-    echo "ASTRO_KEY_ID       Astro cloud service account API key id"
-    echo "ASTRO_KEY_SECRET   Astro cloud service account API key secret"
-    echo "bash deploy.sh <ASTRO_DOCKER_REGISTRY> <ASTRO_ORGANIZATION_ID>  <ASTRO_DEPLOYMENT_ID> <ASTRO_KEY_ID> <ASTRO_KEY_SECRET>"
+    echo "TOKEN     Astro workspace token"
+    echo "bash deploy.sh <ASTRO_DOCKER_REGISTRY> <ASTRO_ORGANIZATION_ID>  <ASTRO_DEPLOYMENT_ID> <TOKEN>"
 }
 
 if [ "$1" == "-h" ]; then
@@ -43,8 +41,7 @@ function clean() {
 ASTRO_DOCKER_REGISTRY=$1
 ASTRO_ORGANIZATION_ID=$2
 ASTRO_DEPLOYMENT_ID=$3
-ASTRO_KEY_ID=$4
-ASTRO_KEY_SECRET=$5
+TOKEN=$4
 MASTER_DAG_DOCKERFILE="Dockerfile"
 
 
@@ -54,58 +51,50 @@ function deploy(){
     docker_registry_astro=$1
     organization_id=$2
     deployment_id=$3
-    key_id=$4
-    key_secret=$5
-    dockerfile=$6
+    TOKEN=$4
+    dockerfile=$5
 
     # Build image and deploy
     BUILD_NUMBER=$(awk 'BEGIN {srand(); print srand()}')
-
     # Enforce registry name to be in lowercase
     docker_registry_astro=$(echo $docker_registry_astro | tr '[:upper:]' '[:lower:]')
     organization_id=$(echo $organization_id | tr '[:upper:]' '[:lower:]')
     deployment_id=$(echo $deployment_id | tr '[:upper:]' '[:lower:]')
     IMAGE_NAME=${docker_registry_astro}/${organization_id}/${deployment_id}:ci-${BUILD_NUMBER}
     docker build --platform=linux/amd64 -t "${IMAGE_NAME}" -f "${SCRIPT_PATH}"/${dockerfile} "${SCRIPT_PATH}"
-    docker login "${docker_registry_astro}" -u "${key_id}" -p "${key_secret}"
+    docker login "${docker_registry_astro}" -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\": \"$key_id\",
-            \"client_secret\": \"$key_secret\",
-            \"audience\": \"astronomer-ee\",
-            \"grant_type\": \"client_credentials\"}" | jq -r '.access_token' )
-
     # Create the Image
-    echo "get image id"
-    IMAGE=$( curl --location --request POST "https://api.astronomer.io/hub/v1" \
-        --header "Authorization: Bearer $TOKEN" \
+      echo "get image id"
+      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')
     # Deploy the Image
     echo "deploy image"
-    curl --location --request POST "https://api.astronomer.io/hub/v1" \
-            --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}\",
-                \"variables\" : {
-                    \"input\" : {
-                        \"id\" : \"$IMAGE\",
-                        \"tag\" : \"ci-$BUILD_NUMBER\",
-                        \"repository\" : \"images.astronomer.cloud/$organization_id/$deployment_id\"
-                        }
-                    }
+    curl --location --request POST "https://api.astronomer-stage.io/hub/graphql" \
+          --header "Authorization: Bearer $TOKEN" \
+          --header "Content-Type: application/json" \
+          --data-raw "{
+              \"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\" : {
+                      \"deploymentId\" : \"$deployment_id\",
+                      \"imageId\" : \"$IMAGE\",
+                      \"tag\" : \"ci-$BUILD_NUMBER\",
+                      \"repository\" : \"images.astronomer-stage.cloud/$organization_id/$deployment_id\",
+                      \"dagDeployEnabled\":false
+                      }
+                  }
             }"
 
 }
@@ -120,6 +109,6 @@ cp -r "${PROJECT_PATH}"/example_dags "${SCRIPT_PATH}"/example_dags
 cp -r "${PROJECT_PATH}"/tests/data "${SCRIPT_PATH}"/tests/data
 
 
-deploy $ASTRO_DOCKER_REGISTRY $ASTRO_ORGANIZATION_ID $ASTRO_DEPLOYMENT_ID $ASTRO_KEY_ID $ASTRO_KEY_SECRET $MASTER_DAG_DOCKERFILE
+deploy $ASTRO_DOCKER_REGISTRY $ASTRO_ORGANIZATION_ID $ASTRO_DEPLOYMENT_ID $TOKEN $MASTER_DAG_DOCKERFILE
 
 clean

python-sdk/tests_integration/databases/test_snowflake.py

@@ -87,7 +87,7 @@ def test_snowflake_create_table_with_columns(database_table_fixture):
         "NUMBER(38,0)",
         "COLUMN",
         "N",
-        "IDENTITY START 1 INCREMENT 1",
+        "IDENTITY START 1 INCREMENT 1 ORDER",
         "Y",
         "N",
         None,