apache-airflow-providers-postgres 6.3.0__py3-none-any.whl → 6.4.1rc1__py3-none-any.whl

airflow/providers/postgres/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "6.3.0"
+__version__ = "6.4.1"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/postgres/dialects/postgres.py

@@ -41,24 +41,21 @@ class PostgresDialect(Dialect):
         """
         if schema is None:
             table, schema = self.extract_schema_from_table(table)
-        pk_columns = [
-            row[0]
-            for row in self.get_records(
-                """
-                  select kcu.column_name
-                  from information_schema.table_constraints tco
-                           join information_schema.key_column_usage kcu
-                                on kcu.constraint_name = tco.constraint_name
-                                    and kcu.constraint_schema = tco.constraint_schema
-                                    and kcu.constraint_name = tco.constraint_name
-                  where tco.constraint_type = 'PRIMARY KEY'
-                    and kcu.table_schema = %s
-                    and kcu.table_name = %s
-                  order by kcu.ordinal_position
-                  """,
-                (self.unescape_word(schema), self.unescape_word(table)),
-            )
-        ]
+        table = self.unescape_word(table) or table
+        schema = self.unescape_word(schema) if schema else None
+        query = """
+            select kcu.column_name
+            from information_schema.table_constraints tco
+                    join information_schema.key_column_usage kcu
+                        on kcu.constraint_name = tco.constraint_name
+                            and kcu.constraint_schema = tco.constraint_schema
+                            and kcu.constraint_name = tco.constraint_name
+            where tco.constraint_type = 'PRIMARY KEY'
+            and kcu.table_schema = %s
+            and kcu.table_name = %s
+            order by kcu.ordinal_position
+        """
+        pk_columns = [row[0] for row in self.get_records(query, (schema, table))]
         return pk_columns or None
 
     @staticmethod
@@ -78,31 +75,27 @@ class PostgresDialect(Dialect):
     ) -> list[str] | None:
         if schema is None:
             table, schema = self.extract_schema_from_table(table)
-
-        column_names = list(
-            row["name"]
-            for row in filter(
-                predicate,
-                map(
-                    self._to_row,
-                    self.get_records(
-                        """
-                        select column_name,
-                               data_type,
-                               is_nullable,
-                               column_default,
-                               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)),
-                    ),
-                ),
-            )
-        )
+        table = self.unescape_word(table) or table
+        schema = self.unescape_word(schema) if schema else None
+        query = """
+            select column_name,
+                    data_type,
+                    is_nullable,
+                    column_default,
+                    is_generated,
+                    is_identity
+            from information_schema.columns
+            where table_schema = %s
+                and table_name = %s
+            order by ordinal_position
+        """
+        column_names = []
+        for row in map(
+            self._to_row,
+            self.get_records(query, (schema, table)),
+        ):
+            if predicate(row):
+                column_names.append(row["name"])
         self.log.debug("Column names for table '%s': %s", table, column_names)
         return column_names
 

airflow/providers/postgres/get_provider_info.py

@@ -65,4 +65,18 @@ def get_provider_info():
                 "handler": "airflow.providers.postgres.assets.postgres.sanitize_uri",
             }
         ],
+        "config": {
+            "postgres": {
+                "description": "Configuration for Postgres hooks and operators.\n",
+                "options": {
+                    "azure_oauth_scope": {
+                        "description": "The scope to use while retrieving Oauth token for Postgres Flexible Server\nfrom Azure Entra authentication.\n",
+                        "version_added": "6.4.0",
+                        "type": "string",
+                        "example": None,
+                        "default": "https://ossrdbms-aad.database.windows.net/.default",
+                    }
+                },
+            }
+        },
     }

airflow/providers/postgres/hooks/postgres.py

@@ -24,16 +24,17 @@ from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeAlias, cast, overload
 
 import psycopg2
-import psycopg2.extensions
 import psycopg2.extras
 from more_itertools import chunked
 from psycopg2.extras import DictCursor, NamedTupleCursor, RealDictCursor, execute_batch
 from sqlalchemy.engine import URL
 
+from airflow.configuration import conf
 from airflow.exceptions import (
     AirflowException,
     AirflowOptionalProviderFeatureException,
 )
+from airflow.providers.common.compat.sdk import Connection
 from airflow.providers.common.sql.hooks.sql import DbApiHook
 from airflow.providers.postgres.dialects.postgres import PostgresDialect
 
@@ -64,11 +65,6 @@ if TYPE_CHECKING:
     if USE_PSYCOPG3:
         from psycopg.errors import Diagnostic
 
-    try:
-        from airflow.sdk import Connection
-    except ImportError:
-        from airflow.models.connection import Connection  # type: ignore[assignment]
-
 CursorType: TypeAlias = DictCursor | RealDictCursor | NamedTupleCursor
 CursorRow: TypeAlias = dict[str, Any] | tuple[Any, ...]
 
@@ -156,7 +152,9 @@ class PostgresHook(DbApiHook):
         "aws_conn_id",
         "sqlalchemy_scheme",
         "sqlalchemy_query",
+        "azure_conn_id",
     }
+    default_azure_oauth_scope = "https://ossrdbms-aad.database.windows.net/.default"
 
     def __init__(
         self, *args, options: str | None = None, enable_log_db_messages: bool = False, **kwargs
@@ -177,6 +175,8 @@ class PostgresHook(DbApiHook):
         query = conn.extra_dejson.get("sqlalchemy_query", {})
         if not isinstance(query, dict):
             raise AirflowException("The parameter 'sqlalchemy_query' must be of type dict!")
+        if conn.extra_dejson.get("iam", False):
+            conn.login, conn.password, conn.port = self.get_iam_token(conn)
         return URL.create(
             drivername="postgresql+psycopg" if USE_PSYCOPG3 else "postgresql",
             username=self.__cast_nullable(conn.login, str),
@@ -441,8 +441,14 @@ class PostgresHook(DbApiHook):
         return PostgresHook._serialize_cell_ppg2(cell, conn)
 
     def get_iam_token(self, conn: Connection) -> tuple[str, str, int]:
+        """Get the IAM token from different identity providers."""
+        if conn.extra_dejson.get("azure_conn_id"):
+            return self.get_azure_iam_token(conn)
+        return self.get_aws_iam_token(conn)
+
+    def get_aws_iam_token(self, conn: Connection) -> tuple[str, str, int]:
         """
-        Get the IAM token.
+        Get the AWS IAM token.
 
         This uses AWSHook to retrieve a temporary password to connect to
         Postgres or Redshift. Port is required. If none is provided, the default
@@ -500,6 +506,34 @@ class PostgresHook(DbApiHook):
             token = rds_client.generate_db_auth_token(conn.host, port, conn.login)
         return cast("str", login), cast("str", token), port
 
+    def get_azure_iam_token(self, conn: Connection) -> tuple[str, str, int]:
+        """
+        Get the Azure IAM token.
+
+        This uses AzureBaseHook to retrieve an OAUTH token to connect to Postgres.
+        Scope for the OAuth token can be set in the config option ``azure_oauth_scope`` under the section ``[postgres]``.
+        """
+        if TYPE_CHECKING:
+            from airflow.providers.microsoft.azure.hooks.base_azure import AzureBaseHook
+
+        azure_conn_id = conn.extra_dejson.get("azure_conn_id", "azure_default")
+        try:
+            azure_conn = Connection.get(azure_conn_id)
+        except AttributeError:
+            azure_conn = Connection.get_connection_from_secrets(azure_conn_id)  # type: ignore[attr-defined]
+        try:
+            azure_base_hook: AzureBaseHook = azure_conn.get_hook()
+        except TypeError as e:
+            if "required positional argument: 'sdk_client'" in str(e):
+                raise AirflowOptionalProviderFeatureException(
+                    "Getting azure token is not supported by current version of 'AzureBaseHook'. "
+                    "Please upgrade apache-airflow-providers-microsoft-azure>=12.8.0"
+                ) from e
+            raise
+        scope = conf.get("postgres", "azure_oauth_scope", fallback=self.default_azure_oauth_scope)
+        token = azure_base_hook.get_token(scope).token
+        return cast("str", conn.login or azure_conn.login), token, conn.port or 5432
+
     def get_table_primary_key(self, table: str, schema: str | None = "public") -> list[str] | None:
         """
         Get the table's primary key.

{apache_airflow_providers_postgres-6.3.0.dist-info → apache_airflow_providers_postgres-6.4.1rc1.dist-info}/METADATA

@@ -1,12 +1,13 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-postgres
-Version: 6.3.0
+Version: 6.4.1rc1
 Summary: Provider package apache-airflow-providers-postgres for Apache Airflow
 Keywords: airflow-provider,postgres,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
 Maintainer-email: Apache Software Foundation <dev@airflow.apache.org>
 Requires-Python: >=3.10
 Description-Content-Type: text/x-rst
+License-Expression: Apache-2.0
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Environment :: Console
 Classifier: Environment :: Web Environment
@@ -14,31 +15,35 @@ Classifier: Intended Audience :: Developers
 Classifier: Intended Audience :: System Administrators
 Classifier: Framework :: Apache Airflow
 Classifier: Framework :: Apache Airflow :: Provider
-Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: System :: Monitoring
-Requires-Dist: apache-airflow>=2.10.0
-Requires-Dist: apache-airflow-providers-common-sql>=1.23.0
+License-File: LICENSE
+License-File: NOTICE
+Requires-Dist: apache-airflow>=2.10.0rc1
+Requires-Dist: apache-airflow-providers-common-compat>=1.8.0rc1
+Requires-Dist: apache-airflow-providers-common-sql>=1.23.0rc1
 Requires-Dist: psycopg2-binary>=2.9.9; python_version < '3.13'
 Requires-Dist: psycopg2-binary>=2.9.10; python_version >= '3.13'
 Requires-Dist: asyncpg>=0.30.0
-Requires-Dist: apache-airflow-providers-amazon>=2.6.0 ; extra == "amazon"
+Requires-Dist: apache-airflow-providers-amazon>=2.6.0rc1 ; extra == "amazon"
+Requires-Dist: apache-airflow-providers-microsoft-azure>=12.8.0rc1 ; extra == "microsoft-azure"
 Requires-Dist: apache-airflow-providers-openlineage ; extra == "openlineage"
 Requires-Dist: pandas>=2.1.2 ; extra == "pandas" and ( python_version <"3.13")
 Requires-Dist: pandas>=2.2.3 ; extra == "pandas" and ( python_version >="3.13")
 Requires-Dist: polars>=1.26.0 ; extra == "polars"
 Requires-Dist: psycopg[binary]>=3.2.9 ; extra == "psycopg"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.3.0/changelog.html
-Project-URL: Documentation, https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.3.0
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-postgres/6.4.1/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-postgres/6.4.1
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
 Project-URL: YouTube, https://www.youtube.com/channel/UCSXwxpWZQ7XZ1WL3wqevChA/
 Provides-Extra: amazon
+Provides-Extra: microsoft-azure
 Provides-Extra: openlineage
 Provides-Extra: pandas
 Provides-Extra: polars
@@ -69,7 +74,7 @@ Provides-Extra: psycopg
 
 Package ``apache-airflow-providers-postgres``
 
-Release: ``6.3.0``
+Release: ``6.4.1``
 
 
 `PostgreSQL <https://www.postgresql.org/>`__
@@ -82,7 +87,7 @@ This is a provider package for ``postgres`` provider. All classes for this provi
 are in ``airflow.providers.postgres`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.3.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.4.1/>`_.
 
 Installation
 ------------
@@ -96,15 +101,16 @@ The package supports the following python versions: 3.10,3.11,3.12,3.13
 Requirements
 ------------
 
-=======================================  ======================================
-PIP package                              Version required
-=======================================  ======================================
-``apache-airflow``                       ``>=2.10.0``
-``apache-airflow-providers-common-sql``  ``>=1.23.0``
-``psycopg2-binary``                      ``>=2.9.9; python_version < "3.13"``
-``psycopg2-binary``                      ``>=2.9.10; python_version >= "3.13"``
-``asyncpg``                              ``>=0.30.0``
-=======================================  ======================================
+==========================================  ======================================
+PIP package                                 Version required
+==========================================  ======================================
+``apache-airflow``                          ``>=2.10.0``
+``apache-airflow-providers-common-compat``  ``>=1.8.0``
+``apache-airflow-providers-common-sql``     ``>=1.23.0``
+``psycopg2-binary``                         ``>=2.9.9; python_version < "3.13"``
+``psycopg2-binary``                         ``>=2.9.10; python_version >= "3.13"``
+``asyncpg``                                 ``>=0.30.0``
+==========================================  ======================================
 
 Cross provider package dependencies
 -----------------------------------
@@ -119,14 +125,30 @@ You can install such cross-provider dependencies when installing from PyPI. For
     pip install apache-airflow-providers-postgres[amazon]
 
 
-==============================================================================================================  ===============
-Dependent package                                                                                               Extra
-==============================================================================================================  ===============
-`apache-airflow-providers-amazon <https://airflow.apache.org/docs/apache-airflow-providers-amazon>`_            ``amazon``
-`apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_    ``common.sql``
-`apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_  ``openlineage``
-==============================================================================================================  ===============
+======================================================================================================================  ===================
+Dependent package                                                                                                       Extra
+======================================================================================================================  ===================
+`apache-airflow-providers-amazon <https://airflow.apache.org/docs/apache-airflow-providers-amazon>`_                    ``amazon``
+`apache-airflow-providers-common-compat <https://airflow.apache.org/docs/apache-airflow-providers-common-compat>`_      ``common.compat``
+`apache-airflow-providers-common-sql <https://airflow.apache.org/docs/apache-airflow-providers-common-sql>`_            ``common.sql``
+`apache-airflow-providers-microsoft-azure <https://airflow.apache.org/docs/apache-airflow-providers-microsoft-azure>`_  ``microsoft.azure``
+`apache-airflow-providers-openlineage <https://airflow.apache.org/docs/apache-airflow-providers-openlineage>`_          ``openlineage``
+======================================================================================================================  ===================
+
+Optional dependencies
+----------------------
+
+===================  =====================================================================================
+Extra                Dependencies
+===================  =====================================================================================
+``amazon``           ``apache-airflow-providers-amazon>=2.6.0``
+``microsoft.azure``  ``apache-airflow-providers-microsoft-azure>=12.8.0``
+``openlineage``      ``apache-airflow-providers-openlineage``
+``pandas``           ``pandas>=2.1.2; python_version <"3.13"``, ``pandas>=2.2.3; python_version >="3.13"``
+``polars``           ``polars>=1.26.0``
+``psycopg``          ``psycopg[binary]>=3.2.9``
+===================  =====================================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.3.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-postgres/6.4.1/changelog.html>`_.
 

apache_airflow_providers_postgres-6.4.1rc1.dist-info/RECORD

@@ -0,0 +1,14 @@
+airflow/providers/postgres/__init__.py,sha256=CbYIEGsVSqkb8BiKTcIUh-DPdfj6HwFI42gQ9bvRcE8,1497
+airflow/providers/postgres/get_provider_info.py,sha256=tLd8Kw5q9bE0XewMKJ4zzwYQoPh_G-E3mmFxK8QzfB0,3269
+airflow/providers/postgres/assets/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/postgres/assets/postgres.py,sha256=XNhOJCbOA_soaaiS73JjULMqAM_7PBryhToe8FJREA0,1522
+airflow/providers/postgres/dialects/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
+airflow/providers/postgres/dialects/postgres.py,sha256=bsliwzooipmgluJsjGHYQHfZNWYxzzZ3WvaBLAy3u3A,5332
+airflow/providers/postgres/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
+airflow/providers/postgres/hooks/postgres.py,sha256=zjAJzkhMEae5cJxVl4KHcmdev2YRuTkpWytuH7UgwPY,28800
+apache_airflow_providers_postgres-6.4.1rc1.dist-info/entry_points.txt,sha256=dhtJi6PTWHd6BwKhmI4OtSPvQVI_p0yYWI0eba83HqY,104
+apache_airflow_providers_postgres-6.4.1rc1.dist-info/licenses/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
+apache_airflow_providers_postgres-6.4.1rc1.dist-info/licenses/NOTICE,sha256=E3-_E02gwwSEFzeeWPKmnIjOoos3hW28CLISV6sYrbQ,168
+apache_airflow_providers_postgres-6.4.1rc1.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+apache_airflow_providers_postgres-6.4.1rc1.dist-info/METADATA,sha256=_EtjUCijBJbk-tgkMGViVqCzaIpIGEwl4RQrYHx-c24,7738
+apache_airflow_providers_postgres-6.4.1rc1.dist-info/RECORD,,

apache_airflow_providers_postgres-6.4.1rc1.dist-info/licenses/NOTICE

@@ -0,0 +1,5 @@
+Apache Airflow
+Copyright 2016-2025 The Apache Software Foundation
+
+This product includes software developed at
+The Apache Software Foundation (http://www.apache.org/).

apache_airflow_providers_postgres-6.3.0.dist-info/RECORD

@@ -1,13 +0,0 @@
-airflow/providers/postgres/LICENSE,sha256=gXPVwptPlW1TJ4HSuG5OMPg-a3h43OGMkZRR1rpwfJA,10850
-airflow/providers/postgres/__init__.py,sha256=EBlBQkd4wsb32h6i--UVow1UxLY3oTKzALQLc1L5W7s,1497
-airflow/providers/postgres/get_provider_info.py,sha256=qEEYbClLY3-NH40dBk2u_nOIfvfEIHdXaWsIJ8J5Z68,2626
-airflow/providers/postgres/assets/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/postgres/assets/postgres.py,sha256=XNhOJCbOA_soaaiS73JjULMqAM_7PBryhToe8FJREA0,1522
-airflow/providers/postgres/dialects/__init__.py,sha256=9hdXHABrVpkbpjZgUft39kOFL2xSGeG4GEua0Hmelus,785
-airflow/providers/postgres/dialects/postgres.py,sha256=WhRMbSQv2u373R62e3buE5J1L3j-XGHhZxrWc8Gtc5s,5551
-airflow/providers/postgres/hooks/__init__.py,sha256=mlJxuZLkd5x-iq2SBwD3mvRQpt3YR7wjz_nceyF1IaI,787
-airflow/providers/postgres/hooks/postgres.py,sha256=YwA99YPz_ZpCvBn5yyw4ykhJ9hMGNvEsMmAf5nsfVE8,26908
-apache_airflow_providers_postgres-6.3.0.dist-info/entry_points.txt,sha256=dhtJi6PTWHd6BwKhmI4OtSPvQVI_p0yYWI0eba83HqY,104
-apache_airflow_providers_postgres-6.3.0.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-apache_airflow_providers_postgres-6.3.0.dist-info/METADATA,sha256=ndnpXUFB9vCAeT4wp6MVzJI4mxgr4KxmCO8NUA4qkIk,6283
-apache_airflow_providers_postgres-6.3.0.dist-info/RECORD,,