From 0d05d411d9da7e70f99203874f4f99252f7a381b Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Tue, 19 Jul 2022 06:40:20 -0400
Subject: [PATCH] feat(duckdb): return the `ir.Table` instance from DuckDB's
 `register` API

---
 ibis/backends/duckdb/__init__.py | 39 ++++++++++++++++++++++----------
 1 file changed, 27 insertions(+), 12 deletions(-)

diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py
index 218e7f270fcc..57e0dd57d12e 100644
--- a/ibis/backends/duckdb/__init__.py
+++ b/ibis/backends/duckdb/__init__.py
@@ -15,6 +15,7 @@
 
 if TYPE_CHECKING:
     import duckdb
+    import ibis.expr.types as ir
 
 import ibis.expr.schema as sch
 from ibis.backends.base.sql.alchemy import BaseAlchemyBackend
@@ -28,7 +29,7 @@ class _ColumnMetadata(NamedTuple):
     type: dt.DataType
 
 
-_register = RegexDispatcher("_register")
+_generate_view_code = RegexDispatcher("_register")
 _dialect = sa.dialects.postgresql.dialect()
 
 
@@ -41,28 +42,36 @@ def _quote(name: str):
     return _dialect.identifier_preparer.quote(name)
 
 
-@_register.register(r"parquet://(?P<path>.+)", priority=10)
+@_generate_view_code.register(r"parquet://(?P<path>.+)", priority=10)
 def _parquet(_, path, table_name=None):
     path = Path(path).absolute()
-    table_name = _quote(table_name or _name_from_path(path))
-    return f"CREATE VIEW {table_name} as SELECT * from read_parquet('{path}')"
+    table_name = table_name or _name_from_path(path)
+    quoted_table_name = _quote(table_name)
+    return (
+        f"CREATE VIEW {quoted_table_name} as SELECT * from read_parquet('{path}')",  # noqa: E501
+        table_name,
+    )
 
 
-@_register.register(r"csv(?:\.gz)?://(?P<path>.+)", priority=10)
+@_generate_view_code.register(r"csv(?:\.gz)?://(?P<path>.+)", priority=10)
 def _csv(_, path, table_name=None):
     path = Path(path).absolute()
-    table_name = _quote(table_name or _name_from_path(path))
-    return f"CREATE VIEW {table_name} as SELECT * from read_csv_auto('{path}')"  # noqa: E501
+    table_name = table_name or _name_from_path(path)
+    quoted_table_name = _quote(table_name)
+    return (
+        f"CREATE VIEW {quoted_table_name} as SELECT * from read_csv_auto('{path}')",  # noqa: E501
+        table_name,
+    )
 
 
-@_register.register(r"(?:file://)?(?P<path>.+)", priority=9)
+@_generate_view_code.register(r"(?:file://)?(?P<path>.+)", priority=9)
 def _file(_, path, table_name=None):
     num_sep_chars = len(os.extsep)
     extension = "".join(Path(path).suffixes)[num_sep_chars:]
-    return _register(f"{extension}://{path}", table_name=table_name)
+    return _generate_view_code(f"{extension}://{path}", table_name=table_name)
 
 
-@_register.register(r".+", priority=1)
+@_generate_view_code.register(r".+", priority=1)
 def _default(_, **kwargs):
     raise ValueError(
         """
@@ -119,7 +128,7 @@ def register(
         self,
         path: str | Path,
         table_name: str | None = None,
-    ) -> None:
+    ) -> ir.Table:
         """Register an external file as a table in the current connection
         database
 
@@ -131,9 +140,15 @@ def register(
             Name for the created table.  Defaults to filename if not given.
             Any dashes in a user-provided or generated name will be
             replaced with underscores.
+
+        Returns
+        -------
+        ir.Table
+            The just-registered table
         """
-        view = _register(path, table_name=table_name)
+        view, table_name = _generate_view_code(path, table_name=table_name)
         self.con.execute(view)
+        return self.table(table_name)
 
     def fetch_from_cursor(
         self,