[AIRFLOW-4417] Add AWS IAM authenication for PostgresHook

airflow/hooks/postgres_hook.py

@@ -33,6 +33,16 @@ class PostgresHook(DbApiHook):
 
     Note: For Redshift, use keepalives_idle in the extra connection parameters
     and set it to less than 300 seconds.
+
+    Note: For AWS IAM authentication, use iam in the extra connection parameters
+    and set it to true. Leave the password field empty. This will use the the
+    "aws_default" connection to get the temporary token unless you override
+    in extras.
+    extras example: ``{"iam":true, "aws_conn_id":"my_aws_conn"}``
+    For Redshift, also use redshift in the extra connection parameters and
+    set it to true. The cluster-identifier is extracted from the beginning of
+    the host field, so is optional. It can however be overridden in the extra field.
+    extras example: ``{"iam":true, "redshift":true, "cluster-identifier": "my_cluster_id"}``
     """
     conn_name_attr = 'postgres_conn_id'
     default_conn_name = 'postgres_default'
@@ -44,6 +54,11 @@ def __init__(self, *args, **kwargs):
 
     def get_conn(self):
         conn = self.get_connection(self.postgres_conn_id)
+
+        # check for authentication via AWS IAM
+        if conn.extra_dejson.get('iam', False):
+            conn.login, conn.password, conn.port = self.get_iam_token(conn)
+
         conn_args = dict(
             host=conn.host,
             user=conn.login,
@@ -111,3 +126,36 @@ def _serialize_cell(cell, conn):
         :rtype: object
         """
         return cell
+
+    def get_iam_token(self, conn):
+        """
+        Uses AWSHook to retrieve a temporary password to connect to Postgres
+        or Redshift. Port is required. If none is provided, default is used for
+        each service
+        """
+        from airflow.contrib.hooks.aws_hook import AwsHook
+
+        redshift = conn.extra_dejson.get('redshift', False)
+        aws_conn_id = conn.extra_dejson.get('aws_conn_id', 'aws_default')
+        aws_hook = AwsHook(aws_conn_id)
+        login = conn.login
+        if conn.port is None:
+            port = 5439 if redshift else 5432
+        else:
+            port = conn.port
+        if redshift:
+            # Pull the custer-identifier from the beginning of the Redshift URL
+            # ex. my-cluster.ccdre4hpd39h.us-east-1.redshift.amazonaws.com returns my-cluster
+            cluster_identifier = conn.extra_dejson.get('cluster-identifier', conn.host.split('.')[0])
+            client = aws_hook.get_client_type('redshift')
+            cluster_creds = client.get_cluster_credentials(
+                DbUser=conn.login,
+                DbName=self.schema or conn.schema,
+                ClusterIdentifier=cluster_identifier,
+                AutoCreate=False)
+            token = cluster_creds['DbPassword']
+            login = cluster_creds['DbUser']
+        else:
+            client = aws_hook.get_client_type('rds')
+            token = client.generate_db_auth_token(conn.host, port, conn.login)
+        return login, token, port

tests/hooks/test_postgres_hook.py

@@ -24,6 +24,51 @@
 from tempfile import NamedTemporaryFile
 
 from airflow.hooks.postgres_hook import PostgresHook
+from airflow.models import Connection
+
+
+class TestPostgresHookConn(unittest.TestCase):
+
+    def setUp(self):
+        super().setUp()
+
+        self.connection = Connection(
+            login='login',
+            password='password',
+            host='host',
+            schema='schema'
+        )
+
+        self.db_hook = PostgresHook()
+        self.db_hook.get_connection = mock.Mock()
+        self.db_hook.get_connection.return_value = self.connection
+
+    @mock.patch('airflow.hooks.postgres_hook.psycopg2.connect')
+    def test_get_conn(self, mock_connect):
+        self.db_hook.get_conn()
+        mock_connect.assert_called_once_with(user='login', password='password', host='host',
+                                             dbname='schema', port=None)
+
+    @mock.patch('airflow.hooks.postgres_hook.psycopg2.connect')
+    @mock.patch('airflow.contrib.hooks.aws_hook.AwsHook.get_client_type')
+    def test_get_conn_rds_iam_postgres(self, mock_client, mock_connect):
+        self.connection.extra = '{"iam":true}'
+        mock_client.return_value.generate_db_auth_token.return_value = 'aws_token'
+        self.db_hook.get_conn()
+        mock_connect.assert_called_once_with(user='login', password='aws_token', host='host',
+                                             dbname='schema', port=5432)
+
+    @mock.patch('airflow.hooks.postgres_hook.psycopg2.connect')
+    @mock.patch('airflow.contrib.hooks.aws_hook.AwsHook.get_client_type')
+    def test_get_conn_rds_iam_redshift(self, mock_client, mock_connect):
+        self.connection.extra = '{"iam":true, "redshift":true}'
+        self.connection.host = 'cluster-identifier.ccdfre4hpd39h.us-east-1.redshift.amazonaws.com'
+        login = 'IAM:{login}'.format(login=self.connection.login)
+        mock_client.return_value.get_cluster_credentials.return_value = {'DbPassword': 'aws_token',
+                                                                         'DbUser': login}
+        self.db_hook.get_conn()
+        mock_connect.assert_called_once_with(user=login, password='aws_token', host=self.connection.host,
+                                             dbname='schema', port=5439)
 
 
 class TestPostgresHook(unittest.TestCase):