feat: Query Feature Settings (#36)

Query feature settings and store in CSV output/feature_settings.csv

datamgt.py

@@ -0,0 +1,51 @@
+import argparse
+import os
+import json
+import library.migrationlogger as nrlogger
+import library.clients.datamgtclient as datamgtclient
+import library.localstore as store
+import library.utils as utils
+
+logger = nrlogger.get_logger(os.path.basename(__file__))
+dmc = datamgtclient.DataManagementClient()
+
+
+def configure_parser():
+    parser = argparse.ArgumentParser(description='Migrate Dashboards')
+    parser.add_argument('--accounts', nargs=1, type=str, required=True, help='Path to file with account IDs')
+    parser.add_argument('--userApiKey', nargs=1, type=str, required=True, help='User API Key')
+    parser.add_argument('--region', type=str, nargs=1, required=False, help='sourceRegion us(default) or eu')
+    parser.add_argument('--featureSettings', dest='featureSettings', required=False, action='store_true',
+                        help='Query Feature Settings')
+    return parser
+
+
+def get_feature_settings(user_api_key, from_file, region):
+    acct_ids = store.load_names(from_file)
+    feature_settings = [['accountId', 'key','name', 'enabled']]
+    for acct_id in acct_ids:
+        result = dmc.get_feature_settings(user_api_key, int(acct_id), region)
+        logger.info(json.dumps(result))
+        for featureSetting in result['response']['data']['actor']['account']['dataManagement']['featureSettings']:
+            feature_settings.append([acct_id, featureSetting['key'], featureSetting['name'], featureSetting['enabled']])
+    logger.info(feature_settings)
+    store.save_feature_settings_csv(feature_settings)
+
+
+
+
+def main():
+    parser = configure_parser()
+    args = parser.parse_args()
+    user_api_key = utils.ensure_user_api_key(args)
+    if not user_api_key:
+        utils.error_and_exit('userApiKey', 'ENV_USER_API_KEY')
+    region = utils.ensure_region(args)
+    if args.featureSettings:
+        get_feature_settings(user_api_key, args.accounts[0], region)
+    else:
+        logger.info("pass --featureSettings to fetch featureSettings")
+
+
+if __name__ == '__main__':
+    main()

library/clients/datamgtclient.py

@@ -0,0 +1,35 @@
+import json
+import os
+import library.utils as utils
+import library.migrationlogger as nrlogger
+import library.clients.gql as nerdgraph
+
+logger = nrlogger.get_logger(os.path.basename(__file__))
+
+
+class DataManagementClient:
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def get_feature_settings(user_api_key, account_id, region):
+        payload = DataManagementClient._query_feature_settings_payload(account_id)
+        logger.debug(json.dumps(payload))
+        return nerdgraph.GraphQl.post(user_api_key, payload, region)
+
+    @staticmethod
+    def _query_feature_settings_payload(account_id):
+        query = '''query($accountId: Int!) { 
+                        actor {
+                            account(id: $accountId) { dataManagement {
+                                featureSettings {
+                                    enabled
+                                    key
+                                    name
+                                }
+                            }}
+                        }
+                    }'''
+        variables = {'accountId': account_id}
+        return {'query': query, 'variables': variables}

library/localstore.py

@@ -117,6 +117,17 @@ def save_host_data_csv(host_name: str, all_host_data: list):
         csv_writer.writerows(host_data + [""] for host_data in all_host_data)
 
 
+def save_feature_settings_csv(fs_data: list):
+    output_dir = Path("output")
+    fs_file_name = 'feature_settings.csv'
+    fs_data_file = output_dir / fs_file_name
+    create_file(fs_data_file)
+    with open(str(fs_data_file), 'w', newline='') as csvfile:
+        csv_writer = csv.writer(csvfile, delimiter=',',
+                                   quotechar='"', quoting=csv.QUOTE_ALL)
+        csv_writer.writerows(host_data + [""] for host_data in fs_data)
+
+
 def load_names(from_file):
     names = []
     with open(from_file) as input_names:

library/utils.py

@@ -25,8 +25,13 @@ def configure_loglevel(args):
     m_logger.set_log_level(log_level)
 
 
+# hack setting up both API key headers so should work for both REST and USER API Keys
 def setup_headers(api_key):
-    return {'Api-Key': api_key, 'Content-Type': 'application/json'}
+    return {'Api-Key': api_key, 'X-Api-Key': api_key, 'Content-Type': 'application/json'}
+
+
+def setup_infra_headers(api_key):
+    return {'X-Api-Key': api_key, 'Content-Type': 'application/json'}
 
 
 def get_next_url(rsp_headers):
@@ -53,8 +58,12 @@ def get_paginated_entities(api_key, fetch_url, entity_key, params={}, pageType=N
     curr_fetch_url = fetch_url
     all_entities = {'response_count': 0, entity_key: []}
     while another_page and error is False:
+        req_headers = setup_headers(api_key)
+        if 'infra-api' in curr_fetch_url:
+            req_headers = setup_infra_headers(api_key)
         resp = requests.get(curr_fetch_url, headers=setup_headers(api_key), params=params)
         if resp.status_code == 200:
+            logger.info(resp.text)
             resp_json = json.loads(resp.text)
             all_entities[entity_key].extend(resp_json[entity_key])
             all_entities['response_count'] = all_entities['response_count'] + len(resp_json[entity_key])
@@ -115,6 +124,16 @@ def ensure_source_api_key(args):
     return api_key
 
 
+def ensure_user_api_key(args):
+    if 'userApiKey' in args and args.userApiKey:
+        api_key = args.userApiKey[0]
+    elif 'user_api_key' in args and args.user_api_key:
+        api_key = args.user_api_key[0]
+    else:
+        api_key = os.environ.get('ENV_USER_API_KEY')
+    return api_key
+
+
 def ensure_insights_query_key(args):
     if 'insightsQueryKey' in args and args.insightsQueryKey:
         insights_query_key = args.insightsQueryKey[0]