fix(sql-parser): prevent bad urns from alter table lineage

[hsheth2]

• Ensures column names aren't in the alter table out tables.
• Add support for drop table/view statements.
• Hook into the standard --update-golden-files mechanism.

Checklist

• The PR conforms to DataHub's Contributing Guideline (particularly Commit Message Format)
• Links to related issues (if applicable)
• Tests for the changes have been added/updated (if applicable)
• Docs related to the changes have been added/updated (if applicable). If a new feature has been added a Usage Guide has been added for the same.
• For any breaking change/potential downtime/deprecation/big changes an entry has been made in Updating DataHub

Summary by CodeRabbit

Summary by CodeRabbit

• New Features

  • Enhanced SQL lineage analysis to support ALTER TABLE operations, improving accuracy in tracking table structure changes.
  • Introduced new JSON structures for handling metadata related to ALTER TABLE, DROP SCHEMA, DROP TABLE, and DROP VIEW operations across various databases.

• Tests

  • Added new tests for validating the functionality of ALTER TABLE, DROP TABLE, DROP VIEW, and DROP SCHEMA operations, ensuring correct processing of SQL commands across different contexts.

• Chores

  • Updated JSON structures to standardize data handling for partition specifications and data process instances in the Airflow plugin.

[coderabbitai]

Walkthrough

The recent changes enhance SQL lineage extraction capabilities by adding support for ALTER TABLE operations, refining tracking of table modifications while excluding irrelevant Drop operations. New test cases for various SQL statements ensure robust validation of these enhancements, broadening the testing coverage across multiple databases.

Changes

Files	Change Summary
.../sqlglot_lineage.py	Modified _table_level_lineage to handle ALTER TABLE operations and filter out Drop operations, improving lineage tracking accuracy.
.../goldens/test_bigquery_alter_table_column.json, .../goldens/test_snowflake_drop_schema.json, .../goldens/test_sqlite_drop_table.json, .../goldens/test_sqlite_drop_view.json	Introduced new JSON structures for representing SQL commands (ALTER TABLE, DROP SCHEMA, DROP TABLE, DROP VIEW) with detailed metadata and processing info.
.../test_sqlglot_lineage.py	Added new test functions to validate handling of ALTER TABLE, DROP TABLE, DROP VIEW, and DROP SCHEMA statements across different databases.
.../check_sql_parser_result.py	Removed an outdated comment about future enhancements, streamlining the code without affecting functionality.
.../airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json	Reordered properties in JSON structures and shifted focus from input to output datasets in data processing instances, enhancing clarity and organization.

Poem

🐇 In fields of code where bunnies hop,
We’ve added lineage, we’ll never stop!
Tables altered, and drops ignored,
Our SQL magic, it’s now restored!
With tests in place, our coverage blooms,
Hopping through queries, dispelling the glooms! 🌼✨

---

Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media?

Share

• X
• Mastodon
• Reddit
• LinkedIn

Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>.
  • Generate unit testing code for this file.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai generate unit testing code for this file.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai generate interesting stats about this repository and render them as a table.
  • @coderabbitai show all the console.log statements in this repository.
  • @coderabbitai read src/utils.ts and generate unit testing code.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (invoked as PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Additionally, you can add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 0

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 903ec87 and 6c313ac.

Files selected for processing (3)

• metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py (2 hunks)
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_alter_table_column.json (1 hunks)
• metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py (1 hunks)

Additional comments not posted (8)

metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_alter_table_column.json (6)

2-2: Verify the query_type.

The query_type is set to UNKNOWN. Ensure this is intentional and correctly represents the type of SQL query being tested.

---

3-3: LGTM!

The query_type_props field is correctly set as an empty dictionary.

---

4-4: Verify the query_fingerprint.

Ensure that the query_fingerprint is correctly generated and matches the expected value for the query.

---

5-8: LGTM!

The in_tables and out_tables fields are correctly set.

---

9-9: LGTM!

The column_lineage field is correctly set to null.

---

10-13: LGTM!

The debug_info field is correctly populated with relevant information.

metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py (1)

Line range hint 207-214:
LGTM!

The test function test_bigquery_alter_table_column correctly asserts the SQL result for the ALTER TABLE operation, using the appropriate dialect and expected file.

metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py (1)

Line range hint 200-222:
LGTM!

The modifications to the _table_level_lineage function correctly handle ALTER TABLE operations and filter out Drop operations.

[coderabbitai]

Actionable comments posted: 0

Outside diff range, codebase verification and nitpick comments (4)

metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py (4)

1218-1225: Add documentation for the new BigQuery test.

The test_bigquery_alter_table_column function is a valuable addition for testing BigQuery SQL parsing. Consider adding a comment explaining the purpose of this test, especially since it involves altering a table.

+ # Test for ALTER TABLE statement in BigQuery, dropping a column.

---

1228-1235: Add documentation for the new SQLite drop table test.

The test_sqlite_drop_table function is a valuable addition for testing SQLite SQL parsing. Consider adding a comment explaining the purpose of this test.

+ # Test for DROP TABLE statement in SQLite.

---

1238-1245: Add documentation for the new SQLite drop view test.

The test_sqlite_drop_view function is a valuable addition for testing SQLite SQL parsing. Consider adding a comment explaining the purpose of this test.

+ # Test for DROP VIEW statement in SQLite.

---

1248-1255: Add documentation for the new Snowflake drop schema test.

The test_snowflake_drop_schema function is a valuable addition for testing Snowflake SQL parsing. Consider adding a comment explaining the purpose of this test.

+ # Test for DROP SCHEMA statement in Snowflake.

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 6c313ac and 0d4cd6a.

Files selected for processing (7)

• metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py (1 hunks)
• metadata-ingestion/src/datahub/testing/check_sql_parser_result.py (1 hunks)
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_alter_table_column.json (1 hunks)
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_drop_schema.json (1 hunks)
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_sqlite_drop_table.json (1 hunks)
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_sqlite_drop_view.json (1 hunks)
• metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py (2 hunks)

Files skipped from review due to trivial changes (2)

• metadata-ingestion/src/datahub/testing/check_sql_parser_result.py
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_sqlite_drop_view.json

Files skipped from review as they are similar to previous changes (2)

• metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_alter_table_column.json

Additional comments not posted (3)

metadata-ingestion/tests/unit/sql_parsing/goldens/test_snowflake_drop_schema.json (1)

1-11: LGTM!

The JSON structure for the test case is well-formed and aligns with the purpose of testing SQL parsing for a "DROP SCHEMA" statement in Snowflake.

metadata-ingestion/tests/unit/sql_parsing/goldens/test_sqlite_drop_table.json (1)

1-13: LGTM!

The JSON structure for the test case is well-formed and aligns with the purpose of testing SQL parsing for a "DROP TABLE" statement in SQLite. The out_tables field correctly captures the expected output.

metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py (1)

11-18: Ensure UPDATE_FILES is correctly set for golden file updates.

The fixture set_update_sql_parser uses monkeypatch to set UPDATE_FILES. Ensure that this is intended for updating golden files during test runs.

Verification successful

---

UPDATE_FILES is correctly set for golden file updates.

The UPDATE_FILES variable is used to determine whether golden files should be updated during test runs. It is set based on an environment variable and correctly manipulated in the test setup using monkeypatch. This setup ensures that golden files are updated when intended.

• metadata-ingestion/src/datahub/testing/check_sql_parser_result.py: Defines UPDATE_FILES based on the environment variable.
• metadata-ingestion/tests/unit/sql_parsing/test_sqlglot_lineage.py: Uses monkeypatch to set UPDATE_FILES for test purposes.

Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Description: Check if the `UPDATE_FILES` attribute is correctly set for golden file updates.

# Test: Search for occurrences of `UPDATE_FILES` in the codebase. Expect: Correct usage in test setup.
rg --type python 'UPDATE_FILES'

Length of output: 62

---

Script:

#!/bin/bash
# Description: Check if the `UPDATE_FILES` attribute is correctly set for golden file updates.

# Test: Search for occurrences of `UPDATE_FILES` in the codebase. Expect: Correct usage in test setup.
rg 'UPDATE_FILES'

Length of output: 378

[coderabbitai]

Actionable comments posted: 0

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 0d4cd6a and 73cb9f0.

Files selected for processing (1)

• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json (19 hunks)

Additional comments not posted (2)

metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json (2)

353-354: Reordering JSON properties for consistency.

The partitionSpec properties have been reordered to place "partition" before "type". This change improves consistency across the JSON structure.

Also applies to: 370-371, 555-556, 745-746, 762-763, 878-879, 1164-1165, 1181-1182, 1423-1424, 1665-1666, 1682-1683, 1833-1834, 1850-1851

---

1486-1487: Shift focus from input to output datasets.

Input datasets are replaced with empty arrays, and output datasets are explicitly listed. This aligns with the goal of ensuring valid URNs in data lineage.

Ensure that this change aligns with the intended data processing logic and does not impact existing workflows.

Also applies to: 1700-1701

[coderabbitai]

Actionable comments posted: 0

Outside diff range, codebase verification and nitpick comments (1)

metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator_no_dag_listener.json (1)

1730-1733: Update aspectName references to ensure consistency.

The change from dataProcessInstanceInput to dataProcessInstanceOutput is not consistently applied across the codebase. Please update the following files to reflect the intended focus on outputs:

• smoke-test/tests/cypress/data.json
• metadata-ingestion/tests/integration/dbt/dbt_test_prefer_sql_parser_lineage_golden.json
• metadata-ingestion/tests/integration/dbt/dbt_test_test_model_performance_golden.json
• metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_golden.json
• metadata-ingestion/tests/integration/fivetran/fivetran_snowflake_empty_connection_user_golden.json
• metadata-ingestion/examples/test_examples/via_node_test_example_fivetran.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag_no_dag_listener.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets_no_dag_listener.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_simple_dag.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_basic_iolets.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_basic_iolets.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_snowflake_operator.json
• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_simple_dag.json

Analysis chain

Update aspectName to reflect output focus.

Changing aspectName from dataProcessInstanceInput to dataProcessInstanceOutput correctly aligns with the shift to output-focused data processing. Ensure that all related logic and downstream systems are updated accordingly.

Also applies to: 1743-1746

Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Description: Verify that all references to `dataProcessInstanceInput` have been updated to `dataProcessInstanceOutput`.

# Test: Search for any remaining references to `dataProcessInstanceInput`. Expect: No matches.
rg --type json '"aspectName": "dataProcessInstanceInput"'

Length of output: 3645

Review details

Configuration used: CodeRabbit UI
Review profile: CHILL

Commits

Files that changed from the base of the PR and between 73cb9f0 and 9e821bf.

Files selected for processing (3)

• metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator_no_dag_listener.json (16 hunks)
• metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py (1 hunks)
• metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_alter_table_column.json (1 hunks)

Files skipped from review as they are similar to previous changes (1)

• metadata-ingestion/tests/unit/sql_parsing/goldens/test_bigquery_alter_table_column.json

Additional comments not posted (4)

metadata-ingestion/src/datahub/sql_parsing/sqlglot_lineage.py (3)

192-202: Include ALTER TABLE in lineage tracking.

The addition of sqlglot.exp.AlterTable in the set of expressions ensures that alterations to table structures are considered in lineage tracking. This aligns with the PR objective to enhance lineage accuracy.

---

220-227: Filter out Drop operations from lineage.

The logic to exclude tables that are dropped is correctly implemented by checking if the statement is an instance of sqlglot.exp.Drop and ensuring the table name is populated. This prevents irrelevant references from affecting lineage.

---

234-234: Exclude tables with Drop parent operation.

The condition if not isinstance(table.parent, sqlglot.exp.Drop) effectively filters out tables that are part of a Drop operation, ensuring they do not contribute to lineage tracking.

metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator_no_dag_listener.json (1)

339-340: Standardize partitionSpec key order.

The reordering of keys in partitionSpec sections enhances consistency across the JSON structure. This change does not affect functionality but improves readability.

Also applies to: 385-386, 438-439, 644-645, 725-726, 1003-1004, 1084-1085, 1290-1291, 1371-1372, 1616-1617, 1633-1634, 1715-1716