[superset] branch master updated: chore(presto): Expose schema and indexes to _partition_query method (#23674)

[jo...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

johnbodley pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 100ccb759c chore(presto): Expose schema and indexes to _partition_query method (#23674)
100ccb759c is described below

commit 100ccb759c16706c9d7af4a264ab056ccefb5e01
Author: John Bodley <45...@users.noreply.github.com>
AuthorDate: Sat Apr 15 06:24:54 2023 +1200

    chore(presto): Expose schema and indexes to _partition_query method (#23674)
    
    Co-authored-by: Michael S. Molina <70...@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 792ef94735..a7d4a13be3 100644
--- a/superset/db_engine_specs/hive.py
+++ b/superset/db_engine_specs/hive.py
@@ -461,6 +461,8 @@ class HiveEngineSpec(PrestoEngineSpec):
     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 12183c2b19..0889dd653b 100644
--- a/superset/db_engine_specs/presto.py
+++ b/superset/db_engine_specs/presto.py
@@ -423,24 +423,30 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta):
         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 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta):
             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 @@ class PrestoBaseEngineSpec(BaseEngineSpec, metaclass=ABCMeta):
                 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 @@ class PrestoEngineSpec(PrestoBaseEngineSpec):
                         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 0a18a3e3e8..c370a3585d 100644
--- a/superset/db_engine_specs/trino.py
+++ b/superset/db_engine_specs/trino.py
@@ -84,6 +84,8 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
                         if schema_name and "." not in table_name
                         else table_name
                     ),
+                    schema=schema_name,
+                    indexes=indexes,
                     database=database,
                 ),
             }