Post Syndicated from Jason Ramos original https://aws.amazon.com/blogs/big-data/accelerate-sql-development-with-sagemaker-data-agent-in-query-editor/

When you develop SQL against Amazon Redshift and Amazon Athena, you spend time finding the right tables across hundreds of databases, writing complex joins and aggregations, debugging failed queries without context from previous attempts, and re-specifying filters for every new question. Amazon SageMaker Data Agent in Query Editor takes a different approach. You describe what you need in natural language, and the Data Agent generates the SQL. It references your actual tables through AWS Glue Data Catalog, proposes step-by-step plans for complex questions, retains context across your session, and offers one-click error recovery with Fix with AI. In this post, you learn how to use Data Agent in Query Editor to explore data, build multi-step analyses, recover from errors, and summarize results using a public education dataset.

Solution overview

You can go from a natural language question to executable SQL in seconds. Data Agent in Query Editor provides a conversational interface with direct access to your AWS data environment, so you spend less time on query mechanics and more time on analysis. Data Agent in Query Editor focuses specifically on SQL development against Amazon Redshift and Amazon Athena. (For Python, SQL, and PySpark across broader analytical and machine learning (ML) workloads, use Data Agent in notebooks.)

Data Agent provides four key capabilities:

• Catalog-aware SQL generation. You don’t need to browse catalog structures or memorize schema details. Data Agent reads your table metadata directly.
• Querybook and session context. You build on previous work. Data Agent uses context from your earlier queries and results.
• Step-by-step planning. You review and approve a structured plan before Data Agent generates SQL.
• Fix with AI. You recover from failed queries with one click.

Data Agent integrates with AWS Glue Data Catalog and reads your actual table names, column types, descriptions, and relationships, so generated SQL references your real tables. Each follow-up question builds on your current Query Editor session—the SQL cells in your querybook, the active connection, your selected cell, and execution results from previously run cells. For complex requests, Data Agent produces a structured plan that specifies which data to retrieve, how to aggregate it, and what filters to apply. You review and approve each step before Data Agent proceeds. When a query fails, choose Fix with AI to get a corrected query based on the error and the failed cell’s context.

[Figure 1: The Query Editor Fix with AI panel, showing a corrected SQL query ready for your review.]

Walkthrough: Education data analysis

In this section, you use Data Agent in Query Editor to analyze California schools data and identify where SAT improvement investment has the most impact. The walkthrough covers four tasks:

• Explore available data.
• Build a multi-step analysis plan.
• Summarize insights from your queries.
• Recover from a failed query.

The same workflow applies to your own data, whether you are analyzing sales figures, operational metrics, or financial records.

The California schools dataset contains SAT score results, school demographic information, and county-level data for public schools across California. The dataset includes tables that organize SAT scores by subject (reading, writing, math), school details (name, address, county, district), and enrollment figures. After you upload the data into your project database, you directly access the tables from Query Editor through your Amazon Athena or Amazon Redshift Lakehouse connection.

Prerequisites

To complete this walkthrough, you need intermediate SQL knowledge and basic familiarity with the AWS Management Console. You don’t need prior AWS Glue experience, but familiarity with data catalogs (centralized metadata repositories) helps.

You can choose one of two setup paths:

• Quick start (5 minutes). SageMaker Unified Studio provides a sample database (sagemaker_sample_db) with pre-loaded data. To explore it, choose Data in the navigation pane, expand AwsDataCatalog, and select sagemaker_sample_db.
• Full setup (30–45 minutes). Upload the California schools dataset into your project’s Lakehouse database. This dataset is publicly available from the California Department of Education. Download the SAT scores, school information, and county-level data files, then upload them through the SageMaker Unified Studio UI. In your project, go to Build, choose Query editor, right-click your project database in the Data explorer, and choose Create table. Drag and drop each CSV file to create the tables. SageMaker Unified Studio stores the data in the project-managed Amazon Simple Storage Service (Amazon S3) location, registers it in AWS Glue Data Catalog, and applies AWS Lake Formation governance automatically.

Running queries against Amazon Athena or Amazon Redshift might incur costs. For pricing details, refer to Amazon Athena pricing and Amazon Redshift pricing. For detailed setup instructions, refer to AWS Identity and Access Management (IAM)-based domains and projects. Before starting the walkthrough, you must have a SageMaker Unified Studio IAM-based domain with a project using the SQL analytics or All Capabilities project profile. The project automatically provisions an AWS Glue database, the required IAM role, and Athena or Redshift Lakehouse connections.

[Figure 2: The Data Explorer panel in Query Editor, showing the california_schools_db and sagemaker_sample_db tables.]

Explore available data. To start, enter the following prompt in the Data Agent panel:

Query my SAT scores from my california_schools_db

Data Agent searches AWS Glue Data Catalog, locates the relevant tables, and generates an initial exploratory query that retrieves SAT score records. It adds a SQL cell directly to your querybook.

• Review the generated SQL in the comparison view, which highlights the proposed code.
• Choose Accept, Reject, or Accept and run.
• After you run the cell, the results appear inline, giving you a view of the data (column names, score ranges, and the number of records) before you write SQL.

[Figure 3: Data Agent returns an exploratory query for the california_schools_db tables, ready for your review.]

[Figure 4: The SQL query results appear beneath the cell after you choose Accept and run.]

Build a multi-step analysis plan. With the data explored, enter a more complex analytical question:

Identify which subjects need investment to improve SAT scores in the lowest-performing counties. Include school-level details with addresses.

Data Agent proposes a step-by-step plan before generating SQL. For this request, Data Agent breaks the question into three steps:

1. Aggregate SAT scores by county and subject to find performance patterns.
2. Filter to counties with a sufficient number of schools and rank the lowest performers.
3. Join school address data to produce a final detailed list.

Review the plan in the Data Agent panel and choose Run step-by-step to proceed.

[Figure 5: Data Agent proposes a multi-step plan with options to Cancel plan or Run step-by-step.]

Data Agent generates SQL for each step and adds it as a separate querybook cell. Review each cell’s SQL in the comparison view, then choose Accept and run to execute it. The results from each step are visible inline, so you can verify the intermediate output (county-level aggregations, the filtered ranking, and the final school list) before moving to the next step. When the steps are complete, your querybook contains the full analytical progression from raw scores to a detailed investment list.

[Figure 6: Each plan step produces a querybook cell that you can review and run independently.]

Summarize insights from your queries. After running the analysis, enter the following prompt:

Summarize the insights from my queries

Data Agent has context on your querybook, including the SQL and the query results from each cell. It generates a natural language summary: which counties are underperforming, which subjects (reading, writing, or math) need the most attention in each county, and how many schools appear on the investment list. This summary provides a starting point for a report or presentation.

[Figure 7: Data Agent summarizes insights from the accumulated query results in the querybook.]

Recover from a failed query. During the analysis, a generated query might produce an error, for example, referencing a column name that doesn’t match the schema or a join condition that returns unexpected results. When a cell fails, Query Editor displays the error message and a Fix with AI option.

Choose Fix with AI, and Data Agent reads the error in the context of the failed cell, then generates corrected SQL and updates the querybook cell. Run the corrected cell to verify the fix.

[Figure 8: After you choose Fix with AI, Data Agent is prompted to generate a corrected query for the failed cell.]

[Figure 9: Data Agent returns corrected SQL for you to review.]

Security and governance

Data Agent operates within your AWS environment and only accesses data that your IAM policies explicitly permit. Your existing IAM access controls and AWS Lake Formation permissions determine what data Data Agent can reach. To use Data Agent, your project role must have permissions to invoke specific Amazon DataZone APIs. For more information, refer to Actions, resources, and condition keys for Amazon DataZone.

Data Agent includes content filtering that prevents it from responding to off-topic requests, requests to reveal its system prompt, and requests for internal technical implementation details. Data Agent is restricted to AWS-related topics and English-language output.

Amazon SageMaker stores your natural language prompts and generated SQL in the AWS Region where you created your SageMaker Unified Studio domain. Data Agent doesn’t store your data, querybook context, or catalog metadata.

To opt out of data usage for service improvement, configure an AI services opt-out policy for Amazon DataZone in AWS Organizations. For more information, refer to Data storage in the SageMaker Data Agent, Service improvement, and AI services opt-out policies.

Clean up

The walkthrough creates querybook cells in your Query Editor session but doesn’t provision standalone infrastructure. To remove the generated SQL cells, delete them from your querybook or delete the querybook itself.

If you uploaded the California schools dataset specifically for this walkthrough, remove the following resources to avoid ongoing charges:

• SageMaker Unified Studio domain. If you created a domain solely for this walkthrough, delete it to stop incurring charges. Refer to the SageMaker Unified Studio administration guide for deletion steps.
• Uploaded tables. In the Data explorer, right-click each table you created and choose Delete table to remove the data from your project database and the underlying S3 storage.
• Amazon Athena query results. Amazon Athena stores query results in an S3 output location. Delete the query result files from that bucket, or delete the bucket if you created it solely for this walkthrough.
• Amazon CloudWatch logs. If Amazon Athena queries generated CloudWatch log groups, delete those log groups to avoid storage charges.

Conclusion

Data Agent in Query Editor brings conversational, catalog-aware SQL development to your Amazon Redshift and Amazon Athena workloads. In this post, you explored unfamiliar data, built a multi-step investment analysis, recovered from query errors, and summarized findings through natural language prompts.

Data Agent works within your existing IAM and AWS Lake Formation security controls, keeps your data within your AWS environment, and retains context across your analytical workflow so each question builds on the last.

Get started with these next steps:

1. Run your first prompt. Open Query Editor in your SageMaker Unified Studio domain and enter Show me the top 10 tables in my catalog with the most columns. For setup, refer to the SageMaker Unified Studio getting started guide.
2. Add descriptions to your AWS Glue Data Catalog. Table descriptions and column-level business metadata improve the quality of generated SQL. For best practices, refer to Populating the AWS Glue Data Catalog.
3. Try a multi-step analysis. Enter Which product categories had declining revenue quarter-over-quarter, and which regions drove the decline? and review Data Agent’s plan step by step.

For more information, refer to the Amazon SageMaker Data Agent documentation, the What’s New blog post, Amazon Redshift documentation, and Amazon Athena documentation. To learn how Data Agent works in notebooks, refer to Accelerate context-aware data analysis and ML workflows with Amazon SageMaker Data Agent.

About the authors

Post Syndicated from Anthony Lempelius, James Mesney original https://aws.amazon.com/blogs/big-data/build-a-trusted-foundation-for-data-and-ai-using-alation-and-amazon-sagemaker-unified-studio/

This post was co-written with Anthony Lempelius and James Mesney from Alation.

When a team wants to reuse a dataset, whether it is to build a new pipeline, launch a dashboard, run an analysis, or power an AI application, the first challenge is rarely the code. Data engineers need to understand lineage, transformations, and operational expectations. Data analysts and BI engineers need consistent definitions, metrics, and trusted sources. Data scientists and AI engineers need to know provenance, quality, access constraints, and how data or features were derived. In many organizations, that context is captured in different places by different teams, often across solutions like Alation and SageMaker Unified Studio, both of which can serve as a system of record for business context depending on who is doing the work and where they operate day to day. When those perspectives are not connected, people revalidate the same information, debate definitions, and duplicate documentation across tools. A unified metadata foundation brings these role specific views together so business context, technical metadata, and governance stay aligned across platforms, making data easier to trust, easier to find, and easier to use across analytics and AI.

The new Alation integration with Amazon SageMaker Unified Studio addresses these challenges by synchronizing catalog metadata between both systems. This synchronization creates a unified metadata experience where technical teams working in SageMaker Unified Studio and business teams working in Alation collaborate on top of the same metadata. You can verify how ML and analytics assets are created, understand dependencies, and maintain traceability across your data lifecycle regardless of which system your teams prefer to use.

In this post, we demonstrate who benefits from this integration, how it works, the specific metadata it synchronizes, and provide a complete deployment guide for your environment.

The value of unified metadata governance

Organizations managing large-scale analytics and ML workloads face critical challenges when metadata is fragmented across multiple systems. When metadata exists in silos, data scientists spend valuable time searching for the right datasets. Teams duplicate metadata management efforts, creating inconsistent definitions and conflicting metrics across the organization.

Regulatory requirements demand clear provenance. Without unified metadata governance, organizations struggle to demonstrate compliance, trace data origins, and maintain audit trails across their ML and analytics pipelines. Data discovery becomes a bottleneck when teams can’t quickly find, understand, and trust the data they need, delaying model development and reducing the overall business value of data investments.

Applying consistent governance policies across disparate systems is nearly impossible without a unified metadata layer. This creates security vulnerabilities, data quality issues, and compliance blind spots. A unified metadata governance approach alleviates these challenges by providing a single source of truth for metadata across ML and analytics systems, enabling faster data discovery, consistent governance, and confident compliance while reducing the operational burden on data and ML teams.

Solution overview

The Alation and SageMaker Unified Studio integration unifies the user experience, synchronizing metadata from cataloged assets between both systems.

This Phase 1 integration extracts metadata from Amazon SageMaker Catalog into Alation, giving you one place to discover assets.

The integration connects through AWS Identity and Access Management (IAM) authentication and synchronizes key metadata elements, including domains, projects, asset names, descriptions, owners, glossary terms, and custom metadata fields. Every metadata update includes provenance information: the originating service, the person who made the change, and the timestamp, creating comprehensive audit trails for compliance.

You can run metadata extractions on demand or schedule them to run automatically. The system performs an initial bulk extraction of your selected domains and projects, then keeps it up-to-date through incremental updates using either event-driven triggers or scheduled polling. Communication uses encrypted APIs with scoped IAM permissions following least-privilege principles.

This integration helps organizations in financial services, telecommunications, retail, manufacturing, and transportation that manage large numbers of analytics and ML workloads across many systems and teams. You can reduce metadata duplication, accelerate data discovery, and enable your data scientists, analysts, and engineers to find trusted data faster so they can focus on building insights rather than validating data quality.

The following diagram illustrates the solution architecture.

The following screenshot showcases the Alation catalog displaying the SageMaker Unified Studio project and its synchronized assets.

Metadata synchronization

This integration automatically synchronizes essential metadata between SageMaker Unified Studio and Alation, facilitating consistent information across both systems. The synchronization brings together the types of metadata you need for discovery, governance, and audit workflows, giving you clearer insight into how datasets, features, and models relate across your services.

The integration synchronizes catalog metadata, including domains, projects, asset names, descriptions, owners, glossary terms, and metadata forms. Additionally, the integration synchronizes provenance metadata, which includes information about the originating service, the actor who made the change, and the timestamp, to support traceability and audit workflows.

Integration mechanics

The integration connects SageMaker Unified Studio and Alation through a scoped IAM role that provides secure, encrypted communication. After you configure this connection within Alation, the system performs an initial extraction of your selected domains and projects, then keeps information current through incremental updates using either event-driven triggers or scheduled polling.

The integration synchronizes metadata forms from SageMaker Unified Studio into Alation through automated field mapping between both systems’ schemas. Metadata forms can capture various asset specific details like feature store references, training run identifiers, model versions, and evaluation metrics.

Every metadata update includes provenance information: the originating service, the person who made the change, and when it occurred. This supports audit and stewardship workflows. Access controls follow least-privilege principles through IAM while applying Alation’s role-based permissions, letting you limit synchronization by project, namespace, or tag as needed.

Security and compliance

Security and compliance are critical when synchronizing metadata across systems. This integration follows enterprise security practices to facilitate safe, controlled metadata synchronization. The connector uses least-privilege access, encrypted transport, and clear separation between metadata and data, so you can maintain governance without disrupting existing workflows.

You configure a scoped IAM role to define which accounts, projects, and namespaces the connector can access, making sure access follows your organization’s security policies. Metadata moves over TLS-protected APIs, and you control which domains and projects to include in Alation. By default, the integration synchronizes only metadata; your data files and artifacts remain in their original AWS locations unless you explicitly choose to export them.

Alation maintains a complete audit trail by recording extraction events, mapping changes, and stewardship activities. These security controls support compliant metadata governance while preserving your existing operational practices.

Prerequisites

Before setting up this integration, ensure you have the following:

• An Alation Cloud Service (ACS) instance
• Alation server admin access
• An AWS account
• A SageMaker Unified Studio domain and project with existing metadata

Configure authentication

Before configuring the Alation connector, you must set up the required AWS resources and permissions. The first step is to configure authentication. The Alation connector supports two authentication methods to access SageMaker Unified Studio. Choose the method that best fits your security requirements.

Option 1: IAM role (Recommended)

Create an IAM role that the Alation connector will assume to access SageMaker Unified Studio. For detailed instructions on creating IAM roles, see IAM role creation.

The following is an example IAM permission policy for SageMaker Catalog access:

{
   "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "AlationSageMakerAccess",
            "Effect": "Allow",
            "Action": [
                "datazone:ListDomains",
                "datazone:GetFormType",
                "datazone:Search",
                "datazone:ListProjects",
                "datazone:GetAsset"
            ],
            "Resource": "arn:aws:datazone:<region>:<account-id>:domain/*”
        }
    ]
}

The following is an example trust policy for the IAM role:

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "AlationSageMakerAccessAssumeRole",
            "Effect": "Allow",
            "Principal": {
                "AWS": "<alation_provided_role_arn>"
            },
            "Action": "sts:AssumeRole"
        }
    ]
}     

Option 2: IAM user with access keys

Create an IAM user with programmatic access and attach the necessary permissions. For detailed instructions on creating IAM users, see Create an IAM user in your AWS account.

Create an IAM user with programmatic access enabled, attach the following policy, and generate access keys for use in Alation configuration:

{
   "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "AlationSageMakerAccess",
            "Effect": "Allow",
            "Action": [
                "datazone:ListDomains",
                "datazone:GetFormType",
                "datazone:Search",
                "datazone:ListProjects",
                "datazone:GetAsset"
            ],
            "Resource": "arn:aws:datazone:<region>:<account-id>:domain/*"
        }
    ]
}

Add IAM role or user to SageMaker Unified Studio domain

Add the IAM role or user you created to the SageMaker Unified Studio domain. For detailed instructions on adding users to a domain, see User management in Amazon SageMaker Unified Studio. The following screenshot shows an example of adding IAM users on the SageMaker dashboard.

Add IAM role or user to SageMaker Unified Studio projects

The IAM role or user must be added as a member to all SageMaker Unified Studio projects that contain metadata you want to synchronize with Alation. Projects without this member will not be included in the synchronization process.

Add the IAM role or user as a project member with Contributor or Owner permissions for each project you want to include in the sync, as illustrated in the following screenshot. For detailed instructions on adding project members, see Add project members.

Install SageMaker enhanced connector

After completing the AWS setup, you can configure the Alation connector to establish the integration. The connector is distributed as a .zip package for upload and installation in the Alation application. To obtain the connector, contact the Forward Deployed Engineering team or your Alation Account Manager.

When you have the .zip package, follow the installation procedures to add the connector.

Create and configure Alation’s data source

Navigate to the Data Sources section in Alation, create a new data source, and select SageMaker Catalog as the source type. Configure the connection settings with the authentication method chosen in the AWS setup.

For IAM role authentication, use the following configuration:

• Connection Type: IAM Role
• Role ARN: ARN of the IAM role created in AWS setup
• External ID: External ID configured in the trust policy
• AWS Region: Region where your SageMaker Unified Studio domain is located

For IAM user authentication, use the following configuration:

• Connection Type: Access Keys
• Access Key ID: Access key from AWS setup
• Secret Access Key: Secret key from AWS setup
• AWS Region: Region where your SageMaker Unified Studio domain is located

Test the connection to verify authentication and network connectivity, as shown in the following screenshot.

Configure metadata extraction settings

Configure the extraction scope by selecting the SageMaker domains and projects to synchronize, as shown in the following screenshot. Only projects where the IAM role or user is a member will be available for synchronization.

Run initial extraction

Execute the first metadata synchronization to import existing metadata from SageMaker Unified Studio into Alation. Monitor the extraction progress through Alation’s status indicators and validate that SageMaker assets appear correctly in the catalog.

The following screenshot shows the job history page with job status Running.

The following screenshot shows the job history page with job status Succeeded.

The following screenshot shows the Alation catalog displaying the SageMaker Unified Studio project and its synchronized assets.

Operate and tune

Configure ongoing operations by setting extraction cadence, configuring reconciliation alerts, and monitoring logs regularly. Add data stewards to synchronized assets, and consider enabling AI-generated descriptions or working with Alation Professional Services for advanced governance design.

Enhanced capabilities

The next phase of the integration introduces three key capabilities: bi-directional metadata synchronization, lineage replication, and data quality metadata replication. The bi-directional capability gives you the flexibility to control where metadata updates originate, either in Alation or in SageMaker Unified Studio, so you can manage metadata changes in the service that best aligns with your organizational workflows and governance processes.

The feature set is rolling out in phases. Phase 1 is available at the time of writing this post and provides extraction from SageMaker Unified Studio into Alation, including initial and incremental updates and audit logging. Phase 2 is coming soon and will offer configurable principal catalogs, advanced scoped syncs, and reconciliation workflows for Alation Cloud Service customers.

These enhancements will support governed, scalable ML operations with increasing depth and automation.

Conclusion

The Alation and SageMaker Unified Studio integration helps organizations bridge the gap between fast analytics and ML development and the governance requirements most enterprises face. By cataloging metadata from SageMaker Unified Studio in Alation, you gain a governed, discoverable view of how assets are created and used. This supports leaders, stewards, compliance teams, and ML practitioners who depend on accurate, well-documented data to scale analytics and AI responsibly.

To learn more about this integration and explore additional resources, refer to the Amazon SageMaker Unified Studio User Guide and Alation Documentation.

About the authors

Post Syndicated from Nira Jaiswal original https://aws.amazon.com/blogs/big-data/accelerate-data-governance-with-custom-subscription-workflows-in-amazon-sagemaker/

Amazon SageMaker provides a single data and AI development environment to discover and build with your data. This unified platform integrates functionality from existing AWS Analytics and Artificial Intelligence and Machine Learning (AI/ML) services, including Amazon EMR, AWS Glue, Amazon Athena, Amazon Redshift, and Amazon Bedrock.

Organizations need to efficiently manage data assets while maintaining governance controls in their data marketplaces. Although manual approval workflows remain important for sensitive datasets and production systems, there’s an increasing need for automated approval processes with less sensitive datasets. In this post, we show you how to automate subscription request approvals within SageMaker, accelerating data access for data consumers.

Prerequisites

For this walkthrough, you must have the following prerequisites:

• An AWS account – If you don’t have an account, you can create one. The account should have permission to do the following:
  • Create and manage SageMaker domains
  • Create and manage IAM roles
  • Create and invoke Lambda functions
• SageMaker domain – For instructions to create a domain, refer to Create an Amazon SageMaker Unified Studio domain – quick setup.
• A demo project – Create a demo project in your SageMaker domain. For instructions, see Create a project. For this example, we choose All capabilities in the project profile section.
• SageMaker domain ID, project ID, and project role ARN – These will be used in later steps to provide permissions for existing datasets and resources, and automatic subscription approval code. To retrieve this information, go to the Project details tab on the project details page on the SageMaker console.
  
• AWS CLI installed – You must have the AWS Command Line Interface (AWS CLI) version 2.11 or later.
• Python installed – You must have Python version 3.8 or later.
• IAM permissions – Sign in as the user with administrative access
• Lambda permissions – Configure the appropriate IAM permissions for the Lambda execution role. The following code is a sample role used for testing this solution. Before implementing this IAM policy in your environment, provide the values for your specific AWS Region and account ID. Adjust them based on the principle of least privilege. To learn more about creating Lambda execution roles, refer to Defining Lambda function permissions with an execution role.

  {
      "Version": "2012-10-17",
      "Statement": [
          {
              "Effect": "Allow",
              "Action": [
                  "datazone:ListSubscriptionRequests",
                  "datazone:AcceptSubscriptionRequest",
                  "datazone:GetSubscriptionRequestDetails",
                  "datazone:GetDomain",
                  "datazone:ListProjects"
              ],
              "Resource": "<<Domain-ARN>>"
          },
          {
              "Effect": "Allow",
              "Action": "sts:AssumeRole",
              "Resource": "<<Domain-ARN>>",
              "Condition": {
                  "StringEquals": {
                      "aws:PrincipalArn": "<<Lambda ARN>>"
                  }
              }
          },
          {
              "Effect": "Allow",
              "Action": "sns:Publish",
              "Resource": "<<SNS-ARN>>"
          },
          {
              "Effect": "Allow",
              "Action": [
                  "logs:CreateLogGroup",
                  "logs:CreateLogStream",
                  "logs:PutLogEvents"
              ],
              "Resource": [
                  "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/*",
                  "arn:aws:logs:${AWS::Region}:${AWS::AccountId}:log-group:/aws/lambda/*:*"
              ]
          }
      ]
  }

Solution overview

Understanding the subscription and approval workflow in Amazon SageMaker is important before diving deep into custom workflow solution. After an asset is published to the SageMaker catalog, data consumers can discover assets. When a data consumer discovers assets in SageMaker catalog, they request access to the asset, by submitting a subscription request with business justification and intended use case. The request enters a pending state and notifies the data producer or asset owner for review. The data producer evaluates the request based on governance policies, consumer credentials, and business context. The data producer can accept, reject, or request additional information from the data consumer. Upon acceptance, SageMaker triggers the AcceptSubscriptionRequest event and begins automated access provisioning. After a subscription is accepted, a subscription fulfilment process gets kicked off to facilitate access to the asset, for the data producer. SageMaker integrates deeply with AWS Lake Formation to manage fine-grained permissions. When a subscription is approved, SageMaker automatically calls Lake Formation APIs to grant specific database, table, and column-level permissions to the subscriber’s IAM role. Lake Formation acts as the central permission engine, translating subscription approvals into actual data access rights without manual intervention. The system provisions and updates resource-based policies on data sources. Once the provisioning completes, the data consumer can immediately access subscribed data through query engines like Athena, Redshift, or EMR, with Lake Formation enforcing permissions at query time.

By default, subscription requests to a published asset require manual approval by a data owner. However, Amazon SageMaker supports automatic approval of subscription requests at asset level: when publishing a data asset, you can choose to not require subscription approval. In this case, all incoming subscription requests to that asset are automatically approved. Let’s first outline the step-by-step process for disabling automatic approval at the asset level.

Configure automatic approval at asset level:

To configure automatic approval, data producers can follow the steps below.

1. Log in to SageMaker Unified Studio portal as data producer. Navigate to Assets and select the target asset
2. Choose Assets → Pick the asset, which you would like to configure for automatic approval.
3. On the asset details page, locate Edit Subscription settings in the right pane.
   
4. Choose Edit next to Subscription Required
   1. Select Not Required in the dialogue box
   2. Confirm your selection

   

Customize SageMaker’s subscription workflow:

While manual approval workflow remains essential for production environments and sensitive data handling, organizations seek to streamline and automate approvals for lower-risk environments and non-sensitive datasets. To achieve this project-level automation, we can enhance SageMaker’s native approval workflow through a custom event-driven solution. This solution leverages AWS’s serverless architecture, combining using AWS Lambda, Amazon EventBridge rules, and Amazon Simple Notification Service (Amazon SNS) to create an automated approval workflow. This customization allows organizations to maintain governance while reducing administrative overhead and accelerating the development cycle in non-critical environments. The event-driven approach ensures real-time processing of approval requests, maintains audit trails, and can be configured to apply different approval rules based on project characteristics and data sensitivity levels.

The custom workflow consists of the following steps:

1. The data consumer submits a subscription request for a published data asset.
2. SageMaker detects the request and generates a subscription event, which is automatically sent to EventBridge.
3. EventBridge triggers the designated Lambda function.
4. The Lambda function sends an AcceptSubscriptionRequest API call to SageMaker.
5. The function also sends a notification through Amazon SNS.
6. AWS Lake Formation processes the approved subscription and updates the relevant access control lists (ACLs) and permission sets.
7. Lake Formation grants access permissions to the data consumer’s project AWS Identity and Access Management (IAM) role.
8. The data consumer now has authorized access to the requested data asset and can begin working with the subscribed data.

The following diagram illustrates the high-level architecture of the solution.

Key benefits

This solution uses AWS Lambda and Amazon EventBridge to automate SageMaker subscription requests approvals, delivering the following benefits for organizations and end-users:

• Scalability – Automatically handles high volumes of subscription requests
• Cost-efficiency – Pay-as-you-go approach with no idle resource costs
• Minimal maintenance – Serverless components require no infrastructure management
• Flexible triggering – Supports event-driven, scheduled, and manual invocation modes
• Audit compliance – Comprehensive logging and traceability through AWS CloudTrail

Step-by-step procedure

This section outlines the detailed process for implementing a custom subscription request approval workflow in Amazon SageMaker

Create Lambda function

Complete the following steps to create your Lambda function:

1. On the Lambda console, choose Functions in the navigation pane.
2. Choose Create function.
3. Select Author from scratch.
4. For Function name, enter a name for the function.
5. For Runtime, choose your runtime (for this post, we use Python version 3.9 or later).
6. Choose Create function.
   
7. On the Lambda function page, choose the Configuration tab and then choose Permissions.
8. Note the execution role to use when configuring the SageMaker project.
   

Create SNS topic

For this solution, we create SNS topic. Complete the following steps to create the SNS topic for automatic approvals:

1. On the Amazon SNS console, choose Topics in the navigation pane.
2. Choose Create topic.
3. For Type, select Standard.
4. For Name, enter a name for the topic.
5. Choose Create topic.
   
6. On the SNS topic details page, note the SNS topic Amazon Resource Name (ARN) to use later in the Lambda function.
   
7. On Subscription tab, choose Create Subscription.
8. For Protocol, choose Email.
9. For Endpoint, enter email address of Data consumers.
   

Create EventBridge rule

Complete the following steps to create an EventBridge rule to capture subscription request events:

1. On the EventBridge console, choose Rules in the navigation pane.
2. Choose Create rule.
3. For Name, enter a name for the rule.
4. For Rule type, select Rule with event pattern.
   This option enables the automatic subscription approval workflow to be triggered when a subscription request is initiated. Alternatively, you can select Schedule to schedule the rule to trigger on a regular basis. Refer to Creating a rule that runs on a schedule in Amazon EventBridge to learn more.
5. Choose Next.
   
6. For Event source, select AWS events or EventBridge partner events.
7. For Creation method, select Use pattern form
8. For Event source, select AWS services
9. For AWS service, select DataZone.
10. For Event type, select Subscription Request Created.
    
11. Configure your target to route events to both the Lambda function and SNS topic.
12. Choose Next.
    
13. For this post, skip configuring tags and choose Next.
    
14. Review the settings and choose Create rule.

Configure automation workflow

Complete the following steps to configure the automation workflow:

1. On the Lambda console, go to the function you created.
2. Configure the EventBridge rule to trigger the Lambda function
3. Configure the destination as SNS topic for event notification.
   

Configure code in Lambda function

Complete the following steps to configure your Lambda function:

1. On the Lambda console, go to the function you created.
2. Add the following code to your function. Provide the domain ID, project ID, and SNS topic ARN that you noted earlier.

   import boto3
   import json
   import logging
   import os
   from botocore.exceptions import ClientError
   
   # Configure logging
   logger = logging.getLogger()
   logger.setLevel(logging.INFO)
   
   def lambda_handler(event, context):
       """Lambda function to auto-approve subscription requests in Amazon SageMaker"""
       try:
           # Initialize clients
           datazone_client = boto3.client('datazone')
           sns_client = boto3.client('sns')
           
           # Get configuration from environment variables or use hardcoded values
           domain_id = os.environ.get('DOMAIN_ID', '<domain_id>')
           project_id = os.environ.get('PROJECT_ID', '<project_id>')
           sns_topic_arn = os.environ.get('SNS_TOPIC_ARN', '<sns_topic_arn>')
           
           # Get pending subscription requests
           pending_requests = get_pending_requests(datazone_client, domain_id, project_id)
           
           if not pending_requests:
               logger.info("No pending subscription requests found")
               return
           
           # Process requests
           for request in pending_requests:
               approve_request(datazone_client, sns_client, domain_id, request, sns_topic_arn)
               
       except Exception as e:
           logger.error(f"Error: {str(e)}")
   
   def get_pending_requests(client, domain_id, project_id):
       """Get all pending subscription requests"""
       requests = []
       next_token = None
       
       try:
           while True:
               params = {
                   'domainIdentifier': domain_id,
                   'status': 'PENDING',
                   'approverProjectId': project_id
               }
               
               if next_token:
                   params['nextToken'] = next_token
               
               response = client.list_subscription_requests(**params)
               
               if 'items' in response:
                   requests.extend(response['items'])
               
               next_token = response.get('nextToken')
               if not next_token:
                   break
                   
           logger.info(f"Found {len(requests)} pending requests")
           return requests
           
       except ClientError as e:
           logger.error(f"Error listing requests: {e}")
           return []
   
   def approve_request(datazone_client, sns_client, domain_id, request, sns_topic_arn):
       """Approve a subscription request and send notification"""
       request_id = request.get('id')
       if not request_id:
           return
           
       try:
           # Approve the request
           datazone_client.accept_subscription_request(
               domainIdentifier=domain_id,
               identifier=request_id,
               decisionComment="Subscription request is auto-approved by Lambda"
           )
           
           # Send notification
           asset_name = request.get('assetName', 'Unknown asset')
           
           message = f"Your subscription request has been auto-approved by Lambda. You can now access this asset."
           
           sns_client.publish(
               TopicArn=sns_topic_arn,
               Subject=f"Subscription Request is auto-approved by Lambda",
               Message=message
           )
           
           logger.info(f"Approved request {request_id} for {asset_name}")
           
       except Exception as e:
           logger.error(f"Error processing request {request_id}: {e}")

3. Choose Test to test the Lambda function code. To learn more about testing Lambda code, refer to Testing Lambda functions in the console.
4. Choose Deploy to deploy the code.

Configure Lambda and project execution roles in SageMaker

Complete the following steps:

1. In SageMaker Unified Studio, go to your publishing project.
2. Choose Members in the navigation pane.
3. Choose Add members.
4. Add the Lambda execution role and project execution roles as Contributor.
   

Test the solution

Complete the following steps to test the solution:

1. In SageMaker Unified Studio, navigate to the data catalog and choose Subscribe on the configured asset to initiate a subscription request.
   
2. Choose Subscription requests in the navigation pane to view the outgoing requests and choose the Approved tab to verify automatic approval.
   
3. Choose View subscription to confirm the approver appears as the Lambda execution role with “Auto-approved by Lambda” as the reason.
   
4. On the CloudTrail console, choose Event history to view the event you created and review the automated approval audit trail.
   

Clean up

To avoid incurring future charges, clean up the resources you created during this walkthrough. The following steps use the AWS Management Console, but you can also use the AWS CLI.

1. Delete the SageMaker domain. To use the AWS CLI, run the following commands:

   aws sagemaker delete-project --project-name <project-name>
   aws datazone delete-domain –identifier <domain_identifier>

2. Delete the SNS topics. To use the AWS CLI, run the following command:

   aws sns delete-topic --topic-arn <topic-arn>

3. Delete the Lambda function. To use the AWS CLI, run the following command:

   aws lambda delete-function --function-name <Lambda function name>

Conclusion

Combining an event-driven architecture with SageMaker creates an automated, cost-effective solution for data governance challenges. This serverless approach automatically handles data access requests while maintaining compliance, so organizations can scale efficiently as their data grows. The solution discussed in this post can help data teams access insights faster with minimal operational costs, making it an excellent choice for businesses that need quick, compliant data access while keeping their systems lean and efficient.

To learn more, visit the Amazon SageMaker Unified Studio page.

About the authors