Repository: tobymao/sqlglot
Pull request: Fix: handling Dialect, t.Type[Dialect] now works in get_or_raise
Fixes #1082 

Problem statement:
Unable to use Dialect instances as a read/write arguments when transpiling
Although typing indicates that functions like `transpile` accept instances of `Dialect` as `read` and `write` arguments, it doesn't seem to work. I'd be happy to submit a merge request to fix this, but it seems like this issue goes deeper, with `Dialect.get_or_raise` actually being the root problem, as it seems to be used both to access a certain Dialect class, as well as generate Dialect objects. Defining typing for that method would probably highlight this issue well.

Since I need transpiling to work with my custom dialects, I'd be happy to work on this, but I don't want to mess something up before someone maintaining the package confirms that this is indeed an issue and that I'm on the right track.

p.s. awesome package!

Hints from issue discussion:
I can confirm that this is an issue. I'll post a fix soon, unless you want to play around with the project and submit a PR.

Thanks for the report! It's interesting that we haven't tested that properly yet.
Thank you, if it's a quick fix I would of course be happier if you do it, since I see there are a lot of places where this `get_or_raise` is used. But I will definitely be contributing to this project in the future, it fits my needs perfectly and I want to help with it :) Great work!

Commit message: More fixups

Patch to review:
diff --git a/sqlglot/__init__.py b/sqlglot/__init__.py
index 6861590a43..5738fd0c51 100644
--- a/sqlglot/__init__.py
+++ b/sqlglot/__init__.py
@@ -36,6 +36,8 @@
 if t.TYPE_CHECKING:
     T = t.TypeVar("T", bound=Expression)
 
+    DialectType = t.Union[str, Dialect, t.Type[Dialect]]
+
 __version__ = "10.6.0"
 
 pretty = False
@@ -45,9 +47,7 @@
 """The default schema used by SQLGlot (e.g. in the optimizer)."""
 
 
-def parse(
-    sql: str, read: t.Optional[str | Dialect] = None, **opts
-) -> t.List[t.Optional[Expression]]:
+def parse(sql: str, read: t.Optional[DialectType] = None, **opts) -> t.List[t.Optional[Expression]]:
     """
     Parses the given SQL string into a collection of syntax trees, one per parsed SQL statement.
 
@@ -76,7 +76,7 @@ def parse_one(
 @t.overload
 def parse_one(
     sql: str,
-    read: str | Dialect,
+    read: t.Optional[DialectType],
     into: t.Type[T],
     **opts,
 ) -> T:
@@ -96,7 +96,7 @@ def parse_one(
 @t.overload
 def parse_one(
     sql: str,
-    read: str | Dialect,
+    read: t.Optional[DialectType],
     into: t.Union[str, t.Collection[t.Union[str, t.Type[Expression]]]],
     **opts,
 ) -> Expression:
@@ -113,7 +113,7 @@ def parse_one(
 
 def parse_one(
     sql: str,
-    read: t.Optional[str | Dialect] = None,
+    read: t.Optional[DialectType] = None,
     into: t.Optional[exp.IntoType] = None,
     **opts,
 ) -> Expression:
@@ -147,8 +147,8 @@ def parse_one(
 
 def transpile(
     sql: str,
-    read: t.Optional[str | Dialect] = None,
-    write: t.Optional[str | Dialect] = None,
+    read: t.Optional[DialectType] = None,
+    write: t.Optional[DialectType] = None,
     identity: bool = True,
     error_level: t.Optional[ErrorLevel] = None,
     **opts,
diff --git a/sqlglot/dialects/dialect.py b/sqlglot/dialects/dialect.py
index 0c2bebabc7..58134fa8cd 100644
--- a/sqlglot/dialects/dialect.py
+++ b/sqlglot/dialects/dialect.py
@@ -122,9 +122,15 @@ class Dialect(metaclass=_Dialect):
     def get_or_raise(cls, dialect):
         if not dialect:
             return cls
+        if isinstance(dialect, _Dialect):
+            return dialect
+        if isinstance(dialect, Dialect):
+            return dialect.__class__
+
         result = cls.get(dialect)
         if not result:
             raise ValueError(f"Unknown dialect '{dialect}'")
+
         return result
 
     @classmethod
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 7c1a116707..ebb0d8a94a 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -32,6 +32,7 @@
 from sqlglot.tokens import Token
 
 if t.TYPE_CHECKING:
+    from sqlglot import DialectType
     from sqlglot.dialects.dialect import Dialect
 
     IntoType = t.Union[
@@ -427,7 +428,7 @@ def __str__(self):
     def __repr__(self):
         return self._to_s()
 
-    def sql(self, dialect: Dialect | str | None = None, **opts) -> str:
+    def sql(self, dialect: t.Optional[DialectType] = None, **opts) -> str:
         """
         Returns SQL string representation of this tree.
 
@@ -2456,7 +2457,7 @@ class Type(AutoName):
 
     @classmethod
     def build(
-        cls, dtype: str | DataType.Type, dialect: t.Optional[str | Dialect] = None, **kwargs
+        cls, dtype: str | DataType.Type, dialect: t.Optional[DialectType] = None, **kwargs
     ) -> DataType:
         from sqlglot import parse_one
 
@@ -3476,7 +3477,7 @@ def maybe_parse(
     sql_or_expression: str | Expression,
     *,
     into: t.Optional[IntoType] = None,
-    dialect: t.Optional[str] = None,
+    dialect: t.Optional[DialectType] = None,
     prefix: t.Optional[str] = None,
     **opts,
 ) -> Expression:
@@ -4373,7 +4374,7 @@ def _expand(node: Expression):
     return expression.transform(_expand, copy=copy)
 
 
-def func(name: str, *args, dialect: t.Optional[Dialect | str] = None, **kwargs) -> Func:
+def func(name: str, *args, dialect: t.Optional[DialectType] = None, **kwargs) -> Func:
     """
     Returns a Func expression.
 
diff --git a/sqlglot/lineage.py b/sqlglot/lineage.py
index 4e7eab847d..c37ed9959c 100644
--- a/sqlglot/lineage.py
+++ b/sqlglot/lineage.py
@@ -36,7 +36,7 @@ def lineage(
     schema: t.Optional[t.Dict | Schema] = None,
     sources: t.Optional[t.Dict[str, str | exp.Subqueryable]] = None,
     rules: t.Sequence[t.Callable] = (qualify_tables, qualify_columns),
-    dialect: t.Optional[str] = None,
+    dialect: t.Optional[DialectType] = None,
 ) -> Node:
     """Build the lineage graph for a column of a SQL query.
 
@@ -126,7 +126,7 @@ class LineageHTML:
     def __init__(
         self,
         node: Node,
-        dialect: t.Optional[str] = None,
+        dialect: t.Optional[DialectType] = None,
         imports: bool = True,
         **opts: t.Any,
     ):
diff --git a/sqlglot/schema.py b/sqlglot/schema.py
index f6f38836dd..c525fccf37 100644
--- a/sqlglot/schema.py
+++ b/sqlglot/schema.py
@@ -153,7 +153,7 @@ def __init__(
         self,
         schema: t.Optional[t.Dict] = None,
         visible: t.Optional[t.Dict] = None,
-        dialect: t.Optional[str] = None,
+        dialect: t.Optional[DialectType] = None,
     ) -> None:
         self.dialect = dialect
         self.visible = visible or {}
diff --git a/tests/dialects/test_dialect.py b/tests/dialects/test_dialect.py
index 5a13655c1a..446247c364 100644
--- a/tests/dialects/test_dialect.py
+++ b/tests/dialects/test_dialect.py
@@ -1,6 +1,7 @@
 import unittest
 
 from sqlglot import Dialect, Dialects, ErrorLevel, UnsupportedError, parse_one
+from sqlglot.dialects import Hive
 
 
 class Validator(unittest.TestCase):
@@ -67,6 +68,11 @@ def test_enum(self):
             self.assertIsNotNone(Dialect.get_or_raise(dialect))
             self.assertIsNotNone(Dialect[dialect.value])
 
+    def test_get_or_raise(self):
+        self.assertEqual(Dialect.get_or_raise(Hive), Hive)
+        self.assertEqual(Dialect.get_or_raise(Hive()), Hive)
+        self.assertEqual(Dialect.get_or_raise("hive"), Hive)
+
     def test_cast(self):
         self.validate_all(
             "CAST(a AS TEXT)",


Code review comments:
diff_hunk: @@ -36,6 +36,8 @@
 if t.TYPE_CHECKING:
     T = t.TypeVar("T", bound=Expression)
 
+    DialectType = t.Union[str, Dialect, t.Type[Dialect]]
path: sqlglot/__init__.py
text: @user1:
```python
DialectType = t.Union[str, Dialect, t.Type[Dialect], None]
```
Then you don't need to wrap it in Optional everywhere.

@author:
I excluded `None` on purpose, so that a user can easily see that the "dialect" parameter is optional in a declaration.

@user1:
dialect=None is pretty clear lol

@user1:
Also, I don't think this should be defined in `sqlglot.__init__.py`. We're a little to easygoing with circular imports lol. Perhaps it should be defined in `dialect.py`?

@author:
You have a point, seems like I didn't take that into account. Making the change, thanks.

@author:
> Also, I don't think this should be defined in sqlglot.__init__.py. We're a little to easygoing with circular imports lol. Perhaps it should be defined in dialect.py?

Yeah agreed, makes sense to put it there.

@author:
Done