From 3713f143b9fe0d71a58b11b8b37f2a9c7a902254 Mon Sep 17 00:00:00 2001 From: T-Santos Date: Mon, 26 Aug 2024 14:47:27 -0400 Subject: [PATCH] Feature/support for truncate table query parsing (#1) (#521) * Feature/support for truncate table query parsing (#1) * failing test for truncate table * fix failing truncate table test --------- Co-authored-by: Santos, Tyler (Boston) * Update test/test_truncate_table.py Co-authored-by: Maciej Brencz * Update test_truncate_table.py Add whitespace --------- Co-authored-by: Santos, Tyler (Boston) Co-authored-by: Maciej Brencz --- sql_metadata/keywords_lists.py | 2 ++ sql_metadata/parser.py | 2 +- test/test_truncate_table.py | 9 +++++++++ 3 files changed, 12 insertions(+), 1 deletion(-) create mode 100644 test/test_truncate_table.py diff --git a/sql_metadata/keywords_lists.py b/sql_metadata/keywords_lists.py index 8521c296..4abeaea0 100644 --- a/sql_metadata/keywords_lists.py +++ b/sql_metadata/keywords_lists.py @@ -81,6 +81,7 @@ class QueryType(str, Enum): CREATE = "CREATE TABLE" ALTER = "ALTER TABLE" DROP = "DROP TABLE" + TRUNCATE = "TRUNCATE TABLE" class TokenType(str, Enum): @@ -109,6 +110,7 @@ class TokenType(str, Enum): "ALTERTABLE": QueryType.ALTER, "DROPTABLE": QueryType.DROP, "CREATEFUNCTION": QueryType.CREATE, + "TRUNCATETABLE": QueryType.TRUNCATE, } # all the keywords we care for - rest is ignored in assigning diff --git a/sql_metadata/parser.py b/sql_metadata/parser.py index 64f40219..d003c547 100644 --- a/sql_metadata/parser.py +++ b/sql_metadata/parser.py @@ -116,7 +116,7 @@ def query_type(self) -> str: ) if tokens[index].normalized == "CREATE": switch = self._get_switch_by_create_query(tokens, index) - elif tokens[index].normalized in ("ALTER", "DROP"): + elif tokens[index].normalized in ("ALTER", "DROP", "TRUNCATE"): switch = tokens[index].normalized + tokens[index + 1].normalized else: switch = tokens[index].normalized diff --git a/test/test_truncate_table.py b/test/test_truncate_table.py new file mode 100644 index 00000000..3662dbbc --- /dev/null +++ b/test/test_truncate_table.py @@ -0,0 +1,9 @@ +from sql_metadata import Parser +from sql_metadata.keywords_lists import QueryType + + +def test_truncate_table(): + parser = Parser("TRUNCATE TABLE foo") + assert parser.query_type == QueryType.TRUNCATE + assert parser.tables == ["foo"] + assert parser.columns == []