mirror of https://github.com/apache/superset.git
parent
6d8870125f
commit
2c564817f1
|
@ -28,7 +28,7 @@ import sqlparse
|
||||||
from sqlalchemy import and_
|
from sqlalchemy import and_
|
||||||
from sqlglot import exp, parse, parse_one
|
from sqlglot import exp, parse, parse_one
|
||||||
from sqlglot.dialects import Dialects
|
from sqlglot.dialects import Dialects
|
||||||
from sqlglot.errors import ParseError
|
from sqlglot.errors import SqlglotError
|
||||||
from sqlglot.optimizer.scope import Scope, ScopeType, traverse_scope
|
from sqlglot.optimizer.scope import Scope, ScopeType, traverse_scope
|
||||||
from sqlparse import keywords
|
from sqlparse import keywords
|
||||||
from sqlparse.lexer import Lexer
|
from sqlparse.lexer import Lexer
|
||||||
|
@ -287,7 +287,7 @@ class ParsedQuery:
|
||||||
"""
|
"""
|
||||||
try:
|
try:
|
||||||
statements = parse(self.stripped(), dialect=self._dialect)
|
statements = parse(self.stripped(), dialect=self._dialect)
|
||||||
except ParseError:
|
except SqlglotError:
|
||||||
logger.warning("Unable to parse SQL (%s): %s", self._dialect, self.sql)
|
logger.warning("Unable to parse SQL (%s): %s", self._dialect, self.sql)
|
||||||
return set()
|
return set()
|
||||||
|
|
||||||
|
@ -319,12 +319,17 @@ class ParsedQuery:
|
||||||
elif isinstance(statement, exp.Command):
|
elif isinstance(statement, exp.Command):
|
||||||
# Commands, like `SHOW COLUMNS FROM foo`, have to be converted into a
|
# Commands, like `SHOW COLUMNS FROM foo`, have to be converted into a
|
||||||
# `SELECT` statetement in order to extract tables.
|
# `SELECT` statetement in order to extract tables.
|
||||||
literal = statement.find(exp.Literal)
|
if not (literal := statement.find(exp.Literal)):
|
||||||
if not literal:
|
|
||||||
return set()
|
return set()
|
||||||
|
|
||||||
pseudo_query = parse_one(f"SELECT {literal.this}", dialect=self._dialect)
|
try:
|
||||||
sources = pseudo_query.find_all(exp.Table)
|
pseudo_query = parse_one(
|
||||||
|
f"SELECT {literal.this}",
|
||||||
|
dialect=self._dialect,
|
||||||
|
)
|
||||||
|
sources = pseudo_query.find_all(exp.Table)
|
||||||
|
except SqlglotError:
|
||||||
|
return set()
|
||||||
else:
|
else:
|
||||||
sources = [
|
sources = [
|
||||||
source
|
source
|
||||||
|
|
|
@ -271,6 +271,7 @@ def test_extract_tables_illdefined() -> None:
|
||||||
assert extract_tables("SELECT * FROM catalogname..tbname") == {
|
assert extract_tables("SELECT * FROM catalogname..tbname") == {
|
||||||
Table(table="tbname", schema=None, catalog="catalogname")
|
Table(table="tbname", schema=None, catalog="catalogname")
|
||||||
}
|
}
|
||||||
|
assert extract_tables('SELECT * FROM "tbname') == set()
|
||||||
|
|
||||||
|
|
||||||
def test_extract_tables_show_tables_from() -> None:
|
def test_extract_tables_show_tables_from() -> None:
|
||||||
|
@ -558,6 +559,10 @@ def test_extract_tables_multistatement() -> None:
|
||||||
Table("t1"),
|
Table("t1"),
|
||||||
Table("t2"),
|
Table("t2"),
|
||||||
}
|
}
|
||||||
|
assert extract_tables(
|
||||||
|
"ADD JAR file:///hive.jar; SELECT * FROM t1;",
|
||||||
|
engine="hive",
|
||||||
|
) == {Table("t1")}
|
||||||
|
|
||||||
|
|
||||||
def test_extract_tables_complex() -> None:
|
def test_extract_tables_complex() -> None:
|
||||||
|
@ -1815,10 +1820,7 @@ def test_extract_table_references(mocker: MockerFixture) -> None:
|
||||||
# test falling back to sqlparse
|
# test falling back to sqlparse
|
||||||
logger = mocker.patch("superset.sql_parse.logger")
|
logger = mocker.patch("superset.sql_parse.logger")
|
||||||
sql = "SELECT * FROM table UNION ALL SELECT * FROM other_table"
|
sql = "SELECT * FROM table UNION ALL SELECT * FROM other_table"
|
||||||
assert extract_table_references(
|
assert extract_table_references(sql, "trino") == {
|
||||||
sql,
|
|
||||||
"trino",
|
|
||||||
) == {
|
|
||||||
Table(table="table", schema=None, catalog=None),
|
Table(table="table", schema=None, catalog=None),
|
||||||
Table(table="other_table", schema=None, catalog=None),
|
Table(table="other_table", schema=None, catalog=None),
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue