From 9692d97e675e0996030bda8cf9d5e3d20a561822 Mon Sep 17 00:00:00 2001 From: macbre Date: Fri, 30 Apr 2021 11:55:03 +0200 Subject: [PATCH 1/6] Add test_create_table - see #35 --- test/test_create_table.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) create mode 100644 test/test_create_table.py diff --git a/test/test_create_table.py b/test/test_create_table.py new file mode 100644 index 00000000..ca4a080f --- /dev/null +++ b/test/test_create_table.py @@ -0,0 +1,15 @@ +from sql_metadata import Parser + + +def test_create_table(): + parser = Parser(""" +CREATE TABLE `new_table` ( + `item_id` int(9) NOT NULL AUTO_INCREMENT, + `foo` varchar(16) NOT NULL DEFAULT '', + PRIMARY KEY (`item_id`,`foo`), + KEY `idx_foo` (`foo`) +) CHARACTER SET utf8; + """) + + assert parser.tables == ["new_table"] + assert parser.columns == ["item_id", "foo"] From f57fccb33bb02273646851ef00dcb7f613a788ed Mon Sep 17 00:00:00 2001 From: macbre Date: Fri, 30 Apr 2021 11:55:58 +0200 Subject: [PATCH 2/6] Black reformat --- test/test_create_table.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/test_create_table.py b/test/test_create_table.py index ca4a080f..29590ac9 100644 --- a/test/test_create_table.py +++ b/test/test_create_table.py @@ -2,14 +2,16 @@ def test_create_table(): - parser = Parser(""" + parser = Parser( + """ CREATE TABLE `new_table` ( `item_id` int(9) NOT NULL AUTO_INCREMENT, `foo` varchar(16) NOT NULL DEFAULT '', PRIMARY KEY (`item_id`,`foo`), KEY `idx_foo` (`foo`) ) CHARACTER SET utf8; - """) + """ + ) assert parser.tables == ["new_table"] assert parser.columns == ["item_id", "foo"] From c1f2783659b9c6fc15bfcb3fdb7565edfcbc7258 Mon Sep 17 00:00:00 2001 From: macbre Date: Wed, 12 May 2021 16:12:53 +0200 Subject: [PATCH 3/6] Parser: add _is_create_table_query property --- sql_metadata/parser.py | 18 ++++++++++++++++++ test/test_create_table.py | 15 ++++++++++++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py index 00938bb8..b2d19247 100644 --- a/sql_metadata/parser.py +++ b/sql_metadata/parser.py @@ -335,6 +335,11 @@ def tables(self) -> List[str]: and token.previous_token.normalized not in ["AS", "WITH"] and token.normalized not in ["AS", "SELECT"] ): + # handle CREATE TABLE queries (#35) + # skip keyword that are withing parenthesis-wrapped list of column + if self._is_create_table_query and token.is_in_parenthesis: + continue + if token.next_token.is_dot: pass # part of the qualified name elif token.is_in_parenthesis and ( @@ -728,3 +733,16 @@ def _preprocess_query(self) -> str: query = re.sub(r"`([^`]+)`\.`([^`]+)`", r"\1.\2", query) return query + + @property + def _is_create_table_query(self) -> bool: + """ + Return True if the query begins with "CREATE TABLE" statement + """ + if ( + self.tokens[0].normalized == "CREATE" + and self.tokens[1].normalized == "TABLE" + ): + return True + + return False diff --git a/test/test_create_table.py b/test/test_create_table.py index 29590ac9..a9ef5138 100644 --- a/test/test_create_table.py +++ b/test/test_create_table.py @@ -1,6 +1,19 @@ from sql_metadata import Parser +def test_is_create_table_query(): + assert Parser("BEGIN")._is_create_table_query is False + assert Parser("SELECT * FROM `foo` ()")._is_create_table_query is False + + assert Parser("CREATE TABLE `foo` ()")._is_create_table_query is True + assert ( + Parser( + "create table abc.foo as SELECT pqr.foo1 , ab.foo2 FROM foo pqr, bar ab" + )._is_create_table_query + is True + ) + + def test_create_table(): parser = Parser( """ @@ -14,4 +27,4 @@ def test_create_table(): ) assert parser.tables == ["new_table"] - assert parser.columns == ["item_id", "foo"] + # assert parser.columns == ["item_id", "foo"] From afcd63a2129c9cc781bde0d1f3c9c5e370d76d45 Mon Sep 17 00:00:00 2001 From: macbre Date: Wed, 12 May 2021 16:14:57 +0200 Subject: [PATCH 4/6] Add test_create_table_as_select --- test/test_create_table.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/test/test_create_table.py b/test/test_create_table.py index a9ef5138..6788c79c 100644 --- a/test/test_create_table.py +++ b/test/test_create_table.py @@ -28,3 +28,16 @@ def test_create_table(): assert parser.tables == ["new_table"] # assert parser.columns == ["item_id", "foo"] + + +def test_create_table_as_select(): + parser = Parser( + """ +create table abc.foo + as SELECT pqr.foo1 , ab.foo2 + FROM foo pqr, bar ab; + """ + ) + + assert parser.tables == ["abc.foo", "foo", "bar"] + assert parser.columns == ["foo.foo1", "bar.foo2"] From fb230c00b9ddaa247bdad85ac03580ad55df0b72 Mon Sep 17 00:00:00 2001 From: macbre Date: Wed, 12 May 2021 16:29:48 +0200 Subject: [PATCH 5/6] Proper handling of columns list in CREATE TABLE queries --- sql_metadata/parser.py | 18 +++++++++++++++++- test/test_create_table.py | 2 +- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py index b2d19247..d23c5cf2 100644 --- a/sql_metadata/parser.py +++ b/sql_metadata/parser.py @@ -18,7 +18,7 @@ TABLE_ADJUSTMENT_KEYWORDS, WITH_ENDING_KEYWORDS, ) -from sql_metadata.token import SQLToken +from sql_metadata.token import SQLToken, EmptyToken from sql_metadata.utils import UniqueList @@ -123,6 +123,22 @@ def columns(self) -> List[str]: subqueries_names = self.subqueries_names for token in self.tokens: + # handle CREATE TABLE queries (#35) + if token.is_name and self._is_create_table_query: + # previous token is either ( or , -> indicates the column name + if token.is_in_parenthesis and token.previous_token.is_punctuation: + columns.append(str(token)) + continue + + # we're in CREATE TABLE query with the columns + # ignore any annotations outside the parenthesis with the list of columns + if ( + not token.is_in_parenthesis + and token.find_nearest_token("SELECT", value_attribute="normalized") + is EmptyToken + ): + continue + if token.is_name and not token.next_token.is_dot: # analyze the name tokens, column names and where condition values if ( diff --git a/test/test_create_table.py b/test/test_create_table.py index 6788c79c..ac0a544b 100644 --- a/test/test_create_table.py +++ b/test/test_create_table.py @@ -27,7 +27,7 @@ def test_create_table(): ) assert parser.tables == ["new_table"] - # assert parser.columns == ["item_id", "foo"] + assert parser.columns == ["item_id", "foo"] def test_create_table_as_select(): From b76c6320a5d3f4a10d991c6dc0855141b0d0ce3b Mon Sep 17 00:00:00 2001 From: macbre Date: Wed, 12 May 2021 16:32:38 +0200 Subject: [PATCH 6/6] Parser: improve docs --- sql_metadata/parser.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py index d23c5cf2..887c5053 100644 --- a/sql_metadata/parser.py +++ b/sql_metadata/parser.py @@ -132,6 +132,7 @@ def columns(self) -> List[str]: # we're in CREATE TABLE query with the columns # ignore any annotations outside the parenthesis with the list of columns + # e.g. ) CHARACTER SET utf8; if ( not token.is_in_parenthesis and token.find_nearest_token("SELECT", value_attribute="normalized")