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

metadata-ingestion/setup.py

@@ -98,7 +98,7 @@
 
 sqlglot_lib = {
     # Using an Acryl fork of sqlglot.
-    # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:hsheth?expand=1
+    # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:main?expand=1
     "acryl-sqlglot[rs]==23.17.1.dev10",
 }
 

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

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
+    )