nits

daft/daft.pyi

@@ -611,7 +611,7 @@ class ScanTask:
         ...
     @staticmethod
     def sql_scan_task(
-        path: str,
+        url: str,
         file_format: FileFormatConfig,
         schema: PySchema,
         num_rows: int | None,

daft/sql/sql_connection.py

@@ -2,22 +2,21 @@
 
 import logging
 import warnings
-from typing import TYPE_CHECKING, Callable
+from typing import Callable
 from urllib.parse import urlparse
 
 import pyarrow as pa
-
-if TYPE_CHECKING:
-    from sqlalchemy.engine import Connection
+from sqlalchemy.engine import Connection
 
 logger = logging.getLogger(__name__)
 
 
 class SQLConnection:
-    def __init__(self, conn: str | Callable[[], Connection], driver: str, dialect: str) -> None:
+    def __init__(self, conn: str | Callable[[], Connection], driver: str, dialect: str, url: str) -> None:
         self.conn = conn
         self.dialect = dialect
         self.driver = driver
+        self.url = url
 
     def __repr__(self) -> str:
         return f"SQLConnection(conn={self.conn})"
@@ -29,17 +28,20 @@ def from_url(cls, url: str) -> SQLConnection:
             dialect, driver = scheme.split("+")
         else:
             dialect, driver = scheme, ""
-        return SQLConnection(url, driver, dialect)
+        return cls(url, driver, dialect, url)
 
     @classmethod
     def from_connection_factory(cls, conn_factory: Callable[[], Connection]) -> SQLConnection:
         try:
             with conn_factory() as connection:
-                if not hasattr(connection, "engine"):
-                    raise ValueError("The connection factory must return a SQLAlchemy connection object.")
+                if not isinstance(connection, Connection):
+                    raise ValueError(
+                        f"Connection factory must return a SQLAlchemy connection object, got: {type(connection)}"
+                    )
                 dialect = connection.engine.dialect.name
                 driver = connection.engine.driver
-            return SQLConnection(conn_factory, driver, dialect)
+                url = connection.engine.url.render_as_string()
+            return cls(conn_factory, driver, dialect, url)
         except Exception as e:
             raise ValueError(f"Unexpected error while calling the connection factory: {e}") from e
 

daft/sql/sql_scan.py

@@ -86,7 +86,7 @@ def to_scan_tasks(self, pushdowns: Pushdowns) -> Iterator[ScanTask]:
 
             scan_tasks.append(
                 ScanTask.sql_scan_task(
-                    path=repr(self.conn),
+                    url=self.conn.url,
                     file_format=file_format_config,
                     schema=self._schema._schema,
                     num_rows=None,
@@ -207,7 +207,7 @@ def _single_scan_task(self, pushdowns: Pushdowns, total_rows: int | None, total_
         return iter(
             [
                 ScanTask.sql_scan_task(
-                    path=repr(self.conn),
+                    url=self.conn.url,
                     file_format=file_format_config,
                     schema=self._schema._schema,
                     num_rows=total_rows,

src/daft-scan/src/python.rs

@@ -322,7 +322,7 @@ pub mod pylib {
         #[allow(clippy::too_many_arguments)]
         #[staticmethod]
         pub fn sql_scan_task(
-            path: String,
+            url: String,
             file_format: PyFileFormatConfig,
             schema: PySchema,
             storage_config: PyStorageConfig,
@@ -335,7 +335,7 @@ pub mod pylib {
                 .map(|s| TableStatistics::from_stats_table(&s.table))
                 .transpose()?;
             let data_source = DataFileSource::DatabaseDataSource {
-                path,
+                path: url,
                 chunk_spec: None,
                 size_bytes,
                 metadata: num_rows.map(|n| TableMetadata { length: n as usize }),