From 933fb321eb0580d66fb8148e5e7636a15233865b Mon Sep 17 00:00:00 2001
From: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
Date: Thu, 17 Aug 2023 04:20:26 -0400
Subject: [PATCH] fix(snowflake): allow backend to choose how to prefix table
 names during compilation

---
 ibis/backends/base/sql/alchemy/__init__.py | 20 +++++---------------
 ibis/backends/snowflake/__init__.py        | 17 +++++++++++++++++
 ibis/backends/trino/__init__.py            |  8 ++++++++
 3 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/ibis/backends/base/sql/alchemy/__init__.py b/ibis/backends/base/sql/alchemy/__init__.py
index 76586b89f28c..7d17f7ef799d 100644
--- a/ibis/backends/base/sql/alchemy/__init__.py
+++ b/ibis/backends/base/sql/alchemy/__init__.py
@@ -930,9 +930,13 @@ def _get_sqla_table(
 
         with self._use_schema(ident, current_db, current_schema):
             result = self._table_from_schema(name, schema=ibis_schema)
-        result.schema = schema
+        result.schema = self._get_schema_for_table(qualname=ident, schema=schema)
         return result
 
+    @abc.abstractmethod
+    def _get_schema_for_table(self, *, qualname: str, schema: str) -> str:
+        """Choose whether to prefix a table with its fully qualified path or schema."""
+
     def drop_table(
         self, name: str, database: str | None = None, force: bool = False
     ) -> None:
@@ -943,17 +947,3 @@ def drop_table(
         drop_stmt = "DROP TABLE" + (" IF EXISTS" * force) + f" {name}"
         with self.begin() as con:
             con.exec_driver_sql(drop_stmt)
-
-
-@compiles(sa.Table, "snowflake")
-def compile_table(element, compiler, **kw):
-    """Override compilation of leaf tables.
-
-    The override is necessary because the dialect does not handle database
-    hierarchies and/or quoting properly.
-    """
-    schema = element.schema
-    name = compiler.preparer.quote_identifier(element.name)
-    if schema is not None:
-        return f"{schema}.{name}"
-    return name
diff --git a/ibis/backends/snowflake/__init__.py b/ibis/backends/snowflake/__init__.py
index 136e0d512e80..393a8fac5c2d 100644
--- a/ibis/backends/snowflake/__init__.py
+++ b/ibis/backends/snowflake/__init__.py
@@ -805,6 +805,23 @@ def read_json(
 
         return self.table(table)
 
+    def _get_schema_for_table(self, *, qualname: str, schema: str) -> str:
+        return qualname
+
+
+@compiles(sa.Table, "snowflake")
+def compile_table(element, compiler, **kw):
+    """Override compilation of leaf tables.
+
+    The override is necessary because snowflake-sqlalchemy does not handle
+    quoting databases and schemas correctly.
+    """
+    schema = element.schema
+    name = compiler.preparer.quote_identifier(element.name)
+    if schema is not None:
+        return f"{schema}.{name}"
+    return name
+
 
 @compiles(sa.sql.Join, "snowflake")
 def compile_join(element, compiler, **kw):
diff --git a/ibis/backends/trino/__init__.py b/ibis/backends/trino/__init__.py
index 9abc2770a6b7..3b9b35ba7c58 100644
--- a/ibis/backends/trino/__init__.py
+++ b/ibis/backends/trino/__init__.py
@@ -363,3 +363,11 @@ def _table_from_schema(
             trino_catalog=database or self.current_database,
             **kwargs,
         )
+
+    def _get_schema_for_table(self, *, qualname: str, schema: str) -> str:
+        """Trino compiles the `trino_catalog` argument into `sa.Table`.
+
+        This means we only need the schema and not the fully qualified
+        $catalog.$schema identifier.
+        """
+        return schema