Skip to content

Commit

Permalink
chore(presto): Expose schema and indexes to _partition_query method (a…
Browse files Browse the repository at this point in the history
…pache#23674)

Co-authored-by: Michael S. Molina <70410625+michael-s-molina@users.noreply.github.com>
  • Loading branch information
2 people authored and sebastianliebscher committed Apr 28, 2023
1 parent c268c4e commit 4f303b2
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 6 deletions.
2 changes: 2 additions & 0 deletions superset/db_engine_specs/hive.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
28 changes: 22 additions & 6 deletions superset/db_engine_specs/presto.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 ""
Expand Down Expand Up @@ -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],
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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,
),
}
Expand Down
2 changes: 2 additions & 0 deletions superset/db_engine_specs/trino.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
),
}
Expand Down

0 comments on commit 4f303b2

Please sign in to comment.