fix datasource selection (#326)

* fix datasource selection

Author: iakov-aws

cid/_version.py

@@ -1,2 +1,2 @@
 
-__version__ = '0.2.2'
+__version__ = '0.2.3'

cid/cli.py

@@ -93,7 +93,7 @@ def deploy(ctx, **kwargs):
      --athena-workgroup TEXT               Athena workgroup
      --glue-data-catalog TEXT              Glue data catalog
      --cur-table-name TEXT                 CUR table name
-     --quicksight-datasource-arn TEXT      QuickSight Datasource ARN (if not found one with provided Athena workgroup)
+     --quicksight-datasource-id TEXT       QuickSight Datasource ID (not needed if only one DataSoruce exisits with a given Athena workgroup)
      --quicksight-user TEXT                QuickSight user
      --dataset-{dataset_name}-id TEXT      QuickSight dataset id for a specific dataset
      --view-{view_name}-{parameter} TEXT   a custom parameter for a view creation, can use variable: {account_id}

cid/common.py

@@ -948,64 +948,89 @@ def get_dataset_data_from_defintion(self, dataset_definition):
         return raw_template
 
 
-
     def create_or_update_dataset(self, dataset_definition: dict, dataset_id: str=None,recursive: bool=True, update: bool=False) -> bool:
         # Read dataset definition from template
         data = self.get_dataset_data_from_defintion(dataset_definition)
         template = Template(json.dumps(data))
         cur_required = dataset_definition.get('dependsOn', dict()).get('cur')
         athena_datasource = None
 
-
         if not len(self.qs.athena_datasources):
             logger.info('No Athena datasources found, attempting to create one')
             self.qs.AthenaWorkGroup = self.athena.WorkGroup
-            self.qs.create_data_source()
+            self.qs.create_data_source() # FIXME: we need to use name/id provided by user if any
+            # FIXME: we need to cleanup if datasource creation fails
 
-        if not len(self.qs.athena_datasources):
-            logger.info('No Athena datasources available, failing')
-            print('No Athena datasources detected and unable to create one. Please create at least one dataset manually if it fails.')
+        if not self.qs.athena_datasources:
+            logger.info('No valid DataSources available, failing')
+            print('No valid DataSources detected and unable to create one. Please create at least one DataSet manually in QuickSight and see why it fails.')
             # Not failing here to let views creation below
         else:
-            pre_compiled_dataset = json.loads(template.safe_substitute())
-            dataset_name = pre_compiled_dataset.get('Name')
-
-            # let's find the schema/database and workgroup name
-            schemas = []
-            datasources = []
-            if dataset_id:
-                schemas = self.qs.get_datasets(id=dataset_id)[0].schemas
-                datasources = self.qs.get_datasets(id=dataset_id)[0].datasources
-            else: # try to find dataset and get athena database
-                found_datasets = self.qs.get_datasets(name=dataset_name)
-                if found_datasets:
-                    schemas = list(set(sum([d.schemas for d in found_datasets], [])))
-                    datasources = list(set(sum([d.datasources for d in found_datasets], [])))
-
-            if len(schemas) == 1:
-                self.athena.DatabaseName = schemas[0]
-            # else user will be suggested to choose database
-            if len(datasources) == 1 and datasources[0] in self.qs.athena_datasources:
-                athena_datasource = self.qs.get_datasources(id=datasources[0])[0]
+            datasource_choices = {
+                f"{datasource.name} {id_} (workgroup={datasource.AthenaParameters.get('WorkGroup')})":id_
+                for id_, datasource in self.qs.athena_datasources.items()
+            }
+            if get_parameters().get('quicksight-datasource-id'):
+                # We have explicit choice of datasource
+                datasource_id = get_parameters().get('quicksight-datasource-id')
+                if datasource_id not in datasource_choices.values():
+                    logger.critical(
+                        f'quicksight-datasource-id={datasource_id} not found or not in a valid state. '
+                        f'Here is a list of available DataSources (Name ID WorkGroup): {datasource_choices.keys()}'
+                    )
+                    exit(1)
+                athena_datasource = self.qs.athena_datasources[datasource_id]
+
             else:
-                #try to find a datasource with defined workgroup
-                workgroup = self.athena.WorkGroup
-                datasources_with_workgroup = self.qs.get_datasources(athena_workgroup_name=workgroup)
-                if len(datasources_with_workgroup) == 1:
-                    athena_datasource = datasources_with_workgroup[0]
+                # Datasources are not obvious for customer so we will try to do our best guess
+                # - if there is just one? -> take that one
+                # - if datasource is references in existing dataset? -> take that one
+                # - if athena workgroup defined -> Try to find a dataset with this workgroup
+                # - and if still nothing -> ask an expicit choice from the user
+                pre_compiled_dataset = json.loads(template.safe_substitute())
+                dataset_name = pre_compiled_dataset.get('Name')
+
+                # let's find the schema/database and workgroup name
+                schemas = []
+                datasources = []
+                if dataset_id:
+                    schemas = self.qs.get_datasets(id=dataset_id)[0].schemas
+                    datasources = self.qs.get_datasets(id=dataset_id)[0].datasources
+                else: # try to find dataset and get athena database
+                    found_datasets = self.qs.get_datasets(name=dataset_name)
+                    if found_datasets:
+                        schemas = list(set(sum([d.schemas for d in found_datasets], [])))
+                        datasources = list(set(sum([d.datasources for d in found_datasets], [])))
+
+                if len(schemas) == 1:
+                    self.athena.DatabaseName = schemas[0]
+                # else user will be suggested to choose database anyway
+
+                if len(datasources) == 1 and datasources[0] in self.qs.athena_datasources:
+                    athena_datasource = self.qs.get_datasources(id=datasources[0])[0]
                 else:
-                    #cannot find the right athena_datasource
-                    athena_datasource = get_parameter(
-                        param_name='quicksight-datasource-arn',
-                        message=f"Please choose DataSource ARN",
-                        choices={f"{arn} (workgroup={datasource.AthenaParameters.get('WorkGroup')})":datasource for arn, datasource in self.qs.athena_datasources.items()},
-                    )
-                    logger.info(f'Found {len(datasources)} Athena datasources, not using {athena_datasource.id}')
-            if isinstance(athena_datasource, Datasource):
-                self.athena.WorkGroup = athena_datasource.AthenaParameters.get('WorkGroup')
+                    #try to find a datasource with defined workgroup
+                    workgroup = self.athena.WorkGroup
+                    datasources_with_workgroup = self.qs.get_datasources(athena_workgroup_name=workgroup)
+                    if len(datasources_with_workgroup) == 1:
+                        athena_datasource = datasources_with_workgroup[0]
+                    else:
+                        #cannot find the right athena_datasource
+                        logger.info('Multiple DataSources found.')
+                        datasource_id = get_parameter(
+                            param_name='quicksight-datasource-id',
+                            message=f"Please choose DataSource (Choose the first one if not sure).",
+                            choices=datasource_choices,
+                        )
+                        athena_datasource = self.qs.athena_datasources[datasource_id]
+                        logger.info(f'Found {len(datasources)} Athena datasources, not using {athena_datasource.id}')
+                if isinstance(athena_datasource, Datasource):
+                    self.athena.WorkGroup = athena_datasource.AthenaParameters.get('WorkGroup')
+                else:
+                    logger.debug('Athena_datasource is not defined. Will only create views')
 
         # Check for required views
-        _views = dataset_definition.get('dependsOn').get('views')
+        _views = dataset_definition.get('dependsOn', {}).get('views', [])
         required_views = [(self.cur.tableName if cur_required and name =='${cur_table_name}' else name) for name in _views]
 
         self.athena.discover_views(required_views)