Querying Data Lake

Query Data Lake

Execute SQL queries on Amazon Athena across default and federated catalogs (Glue, S3 Tables, Redshift) with workgroup selection, statement classification, and error recovery.

Overview

Executes and manages Athena SQL queries across default and federated catalogs. Selects a workgroup, resolves target assets (delegating fuzzy references to finding-data-lake-assets), classifies statements for safety, and reports cost and data scanned. Use the AWS MCP server for sandboxed execution and audit logging; the same AWS CLI commands work directly when the MCP server is not available.

Constraints for parameter acquisition:

• You MUST accept a single optional argument: SQL text, a named-query name, a workgroup name, a catalog name, or profile TABLE_NAME
• You MUST accept the argument as direct text or a pointer to a file containing SQL
• You MUST ask the user for the target AWS region if not already set
• You MUST confirm the output S3 location before executing any non-trivial query
• You MUST respect the user's decision to abort at any step

Common Tasks

1. Verify Dependencies

Check for required tools and AWS access before running queries.

Constraints:

• You MUST verify AWS MCP server tools are available (aws___call_aws) and run queries through them when present; fall back to AWS CLI only if the MCP server is unavailable
• You MUST NOT fall back to shell or Bash for query execution — results must be captured via the MCP tool or aws athena CLI so output location and cost are tracked
• You MUST confirm credentials with aws sts get-caller-identity and inform the user about any missing tools

2. Resolve Workgroup

Check caller identity, list workgroups, auto-select the best one (see workgroup-selection.md).

Constraints:

• You MUST select a workgroup before submitting any query (prevents output-location errors)
• You MUST present the selected workgroup and its output location to the user
• You MUST NOT auto-escalate to a different workgroup on failure without user confirmation

3. Resolve the Target Asset

If the user refers to a table by name, by business concept ("our quarterly report", "the sales data"), by S3 path, or by catalog without specifying the table, delegate to finding-data-lake-assets to return the concrete database.table (and catalog if non-default).

Constraints:

• You MUST NOT attempt to resolve fuzzy asset references with athena list-data-catalogs or by iterating get-tables — those miss federated catalogs and waste tokens
• You SHOULD skip this step only when the user provides a fully-qualified reference (exact database.table) or raw SQL they want executed as-is
• You MUST state the resolved asset explicitly before building the query: "Found [table] in [catalog]. Using this for the query."
• You SHOULD default to the default Glue catalog unless the user mentions "federated", "Redshift", "S3 Tables", or finding-data-lake-assets returns a different catalog

4. Discover Schema

For analytical queries, You SHOULD profile the target table before building the final query. You MUST show sample rows (SELECT ... LIMIT 5) as part of profiling.

5. Build Query

Table addressing depends on catalog type:

• Default Glue catalog: database.table (omit the catalog prefix for single-catalog queries). In cross-catalog queries, qualify default-catalog tables with "awsdatacatalog".database.table.
• Registered data source: datasource.database.table
• Unregistered Glue catalog: "catalog/subcatalog".database.table

6. Classify and Execute

Classify the SQL statement before executing:

Example tool call (via AWS MCP server):

aws___call_aws(command="aws athena start-query-execution --work-group <WORKGROUP_NAME> --query-string '<sql>' --query-execution-context Database=<db>")

For federated or S3 Tables catalogs, also set Catalog=<CATALOG_PATH> in the execution context (e.g. Catalog=s3tablescatalog/<BUCKET_NAME>).

Constraints:

• You MUST warn the user before executing when the target is Redshift-federated ("No partition pruning — every query scans the full table")
• You MUST warn the user before executing a cross-catalog join ("Cross-catalog joins incur network overhead and may be slow")
• You MUST confirm the output S3 location before executing
• You MUST explain which tool is being called before executing
• You MUST respect the user's decision to abort

7. Present and Recover

Present results with cost, data scanned, duration, and actionable insights. On failure, list available workgroups and let the user choose which to retry with.

Argument Routing

Resolve in this order; stop at the first match:

1. Contains SQL keywords (SELECT, SHOW, DESCRIBE, INSERT, etc.) — SQL text, execute directly 2. profile TABLE_NAME — run comprehensive table profiling (see query-patterns.md) 3. Matches a known named query — look up and execute 4. Matches a known workgroup — show workgroup status and recent queries 5. Matches a known catalog — delegate to exploring-data-catalog to enumerate databases and tables 6. No args — show recent query activity and available tables

Principles

• Always select workgroup before executing (prevents output-location errors)
• Profile unfamiliar tables before running analytical queries
• Present cost alongside results so users build cost awareness
• Suggest LIMIT for exploratory queries on large tables
• Never ask domain questions with obvious answers, but always confirm security-relevant actions (workgroup switches, output location changes, non-SELECT statements)

Troubleshooting

Additional Resources

• Workgroup selection logic
• Common query patterns
• Athena best practices
• Athena federated query

Common Query Patterns (Presto/Athena SQL)

Table Profiling

-- Schema discovery
SELECT column_name, data_type
FROM information_schema.columns
WHERE table_schema = '<database>' AND table_name = '<table>';

-- Quick row count and date range
SELECT COUNT(*) as total_rows,
       MIN(created_at) as earliest,
       MAX(created_at) as latest
FROM <table>;

-- Sample data (always do this before analytical queries)
SELECT * FROM <table> LIMIT 5;

-- Null analysis
SELECT
    '<column>' as field,
    COUNT(*) - COUNT(<column>) as null_count,
    ROUND((COUNT(*) - COUNT(<column>)) * 100.0 / COUNT(*), 2) as null_pct
FROM <table>;

Cohort Retention

WITH cohorts AS (
    SELECT
        user_id,
        DATE_TRUNC('month', first_activity_date) as cohort_month
    FROM users
),
activity AS (
    SELECT
        user_id,
        DATE_TRUNC('month', activity_date) as activity_month
    FROM user_activity
)
SELECT
    c.cohort_month,
    COUNT(DISTINCT c.user_id) as cohort_size,
    COUNT(DISTINCT CASE
        WHEN a.activity_month = c.cohort_month THEN a.user_id
    END) as month_0,
    COUNT(DISTINCT CASE
        WHEN a.activity_month = DATE_ADD('month', 1, c.cohort_month) THEN a.user_id
    END) as month_1,
    COUNT(DISTINCT CASE
        WHEN a.activity_month = DATE_ADD('month', 3, c.cohort_month) THEN a.user_id
    END) as month_3,
    COUNT(DISTINCT CASE
        WHEN a.activity_month = DATE_ADD('month', 6, c.cohort_month) THEN a.user_id
    END) as month_6
FROM cohorts c
LEFT JOIN activity a ON c.user_id = a.user_id
GROUP BY c.cohort_month
ORDER BY c.cohort_month;

Funnel Analysis

WITH funnel AS (
    SELECT
        user_id,
        MAX(CASE WHEN event = 'page_view' THEN 1 ELSE 0 END) as step_1_view,
        MAX(CASE WHEN event = 'signup_start' THEN 1 ELSE 0 END) as step_2_start,
        MAX(CASE WHEN event = 'signup_complete' THEN 1 ELSE 0 END) as step_3_complete,
        MAX(CASE WHEN event = 'first_purchase' THEN 1 ELSE 0 END) as step_4_purchase
    FROM events
    WHERE event_date >= DATE_ADD('day', -30, CURRENT_DATE)
    GROUP BY user_id
)
SELECT
    COUNT(*) as total_users,
    SUM(step_1_view) as viewed,
    SUM(step_2_start) as started_signup,
    SUM(step_3_complete) as completed_signup,
    SUM(step_4_purchase) as purchased,
    ROUND(100.0 * SUM(step_2_start) / NULLIF(SUM(step_1_view), 0), 1) as view_to_start_pct,
    ROUND(100.0 * SUM(step_3_complete) / NULLIF(SUM(step_2_start), 0), 1) as start_to_complete_pct,
    ROUND(100.0 * SUM(step_4_purchase) / NULLIF(SUM(step_3_complete), 0), 1) as complete_to_purchase_pct
FROM funnel;

Deduplication

-- Keep the most recent record per key (Presto/Athena syntax)
WITH ranked AS (
    SELECT
        *,
        ROW_NUMBER() OVER (
            PARTITION BY entity_id
            ORDER BY updated_at DESC
        ) as rn
    FROM source_table
)
SELECT * FROM ranked WHERE rn = 1;

Window Functions

-- Running total
SUM(revenue) OVER (ORDER BY event_date ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) as running_total

-- 7-day moving average
AVG(revenue) OVER (ORDER BY event_date ROWS BETWEEN 6 PRECEDING AND CURRENT ROW) as moving_avg_7d

-- Period-over-period comparison
LAG(value, 1) OVER (PARTITION BY entity ORDER BY event_date) as prev_value

-- Percent of total
revenue / SUM(revenue) OVER () as pct_of_total
revenue / SUM(revenue) OVER (PARTITION BY category) as pct_of_category

-- Ranking
ROW_NUMBER() OVER (PARTITION BY category ORDER BY revenue DESC) as rank_in_category

Period Comparison / Growth

When the user asks for "growth", "change", or "comparison" between periods, compute the delta — not raw totals.

WITH quarterly AS (
    SELECT
        category,
        QUARTER(order_date) as q,
        SUM(amount) as revenue
    FROM orders
    WHERE YEAR(order_date) = 2025
    GROUP BY category, QUARTER(order_date)
)
SELECT
    curr.category,
    prev.revenue as prev_period,
    curr.revenue as curr_period,
    ROUND((curr.revenue - prev.revenue) / prev.revenue * 100, 1) as growth_pct
FROM quarterly curr
JOIN quarterly prev ON curr.category = prev.category AND curr.q = prev.q + 1
ORDER BY growth_pct DESC;

Performance-Aware Patterns

-- Always filter on partition keys to reduce scan cost
SELECT region, COUNT(*)
FROM sales
WHERE year = '2024' AND month = '02'
GROUP BY region;

-- Use LIMIT for exploratory queries
SELECT * FROM large_table LIMIT 100;

-- Use approximate functions for large-scale cardinality
SELECT APPROX_DISTINCT(user_id) as approx_unique_users
FROM events;

Data Quality Checks

-- Distinct value counts per column
SELECT
    COUNT(DISTINCT col1) as col1_unique,
    COUNT(DISTINCT col2) as col2_unique
FROM <table>;

-- Detect unexpected values
SELECT column_name, COUNT(*) as cnt
FROM <table>
GROUP BY column_name
ORDER BY cnt DESC
LIMIT 20;

-- Check for join explosion
SELECT COUNT(*) as pre_join FROM table_a;
SELECT COUNT(*) as post_join FROM table_a a JOIN table_b b ON a.id = b.a_id;

Workgroup Selection

Always list workgroups first before executing any query.

Detect Execution Context

Before selecting a workgroup, determine the current IAM identity:

aws sts get-caller-identity --query Arn --output text

The ARN pattern reveals the execution context:

SageMaker Project Role Selection

When running as a SageMaker project role (AmazonDataZone-* or SageMakerUnifiedStudio-*):

1. List all workgroups the role can access:

   aws athena list-work-groups --query 'WorkGroups[].Name' --output json

2. Extract the project ID from the role ARN. Split the role name on -. The first segment is the prefix (e.g., AmazonDataZone), the second segment is the project ID (e.g., abc123def), and subsequent segments form the suffix (e.g., DataLakeAccess). Take the second segment. The project ID is an alphanumeric string (no hyphens). Known suffixes that follow the project ID: DataLakeAccess, SparkAccess, QueryAccess, IngestionAccess. Example:

   arn:aws:sts::123456789012:assumed-role/AmazonDataZone-abc123def-DataLakeAccess/session
                                                         ^^^^^^^^^
                                                         project ID = abc123def

3. Match the workgroup to the project. Project workgroups follow the pattern sagemaker-studio-workgroup-<project-id> or contain the project ID. 4. If exactly one sagemaker-studio-workgroup-* exists, verify its suffix contains the project ID extracted in step 2. If it matches, use it. If it does not match, fall through to step 6. 5. If multiple exist, pick the one whose suffix matches the project ID extracted from the role ARN. Optionally check environment variables SAGEMAKER_PROJECT_ID or SAGEMAKER_PROJECT_NAME if the ARN extraction is ambiguous. 6. If no sagemaker-studio-workgroup-* exists, do not fall back to other workgroups. Inform the user that no project-scoped workgroup was found and ask them to verify their project configuration or IAM permissions.

Project roles typically have IAM permissions scoped to their own workgroup. Attempting to use primary or another project's workgroup will fail with AccessDeniedException. Do not retry with primary in this context.

General Priority Order (Non-Project Roles)

1. sagemaker-studio-workgroup-* workgroups -- most reliable, always have output locations configured 2. Workgroups with explicitly configured output locations 3. primary workgroup (use with caution, may lack output location)

Error Recovery

Anti-patterns

• Never default to primary workgroup without checking others first
• Never hardcode a workgroup name across sessions
• Never retry with primary when running as a SageMaker project role -- it will fail with AccessDeniedException