diff --git a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py index 339bbf4851b9..d934b7ec2a6b 100644 --- a/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py +++ b/metadata-ingestion/src/datahub/sql_parsing/sqlglot_utils.py @@ -342,6 +342,20 @@ def _strip_expression( if isinstance(node, (sqlglot.exp.In, sqlglot.exp.Values)): _simplify_node_expressions(node) elif isinstance(node, sqlglot.exp.Literal): + # A Redshift `COPY ... CREDENTIALS ''` parses the credentials + # as a Literal directly under a Credentials node. sqlglot's + # credentials_sql generator dispatches on isinstance(child, Literal) + # to choose the Redshift scalar form over the Snowflake key=value + # list; a Placeholder there routes into the list path that iterates + # the node and raises "TypeError: 'Placeholder' object is not + # iterable". Redact to a constant literal instead: serialization + # stays on the scalar path, the secret never lands in the generalized + # query text, and fingerprints stay independent of the credentials. + if ( + isinstance(node.parent, sqlglot.exp.Credentials) + and node.arg_key == "credentials" + ): + return sqlglot.exp.Literal.string("**REDACTED**") return sqlglot.exp.Placeholder() return node @@ -367,7 +381,11 @@ def get_query_fingerprint_debug( ) else: expression_sql = generalize_query_fast(expression, dialect=platform) - except (ValueError, sqlglot.errors.SqlglotError) as e: + except (ValueError, TypeError, sqlglot.errors.SqlglotError) as e: + # TypeError guards against sqlglot generator bugs that fail to serialize + # an otherwise-parseable statement (e.g. exotic COPY / credentials + # forms). Fingerprinting is best-effort with a raw-text fallback, so it + # must never propagate and abort ingestion. if not isinstance(expression, str): raise diff --git a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_utils.py b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_utils.py index 73aef98dcd9a..2317862be948 100644 --- a/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_utils.py +++ b/metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_utils.py @@ -5,6 +5,8 @@ import pytest import sqlglot +from datahub.sql_parsing import sqlglot_utils +from datahub.sql_parsing.fingerprint_utils import generate_hash from datahub.sql_parsing.query_types import get_query_type_of_sql from datahub.sql_parsing.schema_resolver import SchemaResolver from datahub.sql_parsing.sql_parsing_common import QueryType @@ -215,6 +217,52 @@ def test_redshift_query_fingerprint(): ) +def test_redshift_copy_credentials_generalization(): + # Regression: `COPY ... CREDENTIALS ''` used to crash generalization + # with `TypeError: 'Placeholder' object is not iterable`. The credentials + # literal was replaced with a Placeholder, which flips sqlglot's + # credentials_sql dispatch onto the Snowflake key=value path that iterates + # the node. + copy_sql = ( + "COPY my_schema.my_table FROM 's3://my-bucket/data' " + "CREDENTIALS 'aws_access_key_id=EXAMPLE;aws_secret_access_key=EXAMPLESECRET' " + "FORMAT AS PARQUET" + ) + + generalized = generalize_query(copy_sql, dialect="redshift") + + # The secret must never leak into the generalized (stored) query text. + assert "EXAMPLESECRET" not in generalized + assert "CREDENTIALS" in generalized + + # Fingerprinting must succeed and be independent of the actual secret, so + # two COPYs differing only in credentials share a fingerprint. + other_creds_sql = ( + "COPY my_schema.my_table FROM 's3://my-bucket/data' " + "CREDENTIALS 'aws_access_key_id=OTHER;aws_secret_access_key=OTHERSECRET' " + "FORMAT AS PARQUET" + ) + assert get_query_fingerprint(copy_sql, "redshift") == get_query_fingerprint( + other_creds_sql, "redshift" + ) + + +def test_get_query_fingerprint_survives_generalization_error(monkeypatch): + # Defense in depth: if generalization raises an unexpected error (e.g. a + # sqlglot generator bug on an exotic statement), fingerprinting must fall + # back to the raw text rather than propagating and killing the pipeline. + def _boom(*args: object, **kwargs: object) -> str: + raise TypeError("simulated sqlglot generator failure") + + monkeypatch.setattr(sqlglot_utils, "generalize_query", _boom) + + raw_query = "SELECT * FROM my_table" + fingerprint = get_query_fingerprint(raw_query, "redshift") + # Pin the fallback path: the fingerprint must be the hash of the raw text, + # proving generalization was bypassed rather than merely "didn't raise". + assert fingerprint == generate_hash(raw_query) + + def test_query_fingerprint_with_secondary_id(): query = "SELECT * FROM users WHERE id = 123"