test(ingest/sql): refactor CLL generator + add tests (#10580)

datahub-project · May 24, 2024 · 2e14f70 · 2e14f70
1 parent e79bdc0
commit 2e14f70
Show file tree

Hide file tree

Showing 16 changed files with 585 additions and 281 deletions.
diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/_extractors.py
@@ -262,7 +262,7 @@ def extract(self) -> Optional[TaskMetadata]:
         )
 
         operator: "BigQueryInsertJobOperator" = self.operator
-        sql = operator.configuration.get("query")
+        sql = operator.configuration.get("query", {}).get("query")
         if not sql:
             self.log.warning("No query found in BigQueryInsertJobOperator")
             return None

diff --git a/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py b/metadata-ingestion-modules/airflow-plugin/src/datahub_airflow_plugin/datahub_listener.py
@@ -247,10 +247,11 @@ def _extract_lineage(
                 SQL_PARSING_RESULT_KEY, None
             )
         if sql_parsing_result:
-            if sql_parsing_result.debug_info.error:
-                datajob.properties["datahub_sql_parser_error"] = str(
-                    sql_parsing_result.debug_info.error
-                )
+            if error := sql_parsing_result.debug_info.error:
+                logger.info(f"SQL parsing error: {error}", exc_info=error)
+                datajob.properties[
+                    "datahub_sql_parser_error"
+                ] = f"{type(error).__name__}: {error}"
             if not sql_parsing_result.debug_info.table_error:
                 input_urns.extend(sql_parsing_result.in_tables)
                 output_urns.extend(sql_parsing_result.out_tables)

diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json b/metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v2_sqlite_operator.json
diff --git a/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py b/metadata-ingestion-modules/airflow-plugin/tests/integration/test_plugin.py
@@ -373,6 +373,7 @@ def test_airflow_plugin(
         golden_path=golden_path,
         ignore_paths=[
             # TODO: If we switched to Git urls, maybe we could get this to work consistently.
+            r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['datahub_sql_parser_error'\]",
             r"root\[\d+\]\['aspect'\]\['json'\]\['customProperties'\]\['openlineage_.*'\]",
         ],
     )

diff --git a/metadata-ingestion-modules/airflow-plugin/tox.ini b/metadata-ingestion-modules/airflow-plugin/tox.ini
@@ -4,7 +4,7 @@
 # and then run "tox" from this directory.
 
 [tox]
-envlist = py38-airflow21, py38-airflow22, py310-airflow24, py310-airflow26, py310-airflow27, py310-airflow28
+envlist = py38-airflow21, py38-airflow22, py310-airflow24, py310-airflow26, py310-airflow27, py310-airflow28, py310-airflow29
 
 [testenv]
 use_develop = true
@@ -19,13 +19,15 @@ deps =
     airflow26: apache-airflow~=2.6.0
     airflow27: apache-airflow~=2.7.0
     airflow28: apache-airflow~=2.8.0
+    airflow29: apache-airflow~=2.9.0
 
     # Respect the Airflow constraints files.
     # We can't make ourselves work with the constraints of Airflow < 2.3.
     py310-airflow24: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.4.3/constraints-3.10.txt
     py310-airflow26: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.6.3/constraints-3.10.txt
     py310-airflow27: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.7.3/constraints-3.10.txt
     py310-airflow28: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.8.1/constraints-3.10.txt
+    py310-airflow29: -c https://raw.githubusercontent.com/apache/airflow/constraints-2.9.1/constraints-3.10.txt
 
     # Before pinning to the constraint files, we previously left the dependencies
     # more open. There were a number of packages for which this caused issues.
@@ -53,6 +55,6 @@ commands =
 [testenv:py310-airflow24]
 extras = dev,integration-tests,plugin-v2,test-airflow24
 
-[testenv:py310-airflow{26,27,28}]
+[testenv:py310-airflow{26,27,28,29}]
 extras = dev,integration-tests,plugin-v2
 
diff --git a/metadata-ingestion/setup.py b/metadata-ingestion/setup.py
@@ -98,8 +98,8 @@
 
 sqlglot_lib = {
     # Using an Acryl fork of sqlglot.
-    # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:hsheth?expand=1
-    "acryl-sqlglot[rs]==23.17.1.dev10",
+    # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:main?expand=1
+    "acryl-sqlglot[rs]==24.0.1.dev7",
 }
 
 classification_lib = {

diff --git a/metadata-ingestion/src/datahub/ingestion/api/report.py b/metadata-ingestion/src/datahub/ingestion/api/report.py
@@ -41,7 +41,7 @@ def to_pure_python_obj(some_val: Any) -> Any:
         if isinstance(some_val, SupportsAsObj):
             return some_val.as_obj()
         elif isinstance(some_val, pydantic.BaseModel):
-            return some_val.dict()
+            return Report.to_pure_python_obj(some_val.dict())
         elif dataclasses.is_dataclass(some_val):
             return dataclasses.asdict(some_val)
         elif isinstance(some_val, list):

diff --git a/metadata-ingestion/src/datahub/sql_parsing/_models.py b/metadata-ingestion/src/datahub/sql_parsing/_models.py
@@ -29,7 +29,13 @@ def __lt__(self, other: "_FrozenModel") -> bool:
         for field in self.__fields__:
             self_v = getattr(self, field)
             other_v = getattr(other, field)
-            if self_v != other_v:
+
+            # Handle None values by pushing them to the end of the ordering.
+            if self_v is None and other_v is not None:
+                return False
+            elif self_v is not None and other_v is None:
+                return True
+            elif self_v != other_v:
                 return self_v < other_v
 
         return False

diff --git a/metadata-ingestion/src/datahub/sql_parsing/query_types.py b/metadata-ingestion/src/datahub/sql_parsing/query_types.py
@@ -0,0 +1,75 @@
+from typing import Optional, Tuple
+
+import sqlglot
+
+from datahub.sql_parsing.sql_parsing_common import QueryType, QueryTypeProps
+from datahub.sql_parsing.sqlglot_utils import (
+    DialectOrStr,
+    get_dialect,
+    is_dialect_instance,
+)
+
+
+def _is_temp_table(table: sqlglot.exp.Table, dialect: sqlglot.Dialect) -> bool:
+    identifier: sqlglot.exp.Identifier = table.this
+
+    return identifier.args.get("temporary") or (
+        is_dialect_instance(dialect, "redshift") and identifier.name.startswith("#")
+    )
+
+
+def _get_create_type_from_kind(kind: Optional[str]) -> QueryType:
+    if kind and "TABLE" in kind:
+        return QueryType.CREATE_TABLE_AS_SELECT
+    elif kind and "VIEW" in kind:
+        return QueryType.CREATE_VIEW
+    else:
+        return QueryType.CREATE_OTHER
+
+
+def get_query_type_of_sql(
+    expression: sqlglot.exp.Expression, dialect: DialectOrStr
+) -> Tuple[QueryType, QueryTypeProps]:
+    dialect = get_dialect(dialect)
+    query_type_props: QueryTypeProps = {}
+
+    # For creates, we need to look at the inner expression.
+    if isinstance(expression, sqlglot.exp.Create):
+        if is_create_table_ddl(expression):
+            return QueryType.CREATE_DDL, query_type_props
+
+        kind = expression.args.get("kind")
+        if kind:
+            kind = kind.upper()
+            query_type_props["kind"] = kind
+
+        target = expression.this
+        if any(
+            isinstance(prop, sqlglot.exp.TemporaryProperty)
+            for prop in (expression.args.get("properties") or [])
+        ) or _is_temp_table(target, dialect=dialect):
+            query_type_props["temporary"] = True
+
+        query_type = _get_create_type_from_kind(kind)
+        return query_type, query_type_props
+
+    # UPGRADE: Once we use Python 3.10, replace this with a match expression.
+    mapping = {
+        sqlglot.exp.Select: QueryType.SELECT,
+        sqlglot.exp.Insert: QueryType.INSERT,
+        sqlglot.exp.Update: QueryType.UPDATE,
+        sqlglot.exp.Delete: QueryType.DELETE,
+        sqlglot.exp.Merge: QueryType.MERGE,
+        sqlglot.exp.Query: QueryType.SELECT,  # unions, etc. are also selects
+    }
+
+    for cls, query_type in mapping.items():
+        if isinstance(expression, cls):
+            return query_type, query_type_props
+    return QueryType.UNKNOWN, {}
+
+
+def is_create_table_ddl(statement: sqlglot.exp.Expression) -> bool:
+    return isinstance(statement, sqlglot.exp.Create) and isinstance(
+        statement.this, sqlglot.exp.Schema
+    )