Merge pull request #30 from Snowflake-Labs/sovereign-cloud-support

Add support for sovereign cloud.

iam.tf

@@ -15,11 +15,11 @@ resource "aws_iam_role" "s3_reader" {
         Action = "sts:AssumeRole"
         Effect = "Allow"
         Principal = {
-          AWS = snowflake_storage_integration.this.storage_aws_iam_user_arn
+          AWS = local.storage_integration_user_arn
         }
         Condition = {
           StringEquals = {
-            "sts:ExternalId" = snowflake_storage_integration.this.storage_aws_external_id
+            "sts:ExternalId" = local.storage_integration_external_id
           }
         }
       }

output.tf

@@ -1,11 +1,11 @@
 output "storage_integration_name" {
   description = "Name of Storage integration"
-  value       = snowflake_storage_integration.this.name
+  value       = local.storage_integration_name
 }
 
 output "bucket_url" {
   description = "GEFF S3 Bucket URL"
-  value       = var.arn_format == "aws-us-gov" ? "s3gov://${aws_s3_bucket.geff_bucket.id}/" : "s3://${aws_s3_bucket.geff_bucket.id}/"
+  value       = "s3://${aws_s3_bucket.geff_bucket.id}/"
 }
 
 output "bucket_arn" {

s3.tf

@@ -64,7 +64,7 @@ data "aws_iam_policy_document" "geff_s3_sns_topic_policy_doc" {
 
     principals {
       type        = "AWS"
-      identifiers = [snowflake_storage_integration.this.storage_aws_iam_user_arn]
+      identifiers = [local.storage_integration_user_arn]
     }
   }
 }

storage_integration.tf

@@ -1,30 +1,76 @@
 locals {
+  storage_provider_map        = lookup(local.snowflake_storage_provider_maps, local.aws_partition, null)
+  snowflake_storage_provider  = local.storage_provider_map["snowflake_storage_provider"]
+  terraform_resource_provider = local.storage_provider_map["terraform_resource_provider"]
+
+  storage_integration_name = "${upper(replace(var.prefix, "-", "_"))}_STORAGE_INTEGRATION"
+
   pipeline_bucket_ids = [
     for bucket_arn in var.data_bucket_arns : element(split(":::", bucket_arn), 1)
   ]
-  storage_provider = length(regexall(".*gov.*", local.aws_region)) > 0 ? "S3GOV" : "S3"
+  storage_allowed_locations = concat(
+    ["${local.snowflake_storage_provider}://${aws_s3_bucket.geff_bucket.id}/"],
+    [for bucket_id in local.pipeline_bucket_ids : "${local.snowflake_storage_provider}://${bucket_id}/"]
+  )
+
+  storage_allowed_locations_snowsql = join(",", [for i in local.storage_allowed_locations : join("", ["'", i, "'"])])
 }
 
 resource "snowflake_storage_integration" "this" {
+  count    = local.terraform_resource_provider == "snowflake" ? 1 : 0
   provider = snowflake.storage_integration_role
 
-  name    = "${upper(replace(var.prefix, "-", "_"))}_STORAGE_INTEGRATION"
+  name    = local.storage_integration_name
   type    = "EXTERNAL_STAGE"
   enabled = true
-  storage_allowed_locations = concat(
-    ["${local.storage_provider}://${aws_s3_bucket.geff_bucket.id}/"],
-    [for bucket_id in local.pipeline_bucket_ids : "s3://${bucket_id}/"]
-  )
-  storage_provider     = local.storage_provider
-  storage_aws_role_arn = "arn:${var.arn_format}:iam::${local.account_id}:role/${local.s3_reader_role_name}"
+
+  storage_allowed_locations = local.storage_allowed_locations
+  storage_provider          = local.snowflake_storage_provider
+  storage_aws_role_arn      = "arn:${local.aws_partition}:iam::${local.account_id}:role/${local.s3_reader_role_name}"
+}
+
+## Create Snowflake storage integration with SnowSQL Terraform provider if the official Snowflake Terraform provider not yet support the specific sovereign cloud.
+resource "snowsql_exec" "snowflake_storage_integration" {
+  count    = local.terraform_resource_provider == "snowsql" ? 1 : 0
+  provider = snowsql.storage_integration_role
+
+  create {
+    statements = <<-EOT
+      CREATE OR REPLACE STORAGE INTEGRATION "${local.storage_integration_name}"
+	TYPE=EXTERNAL_STAGE
+	STORAGE_PROVIDER='${local.snowflake_storage_provider}'
+	STORAGE_AWS_ROLE_ARN="arn:${local.aws_partition}:iam::${local.account_id}:role/${local.s3_reader_role_name}"
+	ENABLED=true
+	STORAGE_ALLOWED_LOCATIONS=(${local.storage_allowed_locations_snowsql});
+    EOT
+  }
+
+  read {
+    statements = "DESCRIBE STORAGE INTEGRATION ${local.storage_integration_name};"
+  }
+
+  delete {
+    statements = "DROP INTEGRATION ${local.storage_integration_name};"
+  }
+}
+
+locals {
+  storage_integration_user_arn = local.terraform_resource_provider == "snowflake" ? snowflake_storage_integration.this[0].storage_aws_iam_user_arn : [for map in jsondecode(nonsensitive(snowsql_exec.snowflake_storage_integration[0].read_results)): map if map.property == "STORAGE_AWS_IAM_USER_ARN"][0]["property_value"]
+
+  storage_integration_external_id = local.terraform_resource_provider == "snowflake" ? snowflake_storage_integration.this[0].storage_aws_external_id : [for map in jsondecode(nonsensitive(snowsql_exec.snowflake_storage_integration[0].read_results)): map if map.property == "STORAGE_AWS_EXTERNAL_ID"][0]["property_value"]
 }
 
 resource "snowflake_integration_grant" "this" {
   provider         = snowflake.storage_integration_role
-  integration_name = snowflake_storage_integration.this.name
+  integration_name = local.storage_integration_name
 
   privilege = "USAGE"
   roles     = var.snowflake_integration_user_roles
 
   with_grant_option = false
+
+  depends_on = [
+    snowflake_storage_integration.this,
+    snowsql_exec.snowflake_storage_integration,
+  ]
 }

variables.tf

@@ -34,12 +34,6 @@ variable "data_bucket_arns" {
   description = "List of Bucket ARNs for the s3_reader role to read from."
 }
 
-variable "arn_format" {
-  type        = string
-  description = "ARN format could be aws or aws-us-gov. Defaults to non-gov."
-  default     = "aws"
-}
-
 variable "bucket_object_ownership_settings" {
   type        = string
   description = "The settings that will impact ACLs and ownership of objects within the bucket."
@@ -52,11 +46,27 @@ data "aws_region" "current" {}
 data "aws_partition" "current" {}
 
 locals {
-  account_id = data.aws_caller_identity.current.account_id
-  aws_region = data.aws_region.current.name
-}
+  account_id     = data.aws_caller_identity.current.account_id
+  aws_region     = data.aws_region.current.name
+  aws_partition  = data.aws_partition.current.partition
+  aws_dns_suffix = data.aws_partition.current.dns_suffix
+
+  # Use SnowSQL Terraform provider if the official Snowflake Terraform provider does not support the specific cloud.
+  snowflake_storage_provider_maps = {
+    aws = {
+      snowflake_storage_provider  = "S3"
+      terraform_resource_provider = "snowflake"
+    }
+    aws-us-gov = {
+      snowflake_storage_provider  = "S3GOV"
+      terraform_resource_provider = "snowflake"
+    }
+    aws-cn = {
+      snowflake_storage_provider  = "S3CHINA"
+      terraform_resource_provider = "snowsql"
+    }
+  }
 
-locals {
   s3_bucket_name        = var.s3_bucket_name == "" ? "${replace(var.prefix, "_", "-")}-${var.env}-bucket" : "${replace(var.s3_bucket_name, "_", "-")}" # Only hiphens + lower alphanumeric are allowed for bucket name
   s3_reader_role_name   = "${var.prefix}-s3-reader"
   s3_sns_policy_name    = "${var.prefix}-s3-sns-topic-policy"

versions.tf

@@ -4,16 +4,26 @@ terraform {
   required_providers {
     aws = {
       source  = "hashicorp/aws"
-      version = ">= 4.38.0"
+      version = ">= 5.72.0"
     }
 
     snowflake = {
       source  = "Snowflake-Labs/snowflake"
-      version = ">= 0.64.0"
+      version = ">= 0.73.0"
 
       configuration_aliases = [
         snowflake.storage_integration_role,
       ]
     }
+
+    snowsql = {
+      source  = "aidanmelen/snowsql"
+      version = ">= 1.3.3"
+
+      configuration_aliases = [
+        snowsql.storage_integration_role,
+      ]
+    }
+
   }
 }