From 4d90a359a21973a4180bfc4f8285ae3a3e123136 Mon Sep 17 00:00:00 2001
From: Damian Owsianny <damian3031@gmail.com>
Date: Wed, 25 Sep 2024 23:24:33 +0200
Subject: [PATCH] Refactor get_indexes to return partition names

---
 trino/sqlalchemy/dialect.py | 20 ++++++++++++++++++--
 1 file changed, 18 insertions(+), 2 deletions(-)

diff --git a/trino/sqlalchemy/dialect.py b/trino/sqlalchemy/dialect.py
index edcc372c..2236c6c0 100644
--- a/trino/sqlalchemy/dialect.py
+++ b/trino/sqlalchemy/dialect.py
@@ -204,6 +204,22 @@ def _get_columns(self, connection: Connection, table_name: str, schema: str = No
             columns.append(column)
         return columns
 
+    def _get_partitions(
+        self,
+        connection: Connection,
+        table_name: str,
+        schema: str = None
+    ) -> List[Dict[str, List[Any]]]:
+        schema = schema or self._get_default_schema_name(connection)
+        query = dedent(
+            f"""
+            SELECT * FROM {schema}."{table_name}$partitions"
+        """
+        ).strip()
+        res = connection.execute(sql.text(query))
+        partition_names = [desc[0] for desc in res.cursor.description]
+        return partition_names
+
     def get_pk_constraint(self, connection: Connection, table_name: str, schema: str = None, **kw) -> Dict[str, Any]:
         """Trino has no support for primary keys. Returns a dummy"""
         return dict(name=None, constrained_columns=[])
@@ -299,7 +315,7 @@ def get_indexes(self, connection: Connection, table_name: str, schema: str = Non
 
         partitioned_columns = None
         try:
-            partitioned_columns = self._get_columns(connection, f"{table_name}$partitions", schema, **kw)
+            partitioned_columns = self._get_partitions(connection, f"{table_name}", schema)
         except Exception as e:
             # e.g. it's not a Hive table or an unpartitioned Hive table
             logger.debug("Couldn't fetch partition columns. schema: %s, table: %s, error: %s", schema, table_name, e)
@@ -307,7 +323,7 @@ def get_indexes(self, connection: Connection, table_name: str, schema: str = Non
             return []
         partition_index = dict(
             name="partition",
-            column_names=[col["name"] for col in partitioned_columns],
+            column_names=partitioned_columns,
             unique=False
         )
         return [partition_index]