From 3940cbdbded17ed35dce504e8b2330ac94f31f40 Mon Sep 17 00:00:00 2001 From: Swapnil Date: Tue, 25 Jun 2024 15:15:03 +0530 Subject: [PATCH] refactor: encaps create/update logic in upsert --- agent.py | 51 +++++++++++++++++++++++++++++++++------------------ 1 file changed, 33 insertions(+), 18 deletions(-) diff --git a/agent.py b/agent.py index ceadc58..2987fd1 100644 --- a/agent.py +++ b/agent.py @@ -244,6 +244,36 @@ class ServerManager: return server['id'] return None + def upsert_server(self, post_data): + server_id = self.existing_server_id(post_data) + + if server_id: + logging.info('Server already exists with id: {}, Updating...'.format(server_id)) + response = self.update_server(post_data, server_id) + else: + logging.info('Server does not exist, Creating...') + response = self.create_server(post_data) + + # Extract the server_id from the response + server_id = response.get('server_id', None) + if server_id is None: + logging.error('Failed to get server_id from response: {}'.format(response)) + raise ValueError('Failed to get server_id from response') + + return server_id + + def upsert_note(self, note_data, server_id): + note_data['service_id'] = server_id + try: + note = self.get_note(server_id) + except urllib.error.HTTPError: + note = None + + if note: + return self.update_note(note_data, server_id) + else: + return self.create_note(note_data) + def validate_env_vars(): api_key = os.getenv('AGENT_API') host = os.getenv('HOST') @@ -263,26 +293,11 @@ def main(): server_manager = ServerManager(host, api_key) - # Check if the server already exists - server_id = server_manager.existing_server_id(post_data) - - # If the server exists, update it - if server_id: - logging.info('Server already exists with id: {}, Updating...'.format(server_id)) - logging.info(server_manager.update_server(post_data, server_id)) - else: - logging.info('Server does not exist, Creating...') - logging.info(server_manager.create_server(post_data)) + server_id = server_manager.upsert_server(post_data) + logging.info('Server id: {}'.format(server_id)) note_data = server_data.create_note_data(server_id) - try: - note = server_manager.get_note(server_id) - except urllib.error.HTTPError: - note = None - if note: - server_manager.update_note(note_data, server_id) - else: - server_manager.create_note(note_data) + server_manager.upsert_note(note_data, server_id) if __name__ == '__main__': main()