From 4f193ed5eabcd2eb37a2905fa001e09cceb2ee33 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jergu=C5=A1=20Lejko?= <jerguslejko@users.noreply.github.com>
Date: Fri, 18 Sep 2026 10:09:43 +0200
Subject: [PATCH] Refactor: replace `dialect_of!` table partition check with
 `Dialect` trait method

---
 src/dialect/generic.rs            |  4 +++
 src/dialect/mod.rs                | 13 ++++++++++
 src/dialect/mysql.rs              |  5 ++++
 src/parser/mod.rs                 |  2 +-
 tests/sqlparser_custom_dialect.rs | 41 +++++++++++++++++++++++++++++++
 5 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/src/dialect/generic.rs b/src/dialect/generic.rs
index d408cb181a..6b6b8b55dc 100644
--- a/src/dialect/generic.rs
+++ b/src/dialect/generic.rs
@@ -285,6 +285,10 @@ impl Dialect for GenericDialect {
         true
     }
 
+    fn supports_table_partitions(&self) -> bool {
+        true
+    }
+
     fn supports_select_format(&self) -> bool {
         true
     }
diff --git a/src/dialect/mod.rs b/src/dialect/mod.rs
index 7c4744c5a7..b4e8bacaa5 100644
--- a/src/dialect/mod.rs
+++ b/src/dialect/mod.rs
@@ -1764,6 +1764,19 @@ pub trait Dialect: Debug + Any {
         false
     }
 
+    /// Returns true if this dialect supports the `PARTITION` clause on a table factor,
+    /// restricting a query to an explicit list of partitions.
+    ///
+    /// Example:
+    /// ```sql
+    /// SELECT * FROM employees PARTITION (p0, p1)
+    /// ```
+    ///
+    /// [MySQL](https://dev.mysql.com/doc/refman/8.4/en/partitioning-selection.html)
+    fn supports_table_partitions(&self) -> bool {
+        false
+    }
+
     /// Returns true if this dialect supports the `FORMAT` clause in `SELECT` statements.
     ///
     /// Example:
diff --git a/src/dialect/mysql.rs b/src/dialect/mysql.rs
index f5c50d0d86..ffd6fb3728 100644
--- a/src/dialect/mysql.rs
+++ b/src/dialect/mysql.rs
@@ -203,6 +203,11 @@ impl Dialect for MySqlDialect {
         true
     }
 
+    /// See: <https://dev.mysql.com/doc/refman/8.4/en/partitioning-selection.html>
+    fn supports_table_partitions(&self) -> bool {
+        true
+    }
+
     fn supports_comment_optimizer_hint(&self) -> bool {
         true
     }
diff --git a/src/parser/mod.rs b/src/parser/mod.rs
index 15f135fffa..754e7f26b6 100644
--- a/src/parser/mod.rs
+++ b/src/parser/mod.rs
@@ -16753,7 +16753,7 @@ impl<'a> Parser<'a> {
                 _ => None,
             };
 
-            let partitions: Vec<Ident> = if dialect_of!(self is MySqlDialect | GenericDialect)
+            let partitions: Vec<Ident> = if self.dialect.supports_table_partitions()
                 && self.parse_keyword(Keyword::PARTITION)
             {
                 self.parse_parenthesized_identifiers()?
diff --git a/tests/sqlparser_custom_dialect.rs b/tests/sqlparser_custom_dialect.rs
index cee604aca7..0a5cd2cc5b 100644
--- a/tests/sqlparser_custom_dialect.rs
+++ b/tests/sqlparser_custom_dialect.rs
@@ -156,6 +156,47 @@ fn test_map_syntax_not_support_default() -> Result<(), ParserError> {
     Ok(())
 }
 
+#[test]
+fn custom_dialect_supports_table_partitions() -> Result<(), ParserError> {
+    #[derive(Debug)]
+    struct MyDialect {}
+
+    impl Dialect for MyDialect {
+        fn is_identifier_start(&self, ch: char) -> bool {
+            is_identifier_start(ch)
+        }
+
+        fn is_identifier_part(&self, ch: char) -> bool {
+            is_identifier_part(ch)
+        }
+
+        fn supports_table_partitions(&self) -> bool {
+            true
+        }
+    }
+
+    let sql = "SELECT * FROM employees PARTITION (p0, p1)";
+    let ast = Parser::parse_sql(&MyDialect {}, sql)?;
+    assert_eq!(sql, &format!("{}", ast[0]));
+
+    // A dialect that does not opt in still rejects the clause.
+    #[derive(Debug)]
+    struct WithoutPartitions {}
+
+    impl Dialect for WithoutPartitions {
+        fn is_identifier_start(&self, ch: char) -> bool {
+            is_identifier_start(ch)
+        }
+
+        fn is_identifier_part(&self, ch: char) -> bool {
+            is_identifier_part(ch)
+        }
+    }
+
+    assert!(Parser::parse_sql(&WithoutPartitions {}, sql).is_err());
+    Ok(())
+}
+
 fn is_identifier_start(ch: char) -> bool {
     ch.is_ascii_lowercase() || ch.is_ascii_uppercase() || ch == '_'
 }