Merge pull request #121 from aws-samples/bugfix/cur-detection

Improve CUR detection

cid/helpers/cur.py

@@ -1,10 +1,25 @@
 import json
+import questionary
 from cid.helpers import Athena
 
 class CUR:
-    defaults = {
-        'TableName': 'customer_all'
-    }
+    requiredColumns = [
+        'identity_line_item_id',
+        'identity_time_interval',
+        'bill_invoice_id',
+        'bill_billing_entity',
+        'bill_bill_type',
+        'bill_payer_account_id',
+        'bill_billing_period_start_date',
+        'bill_billing_period_end_date',
+        'line_item_usage_account_id',
+        'line_item_line_item_type',
+        'line_item_usage_start_date',
+        'line_item_usage_end_date',
+        'line_item_product_code',
+        'line_item_usage_type',
+        'line_item_operation',
+    ]
     _tableName = None
     _metadata = None
     _clients = dict()
@@ -46,20 +61,20 @@ def configured(self) -> bool:
         return self._configured
 
     @property
-    def tableName(self):
+    def tableName(self) -> str:
         if self.metadata is None:
             print('Error: CUR not detected')
             exit(1)
         return self.metadata.get('Name')
 
     @property
-    def hasResourceIDs(self):
+    def hasResourceIDs(self) -> bool:
         if self._configured and self._hasResourceIDs is None:
             self._hasResourceIDs = 'line_item_resource_id' in self.fields
         return self._hasResourceIDs
 
     @property
-    def hasReservations(self):
+    def hasReservations(self) -> bool:
         if self._configured and self._hasReservations is None:
             kwargs = {
                 'cur_table_name': self._tableName
@@ -68,7 +83,7 @@ def hasReservations(self):
         return self._hasReservations
 
     @property
-    def hasSavingsPlans(self):
+    def hasSavingsPlans(self) -> bool:
         if self._configured and self._hasSavingsPlans is None:
             kwargs = {
                 'cur_table_name': self._tableName
@@ -80,24 +95,24 @@ def hasSavingsPlans(self):
     def metadata(self) -> dict:
         if not self._metadata:
             try:
-                # Look for default CUR table
-                self._metadata = self.athena.get_table_metadata(self.defaults.get('TableName'))
-                self._tableName = self.defaults.get('TableName')
-            except self.athena.client.exceptions.MetadataException:
-                # Look based on CUR Athena database name
-                try:
-                    tables = self.athena.list_table_metadata()
-                    tables = [v for v in tables if v.get('TableType') == 'EXTERNAL_TABLE']
-                    for table in tables:
-                        if table.get('Name') in self.athena.DatabaseName:
-                            self._metadata = self.athena.get_table_metadata(table.get('Name'))
-                            self._tableName = table.get('Name')
-                            break
-                    # self._metadata = self.athena.get_table_metadata(self.athena._DatabaseName.rpartition('_')[2])
-                    # self._tableName = self.athena._DatabaseName.rpartition('_')[2]
-                except self.athena.client.exceptions.MetadataException:
-                    # TODO: ask user
-                    print('Error: CUR metadata not found')
+                # Look other tables
+                tables = self.athena.list_table_metadata()
+                # Filter tables with type = 'EXTERNAL_TABLE'
+                tables = [v for v in tables if v.get('TableType') == 'EXTERNAL_TABLE']
+                # Filter tables having CUR structure
+                for table in tables.copy():
+                    columns = [c.get('Name') for c in table.get('Columns')]                    
+                    if not all([c in columns for c in self.requiredColumns]):
+                        tables.remove(table)
+                if len(tables) == 1:
+                    self._metadata = tables[0]
+                    self._tableName = self._metadata.get('Name')
+                elif len(tables) > 1:
+                    self._tableName=questionary.select(
+                        "Multiple CUR tables found, please select one",
+                        choices=[v.get('Name') for v in tables]
+                    ).ask()
+                    self._metadata = self.athena.get_table_metadata(self._tableName)
             except Exception as e:
                 # For other errors dump the message
                 print(json.dumps(e, indent=4, sort_keys=True, default=str))

setup.cfg

@@ -4,7 +4,7 @@ universal = 1
 [metadata]
 name = cid-cmd
 # version = attr: VERSION
-version = 0.1.3
+version = 0.1.4
 keywords = aws, cmd, cli, cost intelligence dashboards
 description = Cloud Intelligence Dashboards deployment helper tool
 long_description = file: README.md