From 76cd060aad73e411d4b4ac2fcb98035af494b841 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Mon, 9 Dec 2024 13:17:50 -0800
Subject: [PATCH 01/19] Add a script to automatically provision a Prefect Cloud
 environment for the debugging tutorial

---
 .gitignore           |  12 ++++-
 infra/destroy_env.sh |  43 ++++++++++++++++++
 infra/main.tf        |  38 ++++++++++++++++
 infra/variables.tf   |  10 +++++
 setup_env.sh         | 103 +++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 205 insertions(+), 1 deletion(-)
 create mode 100755 infra/destroy_env.sh
 create mode 100644 infra/main.tf
 create mode 100644 infra/variables.tf
 create mode 100755 setup_env.sh

diff --git a/.gitignore b/.gitignore
index 151121e..415e09a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,12 @@
-env
+# Python
 __pycache__
+env
+venv
+
+# Terraform files
+.terraform
+*.tfstate
+*.tfstate.*
+.terraform.lock.hcl
+.terraformrc
+terraform.rc
diff --git a/infra/destroy_env.sh b/infra/destroy_env.sh
new file mode 100755
index 0000000..4e33049
--- /dev/null
+++ b/infra/destroy_env.sh
@@ -0,0 +1,43 @@
+#!/bin/bash
+
+###############################################################################
+# This script sets up a Prefect Cloud account with the following environment  #
+#                                                                             #
+# 1. Two workspaces: `production` and `staging`                               #
+# 2. A default Docker work pool in each workspace                             #
+# 3. A flow in each workspace                                                 #
+# 4. The flow in each workspace is run multiple times                         #
+# 5. The flow in `staging` has failures to demonstrate debugging              #
+#                                                                             #
+# NOTE: You must have Docker running on your machine to run this script!!!    #
+###############################################################################
+
+# Exit on any error
+set -e
+
+echo "🔑 Reading Prefect API key and account ID..."
+
+# Get active profile from `profiles.toml`
+ACTIVE_PROFILE=$(awk -F ' = ' '/^active/ {gsub(/"/, "", $2); print $2}' ~/.prefect/profiles.toml)
+
+# Get API key for the active profile from `profiles.toml`
+API_KEY=$(awk -v profile="profiles.$ACTIVE_PROFILE" '
+  $0 ~ "\\[" profile "\\]" {in_section=1; next}
+  in_section && /^\[/ {in_section=0}
+  in_section && /PREFECT_API_KEY/ {
+    gsub(/"/, "", $3)
+    print $3
+    exit
+  }
+' ~/.prefect/profiles.toml)
+export TF_VAR_prefect_api_key=$API_KEY
+
+# Extract account ID from `prefect config view`
+ACCOUNT_ID=$(prefect config view | awk -F'/' '/^https:\/\/app.prefect.cloud\/account\// {print $5}')
+export TF_VAR_prefect_account_id=$ACCOUNT_ID
+
+# Get account handle from the active workspace
+ACCOUNT_HANDLE=$(prefect cloud workspace ls | awk '/^│ \*/ {print $3}' | cut -d'/' -f1)
+
+echo "🏗️ Running Terraform to provision infrastructure..."
+terraform destroy -auto-approve
\ No newline at end of file
diff --git a/infra/main.tf b/infra/main.tf
new file mode 100644
index 0000000..f4d5056
--- /dev/null
+++ b/infra/main.tf
@@ -0,0 +1,38 @@
+terraform {
+  required_providers {
+    prefect = {
+      source = "PrefectHQ/prefect"
+    }
+  }
+}
+
+provider "prefect" {
+  api_key = var.prefect_api_key
+  account_id = var.prefect_account_id
+}
+
+# Create staging workspace
+resource "prefect_workspace" "staging" {
+  name        = "Staging"
+  handle      = "staging"
+}
+
+# Create production workspace
+resource "prefect_workspace" "production" {
+  name        = "Production"
+  handle      = "production"
+}
+
+# Create default work pool in staging workspace
+resource "prefect_work_pool" "staging_default" {
+  name        = "default-work-pool"
+  workspace_id = prefect_workspace.staging.id
+  type        = "docker"
+}
+
+# Create default work pool in production workspace
+resource "prefect_work_pool" "production_default" {
+  name        = "default-work-pool"
+  workspace_id = prefect_workspace.production.id
+  type        = "docker"
+} 
\ No newline at end of file
diff --git a/infra/variables.tf b/infra/variables.tf
new file mode 100644
index 0000000..5bcc6e4
--- /dev/null
+++ b/infra/variables.tf
@@ -0,0 +1,10 @@
+variable "prefect_api_key" {
+  description = "Prefect Cloud API key"
+  type        = string
+  sensitive   = true
+}
+
+variable "prefect_account_id" {
+  description = "Prefect Cloud Account ID"
+  type        = string
+} 
\ No newline at end of file
diff --git a/setup_env.sh b/setup_env.sh
new file mode 100755
index 0000000..430b0dd
--- /dev/null
+++ b/setup_env.sh
@@ -0,0 +1,103 @@
+#!/bin/bash
+
+###############################################################################
+# This script sets up a Prefect Cloud account with the following environment  #
+#                                                                             #
+# 1. Two workspaces: `production` and `staging`                               #
+# 2. A default Docker work pool in each workspace                             #
+# 3. A flow in each workspace                                                 #
+# 4. The flow in each workspace is run multiple times                         #
+# 5. The flow in `staging` has failures to demonstrate debugging              #
+#                                                                             #
+# NOTE: You must have Docker running on your machine to run this script!!!    #
+###############################################################################
+
+# Exit on any error
+set -e
+
+# Check if Docker is running
+echo "🐳 Checking if Docker is running..."
+if ! docker info > /dev/null 2>&1; then
+    echo "❌ Error: Docker is not running. Please start Docker and try again."
+    exit 1
+fi
+
+echo "✅ Docker is running"
+
+echo "🔑 Reading Prefect API key and account ID..."
+
+# Get active profile from `profiles.toml`
+ACTIVE_PROFILE=$(awk -F ' = ' '/^active/ {gsub(/"/, "", $2); print $2}' ~/.prefect/profiles.toml)
+
+# Get API key for the active profile from `profiles.toml`
+API_KEY=$(awk -v profile="profiles.$ACTIVE_PROFILE" '
+  $0 ~ "\\[" profile "\\]" {in_section=1; next}
+  in_section && /^\[/ {in_section=0}
+  in_section && /PREFECT_API_KEY/ {
+    gsub(/"/, "", $3)
+    print $3
+    exit
+  }
+' ~/.prefect/profiles.toml)
+export TF_VAR_prefect_api_key=$API_KEY
+
+# Extract account ID from `prefect config view`
+ACCOUNT_ID=$(prefect config view | awk -F'/' '/^https:\/\/app.prefect.cloud\/account\// {print $5}')
+export TF_VAR_prefect_account_id=$ACCOUNT_ID
+
+# Get account handle from the active workspace
+ACCOUNT_HANDLE=$(prefect cloud workspace ls | awk '/^│ \*/ {print $3}' | cut -d'/' -f1)
+
+echo "🏗️ Running Terraform to provision infrastructure..."
+cd infra/
+terraform init
+terraform apply -auto-approve
+cd ..
+
+###############################################################################
+
+echo "🚀 Populate production workspace..."
+
+# Start worker for production workspace
+prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/production"
+prefect worker start --pool "default-work-pool" &
+PROD_WORKER_PID=$!
+
+# Give workers time to start
+sleep 10
+
+# Run in production workspace
+python simulate_failures.py &
+PROD_SIM_PID=$!
+
+# Wait for simulations to complete
+wait $PROD_SIM_PID
+
+# Kill worker process
+kill $PROD_WORKER_PID
+
+###############################################################################
+
+echo "🚀 Populate staging workspace..."
+
+# Start worker for staging workspace
+prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/staging"
+prefect worker start --pool "default-work-pool" &
+STAGING_WORKER_PID=$!
+
+# Give workers time to start
+sleep 10
+
+# Run in staging workspace
+python simulate_failures.py --fail-at-run 3 &
+STAGING_SIM_PID=$!
+
+# Wait for simulations to complete
+wait $STAGING_SIM_PID
+
+# Kill worker process
+kill $STAGING_WORKER_PID
+
+###############################################################################
+
+echo "✅ All done!" 
\ No newline at end of file

From aeb3b637ce9c32bed6047b0e3fd869aa95208b10 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Mon, 9 Dec 2024 13:30:03 -0800
Subject: [PATCH 02/19] Add new lines at the end of files

---
 infra/destroy_env.sh | 2 +-
 infra/main.tf        | 2 +-
 infra/variables.tf   | 2 +-
 setup_env.sh         | 2 +-
 4 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/infra/destroy_env.sh b/infra/destroy_env.sh
index 4e33049..d08421f 100755
--- a/infra/destroy_env.sh
+++ b/infra/destroy_env.sh
@@ -40,4 +40,4 @@ export TF_VAR_prefect_account_id=$ACCOUNT_ID
 ACCOUNT_HANDLE=$(prefect cloud workspace ls | awk '/^│ \*/ {print $3}' | cut -d'/' -f1)
 
 echo "🏗️ Running Terraform to provision infrastructure..."
-terraform destroy -auto-approve
\ No newline at end of file
+terraform destroy -auto-approve
diff --git a/infra/main.tf b/infra/main.tf
index f4d5056..a22f6e5 100644
--- a/infra/main.tf
+++ b/infra/main.tf
@@ -35,4 +35,4 @@ resource "prefect_work_pool" "production_default" {
   name        = "default-work-pool"
   workspace_id = prefect_workspace.production.id
   type        = "docker"
-} 
\ No newline at end of file
+}
diff --git a/infra/variables.tf b/infra/variables.tf
index 5bcc6e4..5c9f471 100644
--- a/infra/variables.tf
+++ b/infra/variables.tf
@@ -7,4 +7,4 @@ variable "prefect_api_key" {
 variable "prefect_account_id" {
   description = "Prefect Cloud Account ID"
   type        = string
-} 
\ No newline at end of file
+}
diff --git a/setup_env.sh b/setup_env.sh
index 430b0dd..b870835 100755
--- a/setup_env.sh
+++ b/setup_env.sh
@@ -100,4 +100,4 @@ kill $STAGING_WORKER_PID
 
 ###############################################################################
 
-echo "✅ All done!" 
\ No newline at end of file
+echo "✅ All done!"

From 647147c0968f43583a7f74468b030a46356fce42 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Tue, 10 Dec 2024 08:28:32 -0800
Subject: [PATCH 03/19] Update misleading description in the 'destroy_env.sh'
 script

---
 infra/destroy_env.sh | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/infra/destroy_env.sh b/infra/destroy_env.sh
index d08421f..8d7e1c1 100755
--- a/infra/destroy_env.sh
+++ b/infra/destroy_env.sh
@@ -1,15 +1,7 @@
 #!/bin/bash
 
 ###############################################################################
-# This script sets up a Prefect Cloud account with the following environment  #
-#                                                                             #
-# 1. Two workspaces: `production` and `staging`                               #
-# 2. A default Docker work pool in each workspace                             #
-# 3. A flow in each workspace                                                 #
-# 4. The flow in each workspace is run multiple times                         #
-# 5. The flow in `staging` has failures to demonstrate debugging              #
-#                                                                             #
-# NOTE: You must have Docker running on your machine to run this script!!!    #
+# This script destroys any Prefect Cloud resources created by `setup_env.sh`  #
 ###############################################################################
 
 # Exit on any error

From 12d304cbbd8c32540bc30b7fdb69995542f8e493 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Tue, 10 Dec 2024 08:30:13 -0800
Subject: [PATCH 04/19] Update .gitignore to only include the min required
 Terraform files

---
 .gitignore | 2 --
 1 file changed, 2 deletions(-)

diff --git a/.gitignore b/.gitignore
index 415e09a..56254eb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,5 +8,3 @@ venv
 *.tfstate
 *.tfstate.*
 .terraform.lock.hcl
-.terraformrc
-terraform.rc

From 6800f1dae5a570c243143670e3fa8f7a3805f99f Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Tue, 10 Dec 2024 09:37:06 -0800
Subject: [PATCH 05/19] Use a more robust way of getting the account handle for
 the active profile

---
 setup_env.sh | 72 ++++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 67 insertions(+), 5 deletions(-)

diff --git a/setup_env.sh b/setup_env.sh
index b870835..7fb3a17 100755
--- a/setup_env.sh
+++ b/setup_env.sh
@@ -8,13 +8,17 @@
 # 3. A flow in each workspace                                                 #
 # 4. The flow in each workspace is run multiple times                         #
 # 5. The flow in `staging` has failures to demonstrate debugging              #
-#                                                                             #
-# NOTE: You must have Docker running on your machine to run this script!!!    #
+#
+# NOTE: You must have Docker and Terraform installed                          #
 ###############################################################################
 
 # Exit on any error
 set -e
 
+###############################################################################
+# Check for dependencies
+###############################################################################
+
 # Check if Docker is running
 echo "🐳 Checking if Docker is running..."
 if ! docker info > /dev/null 2>&1; then
@@ -24,8 +28,53 @@ fi
 
 echo "✅ Docker is running"
 
+# Check if Terraform is installed
+echo "🔧 Checking if Terraform is installed..."
+if ! command -v terraform &> /dev/null; then
+    echo "❌ Error: Terraform is not installed. Please install Terraform and try again."
+    exit 1
+fi
+
+echo "✅ Terraform is installed"
+
+# Check if Python is installed and determine the Python command
+echo "🐍 Checking if Python is installed..."
+if command -v python3 &> /dev/null; then
+    PYTHON_CMD="python3"
+elif command -v python &> /dev/null; then
+    PYTHON_CMD="python"
+else
+    echo "❌ Error: Python is not installed. Please install Python 3.9 or higher and try again."
+    exit 1
+fi
+
+# Verify Python version is 3.9 or higher
+if ! $PYTHON_CMD -c "import sys; assert sys.version_info >= (3, 9), 'Python 3.9 or higher is required'" &> /dev/null; then
+    echo "❌ Error: Python 3.9 or higher is required. Found $($PYTHON_CMD --version)"
+    exit 1
+fi
+
+echo "✅ Python $(${PYTHON_CMD} --version) is installed"
+
+###############################################################################
+# Set up virtual environment
+###############################################################################
+
+# Create and activate virtual environment
+echo "🌟 Setting up Python virtual environment..."
+$PYTHON_CMD -m venv temp_venv
+source temp_venv/bin/activate
+
+# Install requirements
+echo "📦 Installing Python packages..."
+pip install -r requirements.txt
+
 echo "🔑 Reading Prefect API key and account ID..."
 
+###############################################################################
+# Get auth credentials
+###############################################################################
+
 # Get active profile from `profiles.toml`
 ACTIVE_PROFILE=$(awk -F ' = ' '/^active/ {gsub(/"/, "", $2); print $2}' ~/.prefect/profiles.toml)
 
@@ -45,8 +94,12 @@ export TF_VAR_prefect_api_key=$API_KEY
 ACCOUNT_ID=$(prefect config view | awk -F'/' '/^https:\/\/app.prefect.cloud\/account\// {print $5}')
 export TF_VAR_prefect_account_id=$ACCOUNT_ID
 
-# Get account handle from the active workspace
-ACCOUNT_HANDLE=$(prefect cloud workspace ls | awk '/^│ \*/ {print $3}' | cut -d'/' -f1)
+# Get account handle for the account ID given above
+ACCOUNT_HANDLE=$(curl -s "https://api.prefect.cloud/api/accounts/$ACCOUNT_ID" -H "Authorization: Bearer $API_KEY" | awk -F'"handle":"' '{print $2}' | awk -F'"' '{print $1}')
+
+###############################################################################
+# Provision Prefect Cloud resources
+###############################################################################
 
 echo "🏗️ Running Terraform to provision infrastructure..."
 cd infra/
@@ -54,6 +107,8 @@ terraform init
 terraform apply -auto-approve
 cd ..
 
+###############################################################################
+# Run flows in production
 ###############################################################################
 
 echo "🚀 Populate production workspace..."
@@ -76,6 +131,8 @@ wait $PROD_SIM_PID
 # Kill worker process
 kill $PROD_WORKER_PID
 
+###############################################################################
+# Run flows in production
 ###############################################################################
 
 echo "🚀 Populate staging workspace..."
@@ -99,5 +156,10 @@ wait $STAGING_SIM_PID
 kill $STAGING_WORKER_PID
 
 ###############################################################################
+# Cleanup virtual environment
+###############################################################################
+
+deactivate
+rm -rf temp_venv
 
-echo "✅ All done!"
+echo "✅ All done!"
\ No newline at end of file

From d5cb89f8f3fe2379940be4313e928855340f5a5e Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Tue, 10 Dec 2024 09:43:57 -0800
Subject: [PATCH 06/19] Suppress Docker worker log output

---
 setup_env.sh | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/setup_env.sh b/setup_env.sh
index 7fb3a17..0e367bc 100755
--- a/setup_env.sh
+++ b/setup_env.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 
 ###############################################################################
-# This script sets up a Prefect Cloud account with the following environment  #
+# This script sets up a _paid_ Prefect Cloud account with resources:          #
 #                                                                             #
 # 1. Two workspaces: `production` and `staging`                               #
 # 2. A default Docker work pool in each workspace                             #
@@ -113,13 +113,13 @@ cd ..
 
 echo "🚀 Populate production workspace..."
 
-# Start worker for production workspace
+# Start worker for production workspace with suppressed output
 prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/production"
-prefect worker start --pool "default-work-pool" &
+prefect worker start --pool "default-work-pool" > /dev/null 2>&1 &
 PROD_WORKER_PID=$!
 
 # Give workers time to start
-sleep 10
+sleep 5
 
 # Run in production workspace
 python simulate_failures.py &
@@ -132,18 +132,18 @@ wait $PROD_SIM_PID
 kill $PROD_WORKER_PID
 
 ###############################################################################
-# Run flows in production
+# Run flows in staging
 ###############################################################################
 
 echo "🚀 Populate staging workspace..."
 
-# Start worker for staging workspace
+# Start worker for staging workspace with suppressed output
 prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/staging"
-prefect worker start --pool "default-work-pool" &
+prefect worker start --pool "default-work-pool" > /dev/null 2>&1 &
 STAGING_WORKER_PID=$!
 
 # Give workers time to start
-sleep 10
+sleep 5
 
 # Run in staging workspace
 python simulate_failures.py --fail-at-run 3 &

From 8c1e94de01da4275419c25d03499690d5a5c1a52 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Tue, 10 Dec 2024 09:45:17 -0800
Subject: [PATCH 07/19] Remove the account handle extraction from the
 destroy_env.sh script since it's not needed

---
 infra/destroy_env.sh | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/infra/destroy_env.sh b/infra/destroy_env.sh
index 8d7e1c1..cd94f92 100755
--- a/infra/destroy_env.sh
+++ b/infra/destroy_env.sh
@@ -28,8 +28,5 @@ export TF_VAR_prefect_api_key=$API_KEY
 ACCOUNT_ID=$(prefect config view | awk -F'/' '/^https:\/\/app.prefect.cloud\/account\// {print $5}')
 export TF_VAR_prefect_account_id=$ACCOUNT_ID
 
-# Get account handle from the active workspace
-ACCOUNT_HANDLE=$(prefect cloud workspace ls | awk '/^│ \*/ {print $3}' | cut -d'/' -f1)
-
 echo "🏗️ Running Terraform to provision infrastructure..."
 terraform destroy -auto-approve

From 8ad77a02aaaffad0077be337de6cd300b63653f5 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <djsauble@gmail.com>
Date: Tue, 10 Dec 2024 10:14:56 -0800
Subject: [PATCH 08/19] Confirm that this is a Prefect account that supports
 multiple workspaces

---
 setup_env.sh | 45 ++++++++++++++++++++++++++++-----------------
 1 file changed, 28 insertions(+), 17 deletions(-)

diff --git a/setup_env.sh b/setup_env.sh
index 0e367bc..4ffad86 100755
--- a/setup_env.sh
+++ b/setup_env.sh
@@ -56,21 +56,6 @@ fi
 
 echo "✅ Python $(${PYTHON_CMD} --version) is installed"
 
-###############################################################################
-# Set up virtual environment
-###############################################################################
-
-# Create and activate virtual environment
-echo "🌟 Setting up Python virtual environment..."
-$PYTHON_CMD -m venv temp_venv
-source temp_venv/bin/activate
-
-# Install requirements
-echo "📦 Installing Python packages..."
-pip install -r requirements.txt
-
-echo "🔑 Reading Prefect API key and account ID..."
-
 ###############################################################################
 # Get auth credentials
 ###############################################################################
@@ -94,8 +79,34 @@ export TF_VAR_prefect_api_key=$API_KEY
 ACCOUNT_ID=$(prefect config view | awk -F'/' '/^https:\/\/app.prefect.cloud\/account\// {print $5}')
 export TF_VAR_prefect_account_id=$ACCOUNT_ID
 
-# Get account handle for the account ID given above
-ACCOUNT_HANDLE=$(curl -s "https://api.prefect.cloud/api/accounts/$ACCOUNT_ID" -H "Authorization: Bearer $API_KEY" | awk -F'"handle":"' '{print $2}' | awk -F'"' '{print $1}')
+# Account details
+ACCOUNT_DETAILS=$(curl -s "https://api.prefect.cloud/api/accounts/$ACCOUNT_ID" -H "Authorization: Bearer $API_KEY")
+
+# Get account handle
+ACCOUNT_HANDLE=$(echo $ACCOUNT_DETAILS | awk -F'"handle":"' '{print $2}' | awk -F'"' '{print $1}')
+
+# Get plan type (fail if a personal account, since these do not support multiple workspaces)
+PLAN_TYPE=$(echo $ACCOUNT_DETAILS | awk -F'"plan_type":"' '{print $2}' | awk -F'"' '{print $1}')
+
+if [[ $PLAN_TYPE == "PERSONAL" ]]; then
+    echo "❌ Error: This script requires a paid Prefect Cloud account with support for multiple workspaces."
+    exit 1
+fi
+
+###############################################################################
+# Set up virtual environment
+###############################################################################
+
+# Create and activate virtual environment
+echo "🌟 Setting up Python virtual environment..."
+$PYTHON_CMD -m venv temp_venv
+source temp_venv/bin/activate
+
+# Install requirements
+echo "📦 Installing Python packages..."
+pip install -r requirements.txt
+
+echo "🔑 Reading Prefect API key and account ID..."
 
 ###############################################################################
 # Provision Prefect Cloud resources

From 070f06105b50f1fdbf90a8f0b40f34e3bcbf5fcb Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Wed, 15 Jan 2025 10:36:24 -0800
Subject: [PATCH 09/19] Restructure infra folder

---
 infra/{ => s3}/s3_to_prefect.tf               | 0
 infra/{ => workspaces}/destroy_env.sh         | 0
 infra/{ => workspaces}/main.tf                | 0
 setup_env.sh => infra/workspaces/setup_env.sh | 0
 infra/{ => workspaces}/variables.tf           | 0
 5 files changed, 0 insertions(+), 0 deletions(-)
 rename infra/{ => s3}/s3_to_prefect.tf (100%)
 rename infra/{ => workspaces}/destroy_env.sh (100%)
 rename infra/{ => workspaces}/main.tf (100%)
 rename setup_env.sh => infra/workspaces/setup_env.sh (100%)
 rename infra/{ => workspaces}/variables.tf (100%)

diff --git a/infra/s3_to_prefect.tf b/infra/s3/s3_to_prefect.tf
similarity index 100%
rename from infra/s3_to_prefect.tf
rename to infra/s3/s3_to_prefect.tf
diff --git a/infra/destroy_env.sh b/infra/workspaces/destroy_env.sh
similarity index 100%
rename from infra/destroy_env.sh
rename to infra/workspaces/destroy_env.sh
diff --git a/infra/main.tf b/infra/workspaces/main.tf
similarity index 100%
rename from infra/main.tf
rename to infra/workspaces/main.tf
diff --git a/setup_env.sh b/infra/workspaces/setup_env.sh
similarity index 100%
rename from setup_env.sh
rename to infra/workspaces/setup_env.sh
diff --git a/infra/variables.tf b/infra/workspaces/variables.tf
similarity index 100%
rename from infra/variables.tf
rename to infra/workspaces/variables.tf

From b0656fee99edef6c8f2775dda349142b17a6eac2 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Wed, 15 Jan 2025 11:11:31 -0800
Subject: [PATCH 10/19] Update paths

---
 infra/workspaces/setup_env.sh | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/infra/workspaces/setup_env.sh b/infra/workspaces/setup_env.sh
index 4ffad86..496fa62 100755
--- a/infra/workspaces/setup_env.sh
+++ b/infra/workspaces/setup_env.sh
@@ -104,7 +104,7 @@ source temp_venv/bin/activate
 
 # Install requirements
 echo "📦 Installing Python packages..."
-pip install -r requirements.txt
+pip install -r ../../requirements.txt
 
 echo "🔑 Reading Prefect API key and account ID..."
 
@@ -113,10 +113,8 @@ echo "🔑 Reading Prefect API key and account ID..."
 ###############################################################################
 
 echo "🏗️ Running Terraform to provision infrastructure..."
-cd infra/
 terraform init
 terraform apply -auto-approve
-cd ..
 
 ###############################################################################
 # Run flows in production
@@ -133,7 +131,7 @@ PROD_WORKER_PID=$!
 sleep 5
 
 # Run in production workspace
-python simulate_failures.py &
+python ../../simulate_failures.py &
 PROD_SIM_PID=$!
 
 # Wait for simulations to complete
@@ -157,7 +155,7 @@ STAGING_WORKER_PID=$!
 sleep 5
 
 # Run in staging workspace
-python simulate_failures.py --fail-at-run 3 &
+python ../../simulate_failures.py --fail-at-run 3 &
 STAGING_SIM_PID=$!
 
 # Wait for simulations to complete

From 685528c2f629c968933c5f4e8d30e3ecca0c0880 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 09:19:57 -0800
Subject: [PATCH 11/19] Refactor using modules

---
 .gitignore                                    |  1 +
 infra/workspaces/main.tf                      | 38 -------------------
 infra/workspaces/modules/workspace/main.tf    | 19 ++++++++++
 infra/workspaces/modules/workspace/outputs.tf | 14 +++++++
 .../workspaces/modules/workspace/variables.tf | 21 ++++++++++
 infra/workspaces/provision.tf                 | 37 ++++++++++++++++++
 infra/workspaces/variables.tf                 | 10 -----
 7 files changed, 92 insertions(+), 48 deletions(-)
 delete mode 100644 infra/workspaces/main.tf
 create mode 100644 infra/workspaces/modules/workspace/main.tf
 create mode 100644 infra/workspaces/modules/workspace/outputs.tf
 create mode 100644 infra/workspaces/modules/workspace/variables.tf
 create mode 100644 infra/workspaces/provision.tf
 delete mode 100644 infra/workspaces/variables.tf

diff --git a/.gitignore b/.gitignore
index 56254eb..0eb95e5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,6 +2,7 @@
 __pycache__
 env
 venv
+temp_venv
 
 # Terraform files
 .terraform
diff --git a/infra/workspaces/main.tf b/infra/workspaces/main.tf
deleted file mode 100644
index a22f6e5..0000000
--- a/infra/workspaces/main.tf
+++ /dev/null
@@ -1,38 +0,0 @@
-terraform {
-  required_providers {
-    prefect = {
-      source = "PrefectHQ/prefect"
-    }
-  }
-}
-
-provider "prefect" {
-  api_key = var.prefect_api_key
-  account_id = var.prefect_account_id
-}
-
-# Create staging workspace
-resource "prefect_workspace" "staging" {
-  name        = "Staging"
-  handle      = "staging"
-}
-
-# Create production workspace
-resource "prefect_workspace" "production" {
-  name        = "Production"
-  handle      = "production"
-}
-
-# Create default work pool in staging workspace
-resource "prefect_work_pool" "staging_default" {
-  name        = "default-work-pool"
-  workspace_id = prefect_workspace.staging.id
-  type        = "docker"
-}
-
-# Create default work pool in production workspace
-resource "prefect_work_pool" "production_default" {
-  name        = "default-work-pool"
-  workspace_id = prefect_workspace.production.id
-  type        = "docker"
-}
diff --git a/infra/workspaces/modules/workspace/main.tf b/infra/workspaces/modules/workspace/main.tf
new file mode 100644
index 0000000..3d60652
--- /dev/null
+++ b/infra/workspaces/modules/workspace/main.tf
@@ -0,0 +1,19 @@
+terraform {
+  required_providers {
+    prefect = {
+      source = "PrefectHQ/prefect"
+    }
+  }
+}
+
+# Module for creating a Prefect workspace and its default work pool
+resource "prefect_workspace" "workspace" {
+  name   = var.workspace_name
+  handle = var.workspace_handle
+}
+
+resource "prefect_work_pool" "default" {
+  name         = var.work_pool_name
+  workspace_id = prefect_workspace.workspace.id
+  type         = var.work_pool_type
+} 
\ No newline at end of file
diff --git a/infra/workspaces/modules/workspace/outputs.tf b/infra/workspaces/modules/workspace/outputs.tf
new file mode 100644
index 0000000..548e832
--- /dev/null
+++ b/infra/workspaces/modules/workspace/outputs.tf
@@ -0,0 +1,14 @@
+output "workspace_id" {
+  value       = prefect_workspace.workspace.id
+  description = "ID of the created workspace"
+}
+
+output "workspace_handle" {
+  value       = prefect_workspace.workspace.handle
+  description = "Handle of the created workspace"
+}
+
+output "work_pool_id" {
+  value       = prefect_work_pool.default.id
+  description = "ID of the created work pool"
+} 
\ No newline at end of file
diff --git a/infra/workspaces/modules/workspace/variables.tf b/infra/workspaces/modules/workspace/variables.tf
new file mode 100644
index 0000000..27c7f61
--- /dev/null
+++ b/infra/workspaces/modules/workspace/variables.tf
@@ -0,0 +1,21 @@
+variable "workspace_name" {
+  type        = string
+  description = "Name of the Prefect workspace"
+}
+
+variable "workspace_handle" {
+  type        = string
+  description = "Handle (slug) for the Prefect workspace"
+}
+
+variable "work_pool_name" {
+  type        = string
+  description = "Name of the default work pool"
+  default     = "my-work-pool"
+}
+
+variable "work_pool_type" {
+  type        = string
+  description = "Type of the work pool"
+  default     = "docker"
+} 
\ No newline at end of file
diff --git a/infra/workspaces/provision.tf b/infra/workspaces/provision.tf
new file mode 100644
index 0000000..e49b445
--- /dev/null
+++ b/infra/workspaces/provision.tf
@@ -0,0 +1,37 @@
+terraform {
+  required_providers {
+    prefect = {
+      source = "PrefectHQ/prefect"
+    }
+  }
+}
+
+provider "prefect" {
+  api_key     = var.prefect_api_key
+  account_id  = var.prefect_account_id
+}
+
+# Create staging environment
+module "staging" {
+  source          = "./modules/workspace"
+  workspace_name  = "Staging"
+  workspace_handle = "staging"
+}
+
+# Create production environment
+module "production" {
+  source          = "./modules/workspace"
+  workspace_name  = "Production"
+  workspace_handle = "production"
+}
+
+variable "prefect_api_key" {
+  description = "Prefect Cloud API key"
+  type        = string
+  sensitive   = true
+}
+
+variable "prefect_account_id" {
+  description = "Prefect Cloud Account ID"
+  type        = string
+}
\ No newline at end of file
diff --git a/infra/workspaces/variables.tf b/infra/workspaces/variables.tf
deleted file mode 100644
index 5c9f471..0000000
--- a/infra/workspaces/variables.tf
+++ /dev/null
@@ -1,10 +0,0 @@
-variable "prefect_api_key" {
-  description = "Prefect Cloud API key"
-  type        = string
-  sensitive   = true
-}
-
-variable "prefect_account_id" {
-  description = "Prefect Cloud Account ID"
-  type        = string
-}

From 2be0fd2d7b574479f789666da9629ddb016d2b18 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 09:52:33 -0800
Subject: [PATCH 12/19] Merge variables into the module file

---
 infra/workspaces/modules/workspace/main.tf    | 24 ++++++++++++++++++-
 .../workspaces/modules/workspace/variables.tf | 21 ----------------
 2 files changed, 23 insertions(+), 22 deletions(-)
 delete mode 100644 infra/workspaces/modules/workspace/variables.tf

diff --git a/infra/workspaces/modules/workspace/main.tf b/infra/workspaces/modules/workspace/main.tf
index 3d60652..13a5229 100644
--- a/infra/workspaces/modules/workspace/main.tf
+++ b/infra/workspaces/modules/workspace/main.tf
@@ -16,4 +16,26 @@ resource "prefect_work_pool" "default" {
   name         = var.work_pool_name
   workspace_id = prefect_workspace.workspace.id
   type         = var.work_pool_type
-} 
\ No newline at end of file
+}
+
+variable "workspace_name" {
+  type        = string
+  description = "Name of the Prefect workspace"
+}
+
+variable "workspace_handle" {
+  type        = string
+  description = "Handle (slug) for the Prefect workspace"
+}
+
+variable "work_pool_name" {
+  type        = string
+  description = "Name of the default work pool"
+  default     = "my-work-pool"
+}
+
+variable "work_pool_type" {
+  type        = string
+  description = "Type of the work pool"
+  default     = "docker"
+}
diff --git a/infra/workspaces/modules/workspace/variables.tf b/infra/workspaces/modules/workspace/variables.tf
deleted file mode 100644
index 27c7f61..0000000
--- a/infra/workspaces/modules/workspace/variables.tf
+++ /dev/null
@@ -1,21 +0,0 @@
-variable "workspace_name" {
-  type        = string
-  description = "Name of the Prefect workspace"
-}
-
-variable "workspace_handle" {
-  type        = string
-  description = "Handle (slug) for the Prefect workspace"
-}
-
-variable "work_pool_name" {
-  type        = string
-  description = "Name of the default work pool"
-  default     = "my-work-pool"
-}
-
-variable "work_pool_type" {
-  type        = string
-  description = "Type of the work pool"
-  default     = "docker"
-} 
\ No newline at end of file

From acae843c16e80af07daab9b59ccf61ca9d56f117 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 09:52:58 -0800
Subject: [PATCH 13/19] Revert "Merge variables into the module file"

This reverts commit 2be0fd2d7b574479f789666da9629ddb016d2b18.
---
 infra/workspaces/modules/workspace/main.tf    | 24 +------------------
 .../workspaces/modules/workspace/variables.tf | 21 ++++++++++++++++
 2 files changed, 22 insertions(+), 23 deletions(-)
 create mode 100644 infra/workspaces/modules/workspace/variables.tf

diff --git a/infra/workspaces/modules/workspace/main.tf b/infra/workspaces/modules/workspace/main.tf
index 13a5229..3d60652 100644
--- a/infra/workspaces/modules/workspace/main.tf
+++ b/infra/workspaces/modules/workspace/main.tf
@@ -16,26 +16,4 @@ resource "prefect_work_pool" "default" {
   name         = var.work_pool_name
   workspace_id = prefect_workspace.workspace.id
   type         = var.work_pool_type
-}
-
-variable "workspace_name" {
-  type        = string
-  description = "Name of the Prefect workspace"
-}
-
-variable "workspace_handle" {
-  type        = string
-  description = "Handle (slug) for the Prefect workspace"
-}
-
-variable "work_pool_name" {
-  type        = string
-  description = "Name of the default work pool"
-  default     = "my-work-pool"
-}
-
-variable "work_pool_type" {
-  type        = string
-  description = "Type of the work pool"
-  default     = "docker"
-}
+} 
\ No newline at end of file
diff --git a/infra/workspaces/modules/workspace/variables.tf b/infra/workspaces/modules/workspace/variables.tf
new file mode 100644
index 0000000..27c7f61
--- /dev/null
+++ b/infra/workspaces/modules/workspace/variables.tf
@@ -0,0 +1,21 @@
+variable "workspace_name" {
+  type        = string
+  description = "Name of the Prefect workspace"
+}
+
+variable "workspace_handle" {
+  type        = string
+  description = "Handle (slug) for the Prefect workspace"
+}
+
+variable "work_pool_name" {
+  type        = string
+  description = "Name of the default work pool"
+  default     = "my-work-pool"
+}
+
+variable "work_pool_type" {
+  type        = string
+  description = "Type of the work pool"
+  default     = "docker"
+} 
\ No newline at end of file

From 429cb9076c2c8598ebd3f88d8f855467dcc57724 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 10:00:50 -0800
Subject: [PATCH 14/19] Add newlines to the end of files

---
 infra/workspaces/modules/workspace/main.tf      | 2 +-
 infra/workspaces/modules/workspace/outputs.tf   | 2 +-
 infra/workspaces/modules/workspace/variables.tf | 2 +-
 infra/workspaces/provision.tf                   | 2 +-
 infra/workspaces/setup_env.sh                   | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/infra/workspaces/modules/workspace/main.tf b/infra/workspaces/modules/workspace/main.tf
index 3d60652..d9ec739 100644
--- a/infra/workspaces/modules/workspace/main.tf
+++ b/infra/workspaces/modules/workspace/main.tf
@@ -16,4 +16,4 @@ resource "prefect_work_pool" "default" {
   name         = var.work_pool_name
   workspace_id = prefect_workspace.workspace.id
   type         = var.work_pool_type
-} 
\ No newline at end of file
+}
diff --git a/infra/workspaces/modules/workspace/outputs.tf b/infra/workspaces/modules/workspace/outputs.tf
index 548e832..81f7cc3 100644
--- a/infra/workspaces/modules/workspace/outputs.tf
+++ b/infra/workspaces/modules/workspace/outputs.tf
@@ -11,4 +11,4 @@ output "workspace_handle" {
 output "work_pool_id" {
   value       = prefect_work_pool.default.id
   description = "ID of the created work pool"
-} 
\ No newline at end of file
+}
diff --git a/infra/workspaces/modules/workspace/variables.tf b/infra/workspaces/modules/workspace/variables.tf
index 27c7f61..cd74eea 100644
--- a/infra/workspaces/modules/workspace/variables.tf
+++ b/infra/workspaces/modules/workspace/variables.tf
@@ -18,4 +18,4 @@ variable "work_pool_type" {
   type        = string
   description = "Type of the work pool"
   default     = "docker"
-} 
\ No newline at end of file
+}
diff --git a/infra/workspaces/provision.tf b/infra/workspaces/provision.tf
index e49b445..b50c744 100644
--- a/infra/workspaces/provision.tf
+++ b/infra/workspaces/provision.tf
@@ -34,4 +34,4 @@ variable "prefect_api_key" {
 variable "prefect_account_id" {
   description = "Prefect Cloud Account ID"
   type        = string
-}
\ No newline at end of file
+}
diff --git a/infra/workspaces/setup_env.sh b/infra/workspaces/setup_env.sh
index 496fa62..1972dd1 100755
--- a/infra/workspaces/setup_env.sh
+++ b/infra/workspaces/setup_env.sh
@@ -171,4 +171,4 @@ kill $STAGING_WORKER_PID
 deactivate
 rm -rf temp_venv
 
-echo "✅ All done!"
\ No newline at end of file
+echo "✅ All done!"

From 5f31540cffc54c818ccbb203b24c141b8dd9da7c Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 10:49:37 -0800
Subject: [PATCH 15/19] Remove ugly awk logic and update work pool names

---
 infra/workspaces/setup_env.sh | 47 ++++++++++++++---------------------
 simulate_failures.py          |  2 +-
 2 files changed, 20 insertions(+), 29 deletions(-)

diff --git a/infra/workspaces/setup_env.sh b/infra/workspaces/setup_env.sh
index 1972dd1..15b02b2 100755
--- a/infra/workspaces/setup_env.sh
+++ b/infra/workspaces/setup_env.sh
@@ -4,7 +4,7 @@
 # This script sets up a _paid_ Prefect Cloud account with resources:          #
 #                                                                             #
 # 1. Two workspaces: `production` and `staging`                               #
-# 2. A default Docker work pool in each workspace                             #
+# 2. A Docker work pool in each workspace                             #
 # 3. A flow in each workspace                                                 #
 # 4. The flow in each workspace is run multiple times                         #
 # 5. The flow in `staging` has failures to demonstrate debugging              #
@@ -57,36 +57,29 @@ fi
 echo "✅ Python $(${PYTHON_CMD} --version) is installed"
 
 ###############################################################################
-# Get auth credentials
+# Get account details
 ###############################################################################
 
-# Get active profile from `profiles.toml`
-ACTIVE_PROFILE=$(awk -F ' = ' '/^active/ {gsub(/"/, "", $2); print $2}' ~/.prefect/profiles.toml)
+echo "🔑 Fetching Prefect account details..."
 
-# Get API key for the active profile from `profiles.toml`
-API_KEY=$(awk -v profile="profiles.$ACTIVE_PROFILE" '
-  $0 ~ "\\[" profile "\\]" {in_section=1; next}
-  in_section && /^\[/ {in_section=0}
-  in_section && /PREFECT_API_KEY/ {
-    gsub(/"/, "", $3)
-    print $3
-    exit
-  }
-' ~/.prefect/profiles.toml)
-export TF_VAR_prefect_api_key=$API_KEY
+# Must have set TF_VAR_prefect_api_key and TF_VAR_prefect_account_id environment variables
+if [ -z "$TF_VAR_prefect_api_key" ]; then
+    echo "❌ Error: TF_VAR_prefect_api_key environment variable is not set"
+    exit 1
+fi
 
-# Extract account ID from `prefect config view`
-ACCOUNT_ID=$(prefect config view | awk -F'/' '/^https:\/\/app.prefect.cloud\/account\// {print $5}')
-export TF_VAR_prefect_account_id=$ACCOUNT_ID
+if [ -z "$TF_VAR_prefect_account_id" ]; then
+    echo "❌ Error: TF_VAR_prefect_account_id environment variable is not set"
+    exit 1
+fi
 
 # Account details
-ACCOUNT_DETAILS=$(curl -s "https://api.prefect.cloud/api/accounts/$ACCOUNT_ID" -H "Authorization: Bearer $API_KEY")
-
-# Get account handle
-ACCOUNT_HANDLE=$(echo $ACCOUNT_DETAILS | awk -F'"handle":"' '{print $2}' | awk -F'"' '{print $1}')
+ACCOUNT_DETAILS=$(curl -s "https://api.prefect.cloud/api/accounts/$TF_VAR_prefect_account_id" \
+    -H "Authorization: Bearer $TF_VAR_prefect_api_key")
 
-# Get plan type (fail if a personal account, since these do not support multiple workspaces)
-PLAN_TYPE=$(echo $ACCOUNT_DETAILS | awk -F'"plan_type":"' '{print $2}' | awk -F'"' '{print $1}')
+# Get account handle and plan type using jq
+ACCOUNT_HANDLE=$(echo "$ACCOUNT_DETAILS" | jq -r '.handle')
+PLAN_TYPE=$(echo "$ACCOUNT_DETAILS" | jq -r '.plan_type')
 
 if [[ $PLAN_TYPE == "PERSONAL" ]]; then
     echo "❌ Error: This script requires a paid Prefect Cloud account with support for multiple workspaces."
@@ -106,8 +99,6 @@ source temp_venv/bin/activate
 echo "📦 Installing Python packages..."
 pip install -r ../../requirements.txt
 
-echo "🔑 Reading Prefect API key and account ID..."
-
 ###############################################################################
 # Provision Prefect Cloud resources
 ###############################################################################
@@ -124,7 +115,7 @@ echo "🚀 Populate production workspace..."
 
 # Start worker for production workspace with suppressed output
 prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/production"
-prefect worker start --pool "default-work-pool" > /dev/null 2>&1 &
+prefect worker start --pool "my-work-pool" > /dev/null 2>&1 &
 PROD_WORKER_PID=$!
 
 # Give workers time to start
@@ -148,7 +139,7 @@ echo "🚀 Populate staging workspace..."
 
 # Start worker for staging workspace with suppressed output
 prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/staging"
-prefect worker start --pool "default-work-pool" > /dev/null 2>&1 &
+prefect worker start --pool "my-work-pool" > /dev/null 2>&1 &
 STAGING_WORKER_PID=$!
 
 # Give workers time to start
diff --git a/simulate_failures.py b/simulate_failures.py
index 02cafb7..ba55cf6 100644
--- a/simulate_failures.py
+++ b/simulate_failures.py
@@ -50,7 +50,7 @@ async def create_runs(deployment_id: str, num_runs: int, fail_at_run: int | None
     # Deploy the flow
     deployment_id = data_pipeline.deploy(
         name=args.name,
-        work_pool_name="default-work-pool",
+        work_pool_name="my-work-pool",
         image="prefecthq/prefect:3-latest",
         push=False,
         tags=args.tags.split(',')

From 00263f20a572b709da93c7b1c58da2c5a9be8e5e Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 11:09:50 -0800
Subject: [PATCH 16/19] Let users customize the names of the production and
 staging workspaces

---
 infra/workspaces/destroy_env.sh               | 32 -------------------
 .../workspaces/modules/workspace/variables.tf |  2 +-
 infra/workspaces/provision.tf                 | 16 ++++++++--
 infra/workspaces/setup_env.sh                 | 24 +++++++++-----
 4 files changed, 31 insertions(+), 43 deletions(-)
 delete mode 100755 infra/workspaces/destroy_env.sh

diff --git a/infra/workspaces/destroy_env.sh b/infra/workspaces/destroy_env.sh
deleted file mode 100755
index cd94f92..0000000
--- a/infra/workspaces/destroy_env.sh
+++ /dev/null
@@ -1,32 +0,0 @@
-#!/bin/bash
-
-###############################################################################
-# This script destroys any Prefect Cloud resources created by `setup_env.sh`  #
-###############################################################################
-
-# Exit on any error
-set -e
-
-echo "🔑 Reading Prefect API key and account ID..."
-
-# Get active profile from `profiles.toml`
-ACTIVE_PROFILE=$(awk -F ' = ' '/^active/ {gsub(/"/, "", $2); print $2}' ~/.prefect/profiles.toml)
-
-# Get API key for the active profile from `profiles.toml`
-API_KEY=$(awk -v profile="profiles.$ACTIVE_PROFILE" '
-  $0 ~ "\\[" profile "\\]" {in_section=1; next}
-  in_section && /^\[/ {in_section=0}
-  in_section && /PREFECT_API_KEY/ {
-    gsub(/"/, "", $3)
-    print $3
-    exit
-  }
-' ~/.prefect/profiles.toml)
-export TF_VAR_prefect_api_key=$API_KEY
-
-# Extract account ID from `prefect config view`
-ACCOUNT_ID=$(prefect config view | awk -F'/' '/^https:\/\/app.prefect.cloud\/account\// {print $5}')
-export TF_VAR_prefect_account_id=$ACCOUNT_ID
-
-echo "🏗️ Running Terraform to provision infrastructure..."
-terraform destroy -auto-approve
diff --git a/infra/workspaces/modules/workspace/variables.tf b/infra/workspaces/modules/workspace/variables.tf
index cd74eea..0fd0cb6 100644
--- a/infra/workspaces/modules/workspace/variables.tf
+++ b/infra/workspaces/modules/workspace/variables.tf
@@ -1,6 +1,6 @@
 variable "workspace_name" {
+  description = "Name of the workspace to create"
   type        = string
-  description = "Name of the Prefect workspace"
 }
 
 variable "workspace_handle" {
diff --git a/infra/workspaces/provision.tf b/infra/workspaces/provision.tf
index b50c744..8d5bbda 100644
--- a/infra/workspaces/provision.tf
+++ b/infra/workspaces/provision.tf
@@ -15,14 +15,14 @@ provider "prefect" {
 module "staging" {
   source          = "./modules/workspace"
   workspace_name  = "Staging"
-  workspace_handle = "staging"
+  workspace_handle = var.staging_workspace
 }
 
 # Create production environment
 module "production" {
   source          = "./modules/workspace"
   workspace_name  = "Production"
-  workspace_handle = "production"
+  workspace_handle = var.prod_workspace
 }
 
 variable "prefect_api_key" {
@@ -35,3 +35,15 @@ variable "prefect_account_id" {
   description = "Prefect Cloud Account ID"
   type        = string
 }
+
+variable "prod_workspace" {
+  description = "Name of the production workspace"
+  type        = string
+  default     = "production"
+}
+
+variable "staging_workspace" {
+  description = "Name of the staging workspace"
+  type        = string
+  default     = "staging"
+}
diff --git a/infra/workspaces/setup_env.sh b/infra/workspaces/setup_env.sh
index 15b02b2..d185649 100755
--- a/infra/workspaces/setup_env.sh
+++ b/infra/workspaces/setup_env.sh
@@ -3,12 +3,12 @@
 ###############################################################################
 # This script sets up a _paid_ Prefect Cloud account with resources:          #
 #                                                                             #
-# 1. Two workspaces: `production` and `staging`                               #
-# 2. A Docker work pool in each workspace                             #
+# 1. Two workspaces: `production` and `staging` (customizable via env vars)   #
+# 2. A Docker work pool in each workspace                                     #
 # 3. A flow in each workspace                                                 #
 # 4. The flow in each workspace is run multiple times                         #
 # 5. The flow in `staging` has failures to demonstrate debugging              #
-#
+#                                                                             #
 # NOTE: You must have Docker and Terraform installed                          #
 ###############################################################################
 
@@ -57,7 +57,7 @@ fi
 echo "✅ Python $(${PYTHON_CMD} --version) is installed"
 
 ###############################################################################
-# Get account details
+# Establish account and workspace details
 ###############################################################################
 
 echo "🔑 Fetching Prefect account details..."
@@ -73,6 +73,14 @@ if [ -z "$TF_VAR_prefect_account_id" ]; then
     exit 1
 fi
 
+# Set default workspace names if not provided via environment variables
+PROD_WORKSPACE=${TF_VAR_prod_workspace:-"production"}
+STAGING_WORKSPACE=${TF_VAR_staging_workspace:-"staging"}
+
+# Export for Terraform to use
+export TF_VAR_prod_workspace=$PROD_WORKSPACE
+export TF_VAR_staging_workspace=$STAGING_WORKSPACE
+
 # Account details
 ACCOUNT_DETAILS=$(curl -s "https://api.prefect.cloud/api/accounts/$TF_VAR_prefect_account_id" \
     -H "Authorization: Bearer $TF_VAR_prefect_api_key")
@@ -111,10 +119,10 @@ terraform apply -auto-approve
 # Run flows in production
 ###############################################################################
 
-echo "🚀 Populate production workspace..."
+echo "🚀 Populate $PROD_WORKSPACE workspace..."
 
 # Start worker for production workspace with suppressed output
-prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/production"
+prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/$PROD_WORKSPACE"
 prefect worker start --pool "my-work-pool" > /dev/null 2>&1 &
 PROD_WORKER_PID=$!
 
@@ -135,10 +143,10 @@ kill $PROD_WORKER_PID
 # Run flows in staging
 ###############################################################################
 
-echo "🚀 Populate staging workspace..."
+echo "🚀 Populate $STAGING_WORKSPACE workspace..."
 
 # Start worker for staging workspace with suppressed output
-prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/staging"
+prefect cloud workspace set --workspace "$ACCOUNT_HANDLE/$STAGING_WORKSPACE"
 prefect worker start --pool "my-work-pool" > /dev/null 2>&1 &
 STAGING_WORKER_PID=$!
 

From fb2c95ee4c0ade0a1a20b5b1b744ec6f696cb974 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 11:25:06 -0800
Subject: [PATCH 17/19] Confirm that jq is installed

---
 infra/workspaces/setup_env.sh | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/infra/workspaces/setup_env.sh b/infra/workspaces/setup_env.sh
index d185649..ffbe0ea 100755
--- a/infra/workspaces/setup_env.sh
+++ b/infra/workspaces/setup_env.sh
@@ -56,6 +56,15 @@ fi
 
 echo "✅ Python $(${PYTHON_CMD} --version) is installed"
 
+# Check if jq is installed
+echo "🔧 Checking if jq is installed..."
+if ! command -v jq &> /dev/null; then
+    echo "❌ Error: jq is not installed. Please install jq and try again."
+    exit 1
+fi
+
+echo "✅ jq is installed"
+
 ###############################################################################
 # Establish account and workspace details
 ###############################################################################

From 199c901d5bf6756da23fc185a124ccafe82d8bd6 Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 11:40:35 -0800
Subject: [PATCH 18/19] Move setup_env.sh back to the root since otherwise
 simulate_failures.py will fail

---
 infra/workspaces/setup_env.sh => setup_env.sh | 40 +++++++++++++------
 1 file changed, 28 insertions(+), 12 deletions(-)
 rename infra/workspaces/setup_env.sh => setup_env.sh (88%)

diff --git a/infra/workspaces/setup_env.sh b/setup_env.sh
similarity index 88%
rename from infra/workspaces/setup_env.sh
rename to setup_env.sh
index ffbe0ea..0139de6 100755
--- a/infra/workspaces/setup_env.sh
+++ b/setup_env.sh
@@ -15,6 +15,29 @@
 # Exit on any error
 set -e
 
+cleanup() {
+
+    # Kill any remaining worker processes
+    if [ ! -z "$PROD_WORKER_PID" ]; then
+        kill $PROD_WORKER_PID 2>/dev/null || true
+    fi
+    if [ ! -z "$STAGING_WORKER_PID" ]; then
+        kill $STAGING_WORKER_PID 2>/dev/null || true
+    fi
+
+    # Deactivate and remove virtual environment
+    if [ -d "temp_venv" ]; then
+        deactivate 2>/dev/null || true
+        rm -rf temp_venv
+    fi
+
+    echo "🧹 Cleanup completed"
+
+}
+
+# Set up trap to call cleanup function on script exit (success or failure)
+trap cleanup EXIT
+
 ###############################################################################
 # Check for dependencies
 ###############################################################################
@@ -114,15 +137,15 @@ source temp_venv/bin/activate
 
 # Install requirements
 echo "📦 Installing Python packages..."
-pip install -r ../../requirements.txt
+pip install -r ./requirements.txt
 
 ###############################################################################
 # Provision Prefect Cloud resources
 ###############################################################################
 
 echo "🏗️ Running Terraform to provision infrastructure..."
-terraform init
-terraform apply -auto-approve
+terraform -chdir=infra/workspaces init
+terraform -chdir=infra/workspaces apply -auto-approve
 
 ###############################################################################
 # Run flows in production
@@ -139,7 +162,7 @@ PROD_WORKER_PID=$!
 sleep 5
 
 # Run in production workspace
-python ../../simulate_failures.py &
+python ./simulate_failures.py &
 PROD_SIM_PID=$!
 
 # Wait for simulations to complete
@@ -163,7 +186,7 @@ STAGING_WORKER_PID=$!
 sleep 5
 
 # Run in staging workspace
-python ../../simulate_failures.py --fail-at-run 3 &
+python ./simulate_failures.py --fail-at-run 3 &
 STAGING_SIM_PID=$!
 
 # Wait for simulations to complete
@@ -172,11 +195,4 @@ wait $STAGING_SIM_PID
 # Kill worker process
 kill $STAGING_WORKER_PID
 
-###############################################################################
-# Cleanup virtual environment
-###############################################################################
-
-deactivate
-rm -rf temp_venv
-
 echo "✅ All done!"

From ea7774ba8a9e7eefdd7321fb64648f68d4ae156c Mon Sep 17 00:00:00 2001
From: Daniel Sauble <daniel@prefect.io>
Date: Thu, 16 Jan 2025 11:46:21 -0800
Subject: [PATCH 19/19] Remove extra kill statements, since this is done
 automatically during cleanup

---
 setup_env.sh | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/setup_env.sh b/setup_env.sh
index 0139de6..f8f885f 100755
--- a/setup_env.sh
+++ b/setup_env.sh
@@ -168,9 +168,6 @@ PROD_SIM_PID=$!
 # Wait for simulations to complete
 wait $PROD_SIM_PID
 
-# Kill worker process
-kill $PROD_WORKER_PID
-
 ###############################################################################
 # Run flows in staging
 ###############################################################################
@@ -192,7 +189,4 @@ STAGING_SIM_PID=$!
 # Wait for simulations to complete
 wait $STAGING_SIM_PID
 
-# Kill worker process
-kill $STAGING_WORKER_PID
-
 echo "✅ All done!"