Skip to content

vdk-impala: Refactor query classifier for data lineage #1239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
kostoww merged 4 commits into from
Oct 18, 2022
Merged

vdk-impala: Refactor query classifier for data lineage #1239

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bbf94a6
vdk-impala: Refactor query classifier for data lineage
Oct 13, 2022
8bd7c2d
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 13, 2022
ce44501
vdk-impala: Refactor query classifier for data lineage
Oct 13, 2022
e563113
Merge branch 'main' into topic/pkostov/impala-refactor-lineage-detecting
kostoww Oct 18, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions projects/vdk-plugins/vdk-impala/setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
"vdk-core",
"vdk-lineage-model",
"impyla",
"sqlparse",
"tabulate",
"pydantic",
"pyarrow",
Expand Down
76 changes: 45 additions & 31 deletions projects/vdk-plugins/vdk-impala/src/vdk/plugin/impala/impala_lineage_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from typing import Optional
from typing import Tuple

import sqlparse
from impala._thrift_gen.RuntimeProfile.ttypes import TRuntimeProfileFormat
from impala.hiveserver2 import HiveServer2Cursor
from vdk.api.lineage.model.logger.lineage_logger import ILineageLogger
Expand Down Expand Up @@ -68,7 +69,7 @@ def db_connection_execute_operation(

def _get_lineage_data(self, cursor: HiveServer2Cursor) -> Optional[LineageData]:
query_statement = cursor._cursor.query_string
if not self._is_query_have_lineage(query_statement):
if not self._does_query_have_lineage(query_statement):
return None # do not capture lineage for queries that don't have lineage information
start_time = time.time_ns()
query_profile = cursor.get_profile(profile_format=TRuntimeProfileFormat.STRING)
Expand All @@ -95,7 +96,7 @@ def _get_lineage_data(self, cursor: HiveServer2Cursor) -> Optional[LineageData]:
)

@staticmethod
def _is_query_have_lineage(query_statement: str) -> bool:
def _does_query_have_lineage(query_statement: str) -> bool:
"""
This method checks the query type because not every query
has information that could be classified as lineage data
Expand All @@ -108,35 +109,48 @@ def _is_query_have_lineage(query_statement: str) -> bool:
if query_statement is None:
return False

statement_lines = query_statement.split(
os.linesep
) # TODO if further optimization is needed, consider sqlparse
for line in statement_lines:
line = line.strip().lower()
if line.startswith("--") or (line.startswith("/*") and line.endswith("*/")):
continue # Comments are omitted
if line.startswith("select 1 -- testing if connection is alive."):
return False # managed_connection has a way to open and check connections with keep alive query
if line.startswith(
(
"alter",
"compute",
"create",
"describe",
"drop",
"explain",
"grant",
"invalidate",
"refresh",
"revoke",
"set",
"show",
"truncate",
"use",
)
):
# these commands are not providing lineage data in the profile at the moment
return False
query_statement = query_statement.lower()
# managed_connection has a way to open and check connections with keep alive query
if "select 1 -- testing if connection is alive." in query_statement:
return False

query_statement = sqlparse.format(
Comment thread
kostoww marked this conversation as resolved.
sql=query_statement,
strip_comments=True,
strip_whitespace=True,
keyword_case="lower",
)

if query_statement.startswith("create"):
# some create statements might have lineage
# (create table .. as select ..) this way we check
# if select is present in some form. This might result
# in some corner cases of false positive, but the
# profile of impala will not return any lineage
# info (scan/write hdfs) as our goal of this method is
# to reduce non-lineage query with a non-complex way
return (
"select " in query_statement or "select" + os.linesep in query_statement
)
if query_statement.startswith(
(
"alter",
"compute",
"describe",
"drop",
"explain",
"grant",
"invalidate",
"refresh",
"revoke",
"set",
"show",
"truncate",
"use",
)
):
# these commands are not providing lineage data in the profile at the moment
return False

return True

Expand Down
172 changes: 160 additions & 12 deletions projects/vdk-plugins/vdk-impala/tests/impala_lineage_plugin_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,12 @@ def test_get_lineage_table_from_table_name_valid_name(self):
def test_get_lineage_table_from_table_name_none(self):
self.assertIsNone(ImpalaLineagePlugin._get_lineage_table_from_table_name(None))

def test_is_query_have_lineage(self):
def test_does_query_have_lineage(self):
self.assertTrue(
ImpalaLineagePlugin._is_query_have_lineage("SELECT * FROM table")
ImpalaLineagePlugin._does_query_have_lineage("SELECT * FROM table")
)
self.assertTrue(
ImpalaLineagePlugin._is_query_have_lineage(
ImpalaLineagePlugin._does_query_have_lineage(
"WITH temporaryTable(avgVal) as"
"(SELECT avg(Salary)"
"from Employee)"
Expand All @@ -30,14 +30,14 @@ def test_is_query_have_lineage(self):
)
)
self.assertTrue(
ImpalaLineagePlugin._is_query_have_lineage(
ImpalaLineagePlugin._does_query_have_lineage(
"-- job_name: a-job\n-- op_id: an-op\nINSERT "
"INTO TABLE schema.table /* +SHUFFLE */\n "
"SELECT t1.* FROM schema.table"
)
)
self.assertTrue(
ImpalaLineagePlugin._is_query_have_lineage(
ImpalaLineagePlugin._does_query_have_lineage(
"-- job_name: a-job\n"
"-- /* +SHUFFLE */ below is a query hint to "
"Impala. Do not remove!\n "
Expand All @@ -46,30 +46,178 @@ def test_is_query_have_lineage(self):
)
)

self.assertFalse(ImpalaLineagePlugin._is_query_have_lineage("USE database;"))
self.assertFalse(ImpalaLineagePlugin._does_query_have_lineage("USE database;"))
self.assertFalse(
ImpalaLineagePlugin._is_query_have_lineage("DROP TABLE table;")
ImpalaLineagePlugin._does_query_have_lineage("DROP TABLE table;")
)
self.assertFalse(
ImpalaLineagePlugin._is_query_have_lineage(
ImpalaLineagePlugin._does_query_have_lineage(
"select 1 -- testing if connection is alive."
)
)
self.assertFalse(
ImpalaLineagePlugin._is_query_have_lineage(
ImpalaLineagePlugin._does_query_have_lineage(
"alter table d2.mobile rename to d3.mobile;"
)
)
self.assertFalse(
ImpalaLineagePlugin._is_query_have_lineage(
ImpalaLineagePlugin._does_query_have_lineage(
"-- job_name: a-job\n"
"-- op_id: an-op\n"
"select 1 -- Testing if connection is alive."
)
)
self.assertFalse(
ImpalaLineagePlugin._is_query_have_lineage(
"-- job_name: a-job\n" "-- op_id: an-op\n" "DESCRIBE schema.table"
ImpalaLineagePlugin._does_query_have_lineage(
"-- job_name: a-job\n" "-- op_id: an-op\n" " DESCRIBE schema.table"
)
)
self.assertFalse(
ImpalaLineagePlugin._does_query_have_lineage(
"-- job_name: a-job\n" "-- op_id: an-op\n" " REFRESH schema.table"
)
)
self.assertTrue(
ImpalaLineagePlugin._does_query_have_lineage(
"CREATE TABLE database_one.table_for_prod "
"STORED AS PARQUET AS SELECT * FROM database_two.table_for_prod;"
)
)
self.assertTrue(
ImpalaLineagePlugin._does_query_have_lineage(
"CREATE TABLE database_one.table_for_prod "
"STORED AS PARQUET AS SELECT\n"
"* FROM database_two.table_for_prod;"
)
)
self.assertTrue(
ImpalaLineagePlugin._does_query_have_lineage(
"CREATE TABLE database_one.table_for_prod "
"STORED AS PARQUET AS"
"WITH temporaryTable(avgVal) as"
"(SELECT avg(Salary)"
"from Employee)"
"SELECT EmployeeID,Name, Salary"
"FROM Employee, temporaryTable"
"WHERE Employee.Salary > temporaryTable.avgVal;"
)
)
self.assertFalse(
ImpalaLineagePlugin._does_query_have_lineage(
"CREATE TABLE database_one.table_for_prod " "STORED AS PARQUET"
)
)
self.assertFalse(
ImpalaLineagePlugin._does_query_have_lineage(
"\n"
"-- job_name: job-name\n"
"-- op_id: job-name-1665673200-v9nth\n"
"-- template: template-complex-name\n"
"COMPUTE STATS shop1.users;\n"
)
)

self.assertTrue(
ImpalaLineagePlugin._does_query_have_lineage(
"\n"
"/***\n"
"show create table database.incident;\n"
"\n"
"-- drop table if exists database.incident;\n"
"\n"
"-- invalidate metadata database.incident;\n"
"\n"
"CREATE TABLE database.incident\n"
"(\n"
" id STRING,\n"
" incident STRING,\n"
")\n"
"STORED AS PARQUET\n"
"-- LOCATION check parquet files\n"
";\n"
"***/\n"
"\n"
"with\n"
"---------------------\n"
"indent_tabs_enums\n"
"---------------------\n"
"as\n"
"(\n"
" with\n"
" -------------------------\n"
" indent_tabs\n"
" -------------------------\n"
" as\n"
" (\n"
" select distinct trim( regexp_extract( reason, '\\d+', 0 )) issuenum\n"
" from database.tasks\n"
" where resource_type in ('bundle' )\n"
" and start_time >= '2020-06-01'\n"
" and lower( trim( task_status )) in ( 'finished', 'succeeded' )\n"
" and user_name = 'auto' \n"
" )\n"
")\n"
"select issue.id\n"
"where issue.id > '2019-02-01'\n"
" "
)
)
self.assertTrue(
ImpalaLineagePlugin._does_query_have_lineage(
"\n"
"-- job_name: just-a-job\n"
"-- op_id: just-a-job-1665673200-crhqw\n"
"-- template: vdk.templates.load.dimension\n"
"/* TO DO DROP AND RECREATE TARGET TABLE ON FULL RELOAD OR DATA TYPE CHANGE */\n"
"\n"
"-- DROP TABLE database_staging.a_table;\n"
"-- CREATE TABLE database_staging.a_table STORED AS PARQUET \n"
"-- AS SELECT * FROM database_production.a_table;\n"
"\n"
"-- /* +SHUFFLE */ below is a query hint to Impala. Do not remove!\n"
"-- See https://www.cloudera.com/documentation/enterprise/5-9-x/topics/impala_hints.html for details.\n"
"INSERT OVERWRITE TABLE database_staging.a_table /* +SHUFFLE */\n"
"SELECT * FROM database_production.a_table;\n"
)
)
self.assertTrue(
ImpalaLineagePlugin._does_query_have_lineage(
"\n"
"/***\n"
"show create table database_one.table_incidents;\n"
"\n"
"-- drop table if exists database_one.table_incidents;\n"
"\n"
"CREATE TABLE database_one.table_incidents\n"
"(\n"
" id STRING,\n"
" component STRING,\n"
" category_group STRING,\n"
" category STRING\n"
")\n"
"STORED AS PARQUET\n"
";\n"
"\n"
"***/\n"
"\n"
"---------------------------------------------------\n"
"insert overwrite table database_one.table_incidents\n"
"---------------------------------------------------\n"
"(\n"
" id,\n"
" component,\n"
" category_group,\n"
" category\n"
")\n"
"select isc.id\n"
", isc.component\n"
", 'cat_group' category_group\n"
", cfo.customvalue category\n"
"from database_one.table_components isc\n"
"join database_two.value cfv on isc.id = cfv.issue\n"
"join database_two.option cfo on cfv.f1 = cfo.f2\n"
" and cfv.stringvalue = cfo.id \n"
"where isc.component_group = 'Components' -- (6 of 12)\n"
)
)

Expand Down