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) <[email protected]> * Update test/test_truncate_table.py Co-authored-by: Maciej Brencz <[email protected]> * Update test_truncate_table.py Add whitespace --------- Co-authored-by: Santos, Tyler (Boston) <[email protected]> Co-authored-by: Maciej Brencz <[email protected]>
macbre · Aug 26, 2024 · 3713f14 · 3713f14
1 parent 9a78fbc
commit 3713f14
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 1 deletion.
diff --git 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
@@ -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
@@ -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 == []