From 4f303b2329c85fe533ae4fc43f7f0ee8222760a5 Mon Sep 17 00:00:00 2001
From: John Bodley <4567245+john-bodley@users.noreply.github.com>
Date: Sat, 15 Apr 2023 06:24:54 +1200
Subject: [PATCH] chore(presto): Expose schema and indexes to _partition_query
 method (#23674)

Co-authored-by: Michael S. Molina <70410625+michael-s-molina@users.noreply.github.com>
---
 superset/db_engine_specs/hive.py   |  2 ++
 superset/db_engine_specs/presto.py | 28 ++++++++++++++++++++++------
 superset/db_engine_specs/trino.py  |  2 ++
 3 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/superset/db_engine_specs/hive.py b/superset/db_engine_specs/hive.py
index 792ef947350a..a7d4a13be3b3 100644
--- a/superset/db_engine_specs/hive.py
+++ b/superset/db_engine_specs/hive.py
@@ -461,6 +461,8 @@ def _latest_partition_from_df(cls, df: pd.DataFrame) -> Optional[List[str]]:
     def _partition_query(  # pylint: disable=too-many-arguments
         cls,
         table_name: str,
+        schema: Optional[str],
+        indexes: List[Dict[str, Any]],
         database: "Database",
         limit: int = 0,
         order_by: Optional[List[Tuple[str, bool]]] = None,
diff --git a/superset/db_engine_specs/presto.py b/superset/db_engine_specs/presto.py
index 12183c2b1913..0889dd653b8e 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -423,24 +423,30 @@ def get_function_names(cls, database: Database) -> List[str]:
         return database.get_df("SHOW FUNCTIONS")["Function"].tolist()
 
     @classmethod
-    def _partition_query(  # pylint: disable=too-many-arguments,too-many-locals
+    def _partition_query(  # pylint: disable=too-many-arguments,too-many-locals,unused-argument
         cls,
         table_name: str,
+        schema: Optional[str],
+        indexes: List[Dict[str, Any]],
         database: Database,
         limit: int = 0,
         order_by: Optional[List[Tuple[str, bool]]] = None,
         filters: Optional[Dict[Any, Any]] = None,
     ) -> str:
-        """Returns a partition query
+        """
+        Return a partition query.
+
+        Note the unused arguments are exposed for sub-classing purposes where custom
+        integrations may require the schema, indexes, etc. to build the partition query.
 
         :param table_name: the name of the table to get partitions from
-        :type table_name: str
+        :param schema: the schema name
+        :param indexes: the indexes associated with the table
+        :param database: the database the query will be run against
         :param limit: the number of partitions to be returned
-        :type limit: int
         :param order_by: a list of tuples of field name and a boolean
             that determines if that field should be sorted in descending
             order
-        :type order_by: list of (str, bool) tuples
         :param filters: dict of field name and filter value combinations
         """
         limit_clause = "LIMIT {}".format(limit) if limit else ""
@@ -566,6 +572,8 @@ def latest_partition(
             df=database.get_df(
                 sql=cls._partition_query(
                     table_name,
+                    schema,
+                    indexes,
                     database,
                     limit=1,
                     order_by=[(column_name, True) for column_name in column_names],
@@ -620,7 +628,13 @@ def latest_sub_partition(
                 field_to_return = field
 
         sql = cls._partition_query(
-            table_name, database, 1, [(field_to_return, True)], kwargs
+            table_name,
+            schema,
+            indexes,
+            database,
+            limit=1,
+            order_by=[(field_to_return, True)],
+            filters=kwargs,
         )
         df = database.get_df(sql, schema)
         if df.empty:
@@ -1214,6 +1228,8 @@ def extra_table_metadata(
                         if schema_name and "." not in table_name
                         else table_name
                     ),
+                    schema=schema_name,
+                    indexes=indexes,
                     database=database,
                 ),
             }
diff --git a/superset/db_engine_specs/trino.py b/superset/db_engine_specs/trino.py
index 0a18a3e3e8fe..c370a3585d5b 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -84,6 +84,8 @@ def extra_table_metadata(
                         if schema_name and "." not in table_name
                         else table_name
                     ),
+                    schema=schema_name,
+                    indexes=indexes,
                     database=database,
                 ),
             }