Improve read_database_uri overloads

pola-rs · Oct 20, 2024 · 48a8531 · 48a8531
1 parent f3eba22
commit 48a8531
Showing 1 changed file with 30 additions and 0 deletions.
diff --git a/py-polars/polars/io/database/functions.py b/py-polars/polars/io/database/functions.py
@@ -263,6 +263,36 @@ def read_database(
         )
 
 
+@overload
+def read_database_uri(
+    query: str,
+    uri: str,
+    *,
+    partition_on: str | None = None,
+    partition_range: tuple[int, int] | None = None,
+    partition_num: int | None = None,
+    protocol: str | None = None,
+    engine: Literal["adbc"],
+    schema_overrides: SchemaDict | None = None,
+    execute_options: dict[str, Any] | None = None,
+) -> DataFrame: ...
+
+
+@overload
+def read_database_uri(
+    query: list[str] | str,
+    uri: str,
+    *,
+    partition_on: str | None = None,
+    partition_range: tuple[int, int] | None = None,
+    partition_num: int | None = None,
+    protocol: str | None = None,
+    engine: Literal["connectorx"] | None = None,
+    schema_overrides: SchemaDict | None = None,
+    execute_options: None = None,
+) -> DataFrame: ...
+
+
 def read_database_uri(
     query: list[str] | str,
     uri: str,