Repository: zalando/patroni
Pull request #2590 (2023-03-09 15:07:17)

Use IMDSv2 in aws callback example script
Use the new AWS session-oriented instance metadata access method

CyberDem0n: :+1:

Patch:
diff --git a/features/citus.feature b/features/citus.feature
index a7f53ef0..2cf295eb 100644
--- a/features/citus.feature
+++ b/features/citus.feature
@@ -10,20 +10,20 @@ Feature: citus
     And I start postgres3 in citus group 1
     Then replication works from postgres0 to postgres1 after 15 seconds
     Then replication works from postgres2 to postgres3 after 15 seconds
-    And postgres0 is registered in the coordinator postgres0 as the worker in group 0
-    And postgres2 is registered in the coordinator postgres0 as the worker in group 1
+    And postgres0 is registered in the postgres0 as the worker in group 0
+    And postgres2 is registered in the postgres0 as the worker in group 1
 
   Scenario: coordinator failover updates pg_dist_node
     Given I run patronictl.py failover batman --group 0 --candidate postgres1 --force
     Then postgres1 role is the primary after 10 seconds
     And replication works from postgres1 to postgres0 after 15 seconds
     And "sync" key in a group 0 in DCS has sync_standby=postgres0 after 15 seconds
-    And postgres1 is registered in the coordinator postgres1 as the worker in group 0
+    And postgres1 is registered in the postgres2 as the worker in group 0
     When I run patronictl.py failover batman --group 0 --candidate postgres0 --force
     Then postgres0 role is the primary after 10 seconds
     And replication works from postgres0 to postgres1 after 15 seconds
     And "sync" key in a group 0 in DCS has sync_standby=postgres1 after 15 seconds
-    And postgres0 is registered in the coordinator postgres0 as the worker in group 0
+    And postgres0 is registered in the postgres2 as the worker in group 0
 
   Scenario: worker switchover doesn't break client queries on the coordinator
     Given I create a distributed table on postgres0
@@ -33,14 +33,14 @@ Feature: citus
     And postgres3 role is the primary after 10 seconds
     And replication works from postgres3 to postgres2 after 15 seconds
     And "sync" key in a group 1 in DCS has sync_standby=postgres2 after 15 seconds
-    And postgres3 is registered in the coordinator postgres0 as the worker in group 1
+    And postgres3 is registered in the postgres0 as the worker in group 1
     And a thread is still alive
     When I run patronictl.py switchover batman --group 1 --force
     Then I receive a response returncode 0
     And postgres2 role is the primary after 10 seconds
     And replication works from postgres2 to postgres3 after 15 seconds
     And "sync" key in a group 1 in DCS has sync_standby=postgres3 after 15 seconds
-    And postgres2 is registered in the coordinator postgres0 as the worker in group 1
+    And postgres2 is registered in the postgres0 as the worker in group 1
     And a thread is still alive
     When I stop a thread
     Then a distributed table on postgres0 has expected rows
@@ -52,7 +52,7 @@ Feature: citus
     Then I receive a response returncode 0
     And postgres2 role is the primary after 10 seconds
     And replication works from postgres2 to postgres3 after 15 seconds
-    And postgres2 is registered in the coordinator postgres0 as the worker in group 1
+    And postgres2 is registered in the postgres0 as the worker in group 1
     And a thread is still alive
     When I stop a thread
     Then a distributed table on postgres0 has expected rows
@@ -65,7 +65,7 @@ Feature: citus
     Then I receive a response returncode 0
     And I receive a response output "+ttl: 20"
     When I sleep for 2 seconds
-    Then postgres4 is registered in the coordinator postgres0 as the worker in group 2
+    Then postgres4 is registered in the postgres2 as the worker in group 2
     When I shut down postgres4
     Then There is a transaction in progress on postgres0 changing pg_dist_node
     When I run patronictl.py restart batman postgres2 --group 1 --force
diff --git a/features/steps/citus.py b/features/steps/citus.py
index 5144f466..01f5ef2f 100644
--- a/features/steps/citus.py
+++ b/features/steps/citus.py
@@ -44,7 +44,7 @@ def start_citus(context, name, group):
     return context.pctl.start(name, custom_config={"citus": {"database": "postgres", "group": int(group)}})
 
 
-@step('{name1:w} is registered in the coordinator {name2:w} as the worker in group {group:d}')
+@step('{name1:w} is registered in the {name2:w} as the worker in group {group:d}')
 def check_registration(context, name1, name2, group):
     worker_port = int(context.pctl.query(name1, "SHOW port").fetchone()[0])
     r = context.pctl.query(name2, "SELECT nodeport FROM pg_catalog.pg_dist_node WHERE groupid = {0}".format(group))
diff --git a/patroni/postgresql/citus.py b/patroni/postgresql/citus.py
index 0866eb9d..69fa96a7 100644
--- a/patroni/postgresql/citus.py
+++ b/patroni/postgresql/citus.py
@@ -196,17 +196,12 @@ class CitusHandler(Thread):
             return i, task
 
     def update_node(self, task):
-        if task.group == CITUS_COORDINATOR_GROUP_ID:
-            return self.query("SELECT pg_catalog.citus_set_coordinator_host(%s, %s, 'primary', 'default')",
-                              task.host, task.port)
-
-        if task.nodeid is None and task.event != 'before_demote':
-            task.nodeid = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, 'primary', 'default')",
-                                     task.host, task.port, task.group).fetchone()[0]
-        elif task.nodeid is not None:
-            # XXX: statement_timeout?
+        if task.nodeid is not None:
             self.query('SELECT pg_catalog.citus_update_node(%s, %s, %s, true, %s)',
                        task.nodeid, task.host, task.port, task.cooldown)
+        elif task.event != 'before_demote':
+            task.nodeid = self.query("SELECT pg_catalog.citus_add_node(%s, %s, %s, 'primary', 'default')",
+                                     task.host, task.port, task.group).fetchone()[0]
 
     def process_task(self, task):
         """Updates a single row in `pg_dist_node` table, optionally in a transaction.
@@ -361,6 +356,11 @@ class CitusHandler(Thread):
                     cur.execute("INSERT INTO pg_catalog.pg_dist_authinfo VALUES"
                                 "(0, pg_catalog.current_user(), %s)",
                                 (self._postgresql.config.format_dsn(params),))
+
+                if self.is_coordinator():
+                    r = urlparse(self._postgresql.connection_string)
+                    cur.execute("SELECT pg_catalog.citus_set_coordinator_host(%s, %s, 'primary', 'default')",
+                                (r.hostname, r.port or 5432))
         finally:
             conn.close()
 
diff --git a/patroni/scripts/aws.py b/patroni/scripts/aws.py
index 31e7e398..7d18e828 100755
--- a/patroni/scripts/aws.py
+++ b/patroni/scripts/aws.py
@@ -6,9 +6,9 @@ import sys
 import boto3
 
 from ..utils import Retry, RetryFailedError
-from ..request import get as requests_get
 
 from botocore.exceptions import ClientError
+from botocore.utils import IMDSFetcher
 
 logger = logging.getLogger(__name__)
 
@@ -21,14 +21,16 @@ class AWSConnection(object):
         self._retry = Retry(deadline=300, max_delay=30, max_tries=-1, retry_exceptions=(ClientError,))
         try:
             # get the instance id
-            r = requests_get('http://169.254.169.254/latest/dynamic/instance-identity/document', timeout=2.1)
+            fetcher = IMDSFetcher(timeout=2.1)
+            token = fetcher._fetch_metadata_token()
+            r = fetcher._get_request("/latest/dynamic/instance-identity/document", None, token)
         except Exception:
             logger.error('cannot query AWS meta-data')
             return
 
-        if r.status < 400:
+        if r.status_code < 400:
             try:
-                content = json.loads(r.data.decode('utf-8'))
+                content = json.loads(r.text)
                 self.instance_id = content['instanceId']
                 self.region = content['region']
             except Exception:


Test patch:
diff --git a/tests/test_aws.py b/tests/test_aws.py
index 91345f33..c291dbd2 100644
--- a/tests/test_aws.py
+++ b/tests/test_aws.py
@@ -1,9 +1,9 @@
 import botocore
+import botocore.awsrequest
 import sys
 import unittest
-import urllib3
 
-from mock import Mock, patch
+from mock import Mock, PropertyMock, patch
 from collections import namedtuple
 from patroni.scripts.aws import AWSConnection, main as _main
 
@@ -28,31 +28,45 @@ class MockEc2Connection(object):
         return True
 
 
+class MockIMDSFetcher(object):
+
+    def __init__(self, timeout):
+        pass
+
+    @staticmethod
+    def _fetch_metadata_token():
+        return ''
+
+    @staticmethod
+    def _get_request(*args):
+        return botocore.awsrequest.AWSResponse(url='', status_code=200, headers={}, raw=None)
+
+
 @patch('boto3.resource', Mock(return_value=MockEc2Connection()))
+@patch('patroni.scripts.aws.IMDSFetcher', MockIMDSFetcher)
 class TestAWSConnection(unittest.TestCase):
 
-    @patch('patroni.scripts.aws.requests_get', Mock(return_value=urllib3.HTTPResponse(
-        status=200, body=b'{"instanceId": "012345", "region": "eu-west-1"}')))
-    def setUp(self):
-        self.conn = AWSConnection('test')
-
+    @patch.object(botocore.awsrequest.AWSResponse, 'text',
+                  PropertyMock(return_value='{"instanceId": "012345", "region": "eu-west-1"}'))
     def test_on_role_change(self):
-        self.assertTrue(self.conn.on_role_change('primary'))
+        conn = AWSConnection('test')
+        self.assertTrue(conn.on_role_change('primary'))
         with patch.object(MockVolumes, 'filter', Mock(return_value=[])):
-            self.conn._retry.max_tries = 1
-            self.assertFalse(self.conn.on_role_change('primary'))
+            conn._retry.max_tries = 1
+            self.assertFalse(conn.on_role_change('primary'))
 
-    @patch('patroni.scripts.aws.requests_get', Mock(side_effect=Exception('foo')))
+    @patch.object(MockIMDSFetcher, '_get_request', Mock(side_effect=Exception('foo')))
     def test_non_aws(self):
         conn = AWSConnection('test')
         self.assertFalse(conn.on_role_change("primary"))
 
-    @patch('patroni.scripts.aws.requests_get', Mock(return_value=urllib3.HTTPResponse(status=200, body=b'foo')))
+    @patch.object(botocore.awsrequest.AWSResponse, 'text', PropertyMock(return_value='boo'))
     def test_aws_bizare_response(self):
         conn = AWSConnection('test')
         self.assertFalse(conn.aws_available())
 
-    @patch('patroni.scripts.aws.requests_get', Mock(return_value=urllib3.HTTPResponse(status=503, body=b'Error')))
+    @patch.object(MockIMDSFetcher, '_get_request', Mock(return_value=botocore.awsrequest.AWSResponse(
+        url='', status_code=503, headers={}, raw=None)))
     @patch('sys.exit', Mock())
     def test_main(self):
         self.assertIsNone(_main())