- Introduction
- Create S3 bucket and cross-account instance profile
- Use Autoloader to create SNS-SQS across accounts
- Manually create SNS-SQS for cross-account Autoloader
- Test cross-account Autoloader connection
Introduction
This is a step-by-step guide to set up an AWS cross-account Databricks Autoloader connection in the File Notification mode. This will allow you to automatically load data from an S3 bucket in one AWS account (Account A) into a Databricks workspace in another AWS account (Account B).
Databricks Autoloader can either automatically set up SNS and SQS, or we can manually create the resources and then use them in the Autoloader. In either case, we will need an instance profile in Account B to access the SNS and SQS in Account A.
Understanding Key Services:
Before proceeding, let's clarify the purpose of the relevant services in AWS.
SNS (Simple Notification Service)
SNS is a fully managed pub/sub messaging service that enables seamless message delivery from publishers to subscribers. Autoloader file notification uses an SNS to get notifications whenever a file lands in S3.
SQS (Simple Queue Service)
SQS is a fully managed message queuing service that offers scalable, reliable, and distributed message queues. Autoloader uses an SQS to durably store messages from SNS. When the autoloader stream is started, it processes messages from the SQS to identify the new files.
Role-Based Access to Buckets
To access S3 buckets in another AWS account, you need to define IAM roles with policies that grant the necessary permissions. These roles establish trust relationships and ensure secure access to resources across accounts.
Trust Relationship
A trust relationship in AWS IAM defines which entities are trusted to assume a particular IAM role. When setting up cross-account access, trust relationships determine which accounts or entities can assume roles in other accounts.
Bucket Policy
A bucket policy in AWS S3 sets permissions for objects within a bucket, controlling access at the bucket and object level. It's written in JSON format and specifies who can access the bucket and what actions they can perform.
Create S3 bucket and cross-account instance profile
Now that we understand the relevant AWS services, we can get started setting up a cross-account Autoloader connection.
- Create an S3 bucket in Account A. An example name could be acc-a-autol-input.
- Create an instance profile in Account B (refer steps 1 to 4 under Step 1: Create an instance profile using the AWS console).
- Add the policy provided below to the Account B instance profile role to access the bucket in Account A. This policy gives the instance profile created in Step 2 access to the S3 bucket created in Step 1.
{ "Version": "2012-10-17", "Statement": [ { "Effect": "Allow", "Action": [ "s3:ListBucket" ], "Resource": [ "arn:aws:s3:::acc-a-autol-input" ] }, { "Effect": "Allow", "Action": [ "s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:PutObjectAcl" ], "Resource": [ "arn:aws:s3:::acc-a-autol-input/*" ] } ] } - Add a pass role for this instance profile role. This must be added to the role that created the Databricks deployment (refer steps 1 to 8 under Step 5: Add the S3 IAM role to the EC2 policy).
- Go to Account A and add a bucket policy to the bucket in Account A. This S3 bucket policy will grant the instance profile created in Step 2 access to the S3 bucket created in Step 1 in a different account.
{ "Version": "2012-10-17", "Statement": [ { "Sid": "Example permissions", "Effect": "Allow", "Principal": { "AWS": "arn:aws:iam::<account_b_id>:role/acc_b_instance_profile" }, "Action": [ "s3:GetBucketLocation", "s3:ListBucket" ], "Resource": "arn:aws:s3:::acc-a-autol-input" }, { "Effect": "Allow", "Principal": { "AWS": "arn:aws:iam::<account_b_id>:role/acc_b_instance_profile" }, "Action": [ "s3:PutObject", "s3:GetObject", "s3:DeleteObject", "s3:PutObjectAcl" ], "Resource": "arn:aws:s3:::acc-a-autol-input/*" } ] } - Execute steps 1 to 5 under Step 6: Add the instance profile to Databricks in the Databricks Workspace console.
Use Autoloader to create SNS-SQS across accounts
For Autoloader to automatically create SNS-SQS, you will need an IAM role with access to create SNS and SQS:
- Create an IAM role in Account A. This will have permissions to auto-create SNS-SQS for Autoloader.
- Role name example: acc_a_autol_auto_create_role
- Policy name example: acc_a_autol_auto_create_policy – This policy lets Databricks autoloader create an SNS and an SQS with prefix databricks-auto-ingest-* during the stream initialization of the autoloader job.
{ "Version": "2012-10-17", "Statement": [ { "Sid": "DatabricksAutoLoaderSetup", "Effect": "Allow", "Action": [ "s3:GetBucketNotification", "s3:PutBucketNotification", "sns:ListSubscriptionsByTopic", "sns:GetTopicAttributes", "sns:SetTopicAttributes", "sns:CreateTopic", "sns:TagResource", "sns:Publish", "sns:Subscribe", "sqs:CreateQueue", "sqs:DeleteMessage", "sqs:ReceiveMessage", "sqs:SendMessage", "sqs:GetQueueUrl", "sqs:GetQueueAttributes", "sqs:SetQueueAttributes", "sqs:TagQueue", "sqs:ChangeMessageVisibility" ], "Resource": [ "arn:aws:s3:::acc-a-autol-input", "arn:aws:sqs:<aws_region>:<account_a_id>:databricks-auto-ingest-*", "arn:aws:sns:<aws_region>:<account_a_id>:databricks-auto-ingest-*" ] }, { "Sid": "DatabricksAutoLoaderList", "Effect": "Allow", "Action": [ "sqs:ListQueues", "sqs:ListQueueTags", "sns:ListTopics" ], "Resource": "*" }, { "Sid": "DatabricksAutoLoaderTeardown", "Effect": "Allow", "Action": [ "sns:Unsubscribe", "sns:DeleteTopic", "sqs:DeleteQueue" ], "Resource": [ "arn:aws:sqs:<aws_region>:<account_a_id>:databricks-auto-ingest-*", "arn:aws:sns:<aws_region>:<account_a_id>:databricks-auto-ingest-*" ] } ] } - Trust Relationship – lets the instance profile assume this role.
{ "Version": "2012-10-17", "Statement": [ { "Sid": "Statement", "Effect": "Allow", "Principal": { "AWS": "arn:aws:iam::<account_b_id>:role/acc_b_instance_profile" }, "Action": "sts:AssumeRole" } ] }
- It needs to be added to the instance profile In Account B. Add this to the existing policy in instance profile IAM role acc_b_instance_profile:
{ "Sid": "AssumeRoleAccA", "Effect": "Allow", "Action": "sts:AssumeRole", "Resource": "arn:aws:iam::<account_a_id>:role/acc_a_autol_auto_create_role" } - Run and test the Autoloader code from a Notebook in your Databricks Workspace:
options = { "cloudFiles.format": "parquet", "cloudFiles.schemaLocation": "s3://acc-a-autol-input/schema/auto/", "cloudFiles.useNotifications": True, "cloudFiles.roleArn": "arn:aws:iam::<account_a_id>:role/acc_a_autol_auto_create_role", } df = ( spark.readStream.format("cloudFiles") .options(**options) .load("s3://acc-a-autol-input/data/auto/") ) ( df.writeStream.option("checkpointLocation", "s3://acc-a-autol-input/checkpoint/auto1/1/") .start("s3://acc-a-autol-input/output/auto/") )
This Autoloader code will auto-create an SNS and an SQS with names similar to:
- databricks-auto-ingest-query-fb599288-a1ea-4443-8d31-ed66d86fd9aa-source-0
- databricks-auto-ingest-query-fb599288-a1ea-4443-8d31-ed66d86fd9aa-source-0
These will get created once during the Stream Initialization of the Autoloader.
Manually create SNS-SQS for cross-account Autoloader
If you want to manually create SNS-SQS and link this to the Autoloader, follow these steps:
- Create an SNS in Account A called acc_a_autol_sns.
- Type: Standard
- Access policy of the SNS to read events from S3 bucket in Acc A:
{ "Version": "2008-10-17", "Id": "notificationPolicy", "Statement": [ { "Sid": "allowS3Notification", "Effect": "Allow", "Principal": { "AWS": "*" }, "Action": "SNS:Publish", "Resource": "arn:aws:sns:<aws_region>:<account_a_id>:acc_a_autol_sns", "Condition": { "ArnLike": { "aws:SourceArn": "arn:aws:s3:*:*:acc-a-autol-input" } } } ] }
- Create an SQS in Account A called acc_a_autol_sqs.
- Type: Standard
- Access policy of the SQS to get messages from the SNS:
{ "Version": "2008-10-17", "Id": "notificationPolicy", "Statement": [ { "Sid": "allowS3Notification", "Effect": "Allow", "Principal": { "AWS": "*" }, "Action": "SQS:SendMessage", "Resource": "arn:aws:sqs:<aws_region>:<account_a_id>:acc_a_autol_sqs", "Condition": { "ArnLike": { "aws:SourceArn": "arn:aws:sns:<aws_region>:<account_a_id>:acc_a_autol_sns" } } } ] } - Subscribe to the above SNS
- Create an S3 bucket event notification
- Name: databricks-autol-events
- Prefix: data/manual/
- All object create events
- Destination: Put the SNS ARN
- Create an IAM role with SNS-SQS Autoloader auto-create policy in Account A
- Role Name: acc_a_autol_manual_create_role
- Policy Name: acc_a_autol_manual_create_policy – This policy lets Databricks autoloader access to the SNS and SQS created in Steps 1 and 2.
{ "Version": "2012-10-17", "Statement": [ { "Sid": "DatabricksAutoLoaderUse", "Effect": "Allow", "Action": [ "s3:GetBucketNotification", "sns:ListSubscriptionsByTopic", "sns:GetTopicAttributes", "sns:TagResource", "sns:Publish", "sqs:DeleteMessage", "sqs:ReceiveMessage", "sqs:SendMessage", "sqs:GetQueueUrl", "sqs:GetQueueAttributes", "sqs:TagQueue", "sqs:ChangeMessageVisibility" ], "Resource": [ "arn:aws:sqs:<aws_region>:<account_a_id>:acc_a_autol_sqs", "arn:aws:sns:<aws_region>:<account_a_id>:acc_a_autol_sns", "arn:aws:s3:::acc-a-autol-input" ] } ] } - Trust Relationship – lets the instance profile assume this role.
{ "Version": "2012-10-17", "Statement": [ { "Sid": "Statement", "Effect": "Allow", "Principal": { "AWS": "arn:aws:iam::<account_b_id>:role/acc_b_instance_profile" }, "Action": "sts:AssumeRole" } ] } -
Add this to the existing policy in instance profile IAM role acc_b_instance_profile:
{ "Sid": "AssumeManualRoleAccA", "Effect": "Allow", "Action": "sts:AssumeRole", "Resource": "arn:aws:iam::<account_a_id>:role/acc_a_autol_manual_create_role" }
Test cross-account Autoloader connection
Run and test the Autoloader code from a notebook in Databricks Workspace:
options = {
"cloudFiles.format": "parquet",
"cloudFiles.schemaLocation": "s3://acc-a-autol-input/schema/manual/",
"cloudFiles.useNotifications": True,
"cloudFiles.roleArn": "arn:aws:iam::<account_a_id>:role/acc_a_autol_manual_create_role",
"cloudFiles.queueUrl": "https://sqs.<aws_region>.amazonaws.com/<account_a_id>/acc_a_autol_sqs"
}
df = (
spark.readStream.format("cloudFiles")
.options(**options)
.load("s3://acc-a-autol-input/data/manual/")
)
(
df.writeStream.option("checkpointLocation", "s3://acc-a-autol-input/checkpoint/manual1/1/")
.start("s3://acc-a-autol-input/output/manual/")
)
By following these steps, you should be able to successfully set up a cross-account Autoloader connection with Databricks.
2 Comments
