(airflow) branch main updated: PostgresDialect should use index instead of name in get_column_names and get_primary_keys (#54832)

[po...@apache.org]

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

potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git


The following commit(s) were added to refs/heads/main by this push:
     new 2f58d7635eb PostgresDialect should use index instead of name in get_column_names and get_primary_keys (#54832)
2f58d7635eb is described below

commit 2f58d7635eb9d7f75b5f7f4eb7b9c68d214d9ab5
Author: David Blain <in...@dabla.be>
AuthorDate: Fri Aug 22 17:22:57 2025 +0200

    PostgresDialect should use index instead of name in get_column_names and get_primary_keys (#54832)
    
    * fix: PostgresDialect should use index instead of name as returned rows are list of tuples, not list of dicts
    
    * refactor: Updated TestPostgresDialect
    
    * Revert "refactor: Updated TestPostgresDialect"
    
    This reverts commit 244b889b49cfdc376080c095024e14d2d2898bdf.
    
    * refactor: Refactored get_column_names in PostgresDialect so it allows back filtering through predicate
    
    * refactor: Also take into account autoincrement columns in PostgresDialect
    
    * refactor: Fixed mocking of records in TestPostgresDialect
    
    * refactor: Reformatted get_column_names
---
 .../providers/postgres/dialects/postgres.py        | 30 +++++++++++++++++-----
 .../tests/unit/postgres/dialects/test_postgres.py  | 10 ++++----
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/providers/postgres/src/airflow/providers/postgres/dialects/postgres.py b/providers/postgres/src/airflow/providers/postgres/dialects/postgres.py
index 6c1d3104d10..9e31d39b5e8 100644
--- a/providers/postgres/src/airflow/providers/postgres/dialects/postgres.py
+++ b/providers/postgres/src/airflow/providers/postgres/dialects/postgres.py
@@ -42,10 +42,10 @@ class PostgresDialect(Dialect):
         if schema is None:
             table, schema = self.extract_schema_from_table(table)
         pk_columns = [
-            row["column_name"]
+            row[0]
             for row in self.get_records(
                 """
-                  select kcu.column_name as column_name
+                  select kcu.column_name
                   from information_schema.table_constraints tco
                            join information_schema.key_column_usage kcu
                                 on kcu.constraint_name = tco.constraint_name
@@ -61,29 +61,45 @@ class PostgresDialect(Dialect):
         ]
         return pk_columns or None
 
+    @staticmethod
+    def _to_row(row):
+        return {
+            "name": row[0],
+            "type": row[1],
+            "nullable": row[2].casefold() == "yes",
+            "default": row[3],
+            "autoincrement": row[4].casefold() == "always",
+            "identity": row[5].casefold() == "yes",
+        }
+
     @lru_cache(maxsize=None)
     def get_column_names(
         self, table: str, schema: str | None = None, predicate: Callable[[T], bool] = lambda column: True
     ) -> list[str] | None:
         if schema is None:
             table, schema = self.extract_schema_from_table(table)
+
         column_names = list(
-            row["column_name"]
+            row["name"]
             for row in filter(
                 predicate,
-                self.get_records(
-                    """
+                map(
+                    self._to_row,
+                    self.get_records(
+                        """
                         select column_name,
                                data_type,
                                is_nullable,
                                column_default,
-                               ordinal_position
+                               is_generated,
+                               is_identity
                         from information_schema.columns
                         where table_schema = %s
                           and table_name = %s
                         order by ordinal_position
                         """,
-                    (self.unescape_word(schema), self.unescape_word(table)),
+                        (self.unescape_word(schema), self.unescape_word(table)),
+                    ),
                 ),
             )
         )
diff --git a/providers/postgres/tests/unit/postgres/dialects/test_postgres.py b/providers/postgres/tests/unit/postgres/dialects/test_postgres.py
index 5406805420b..999386baac1 100644
--- a/providers/postgres/tests/unit/postgres/dialects/test_postgres.py
+++ b/providers/postgres/tests/unit/postgres/dialects/test_postgres.py
@@ -30,12 +30,12 @@ class TestPostgresDialect:
             assert "hollywood" in parameters, "Missing 'schema' in parameters"
             assert "actors" in parameters, "Missing 'table' in parameters"
             if "kcu." in sql:
-                return [{"column_name": "id"}]
+                return [("id",)]
             return [
-                {"column_name": "id", "identity": True},
-                {"column_name": "name"},
-                {"column_name": "firstname"},
-                {"column_name": "age"},
+                ("id", None, "NO", None, "ALWAYS", "YES"),
+                ("name", None, "YES", None, "NEVER", "NO"),
+                ("firstname", None, "YES", None, "NEVER", "NO"),
+                ("age", None, "YES", None, "NEVER", "NO"),
             ]
 
         self.test_db_hook = MagicMock(placeholder="?", spec=DbApiHook)