Sources GCP BigQuery Table ID should be maintained while updating app-settings

koku/koku/settings.py

@@ -418,7 +418,7 @@
 RABBITMQ_HOST = os.getenv("RABBITMQ_HOST", "localhost")
 RABBITMQ_PORT = os.getenv("RABBITMQ_PORT", "5672")
 
-CELERY_BROKER_URL = f"redis://{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB}"
+CELERY_BROKER_URL = f"amqp://{RABBITMQ_HOST}:{RABBITMQ_PORT}"
 CELERY_RESULTS_URL = f"redis://{REDIS_HOST}:{REDIS_PORT}/{REDIS_DB}"
 CELERY_IMPORTS = ("masu.processor.tasks", "masu.celery.tasks", "koku.metrics")
 CELERY_BROKER_POOL_LIMIT = None

koku/providers/gcp/provider.py

@@ -107,8 +107,7 @@ def cost_usage_source_is_reachable(self, credentials, data_source):
             LOG.info(error_obj(key, reason))
             raise serializers.ValidationError(error_obj(key, reason))
 
-        if not data_source.get("table_id"):
-            self._detect_billing_export_table(data_source, credentials)
+        self._detect_billing_export_table(data_source, credentials)
 
         return True
 

koku/providers/test/gcp/tests_provider.py

@@ -50,16 +50,20 @@ def test_name(self):
         provider = GCPProvider()
         self.assertEqual(provider.name(), Provider.PROVIDER_GCP)
 
+    @patch("providers.gcp.provider.bigquery")
     @patch("providers.gcp.provider.discovery")
     @patch("providers.gcp.provider.google.auth.default")
-    def test_cost_usage_source_is_reachable_valid(self, mock_auth, mock_discovery):
+    def test_cost_usage_source_is_reachable_valid(self, mock_auth, mock_discovery, mock_bigquery):
         """Test that cost_usage_source_is_reachable succeeds."""
+        mock_bigquery.Client.return_value = MockBigQueryClient(
+            "test_project", "test_dataset", ["gcp_billing_export_1234"]
+        )
         gcp_creds = MagicMock()
         mock_auth.return_value = (gcp_creds, MagicMock())
         mock_discovery.build.return_value.projects.return_value.testIamPermissions.return_value.execute.return_value.get.return_value = (  # noqa: E501
             REQUIRED_IAM_PERMISSIONS
         )
-        billing_source_param = {"dataset": FAKE.word(), "table_id": FAKE.word()}
+        billing_source_param = {"dataset": FAKE.word()}
         credentials_param = {"project_id": FAKE.word()}
         provider = GCPProvider()
         self.assertTrue(provider.cost_usage_source_is_reachable(credentials_param, billing_source_param))

koku/sources/storage.py

@@ -548,6 +548,9 @@ def _update_billing_source(instance, billing_source):
         if data_source.get("resource_group") or data_source.get("storage_account"):
             billing_copy.update(billing_source.get("data_source"))
             billing_source["data_source"] = billing_copy
+        if billing_copy.get("table_id"):
+            billing_copy.update(billing_source.get("data_source"))
+            billing_source["data_source"]["table_id"] = billing_copy.get("table_id")
     _validate_billing_source(instance.source_type, billing_source)
     # This if statement can also be removed if UI is updated to send "data_source" field
     if instance.source_type in (Provider.PROVIDER_AWS, Provider.PROVIDER_AWS_LOCAL) and not billing_source.get(

koku/sources/test/test_storage.py

@@ -848,6 +848,59 @@ def test_update_application_settings_billing_and_auth_starting_with_billing(self
         self.assertEqual(db_obj.billing_source.get("data_source").get("resource_group"), resource_group)
         self.assertEqual(db_obj.billing_source.get("data_source").get("storage_account"), storage_account)
 
+    def test_update_application_settings_gcp_table_id_present(self):
+        """Test to update application settings for GCP when table ID is known."""
+        test_source_id = 3
+        old_dataset = "olddataset"
+        new_dataset = "newdataset"
+        project_id = "myproject"
+        table_id = "bigquerytableid"
+        settings = {"billing_source": {"data_source": {"dataset": new_dataset}}}
+
+        gcp_obj = Sources(
+            source_id=test_source_id,
+            auth_header=self.test_header,
+            offset=3,
+            authentication={"credentials": {"project_id": project_id}},
+            billing_source={"data_source": {"dataset": old_dataset, "table_id": table_id}},
+            source_type=Provider.PROVIDER_GCP,
+            name="Test GCP Source",
+        )
+        gcp_obj.save()
+
+        storage.update_application_settings(test_source_id, settings)
+        db_obj = Sources.objects.get(source_id=test_source_id)
+
+        self.assertEqual(db_obj.authentication.get("credentials").get("project_id"), project_id)
+        self.assertEqual(db_obj.billing_source.get("data_source").get("dataset"), new_dataset)
+        self.assertEqual(db_obj.billing_source.get("data_source").get("table_id"), table_id)
+
+    def test_update_application_settings_gcp_table_id_not_present(self):
+        """Test to update application settings for GCP when table ID is not known."""
+        test_source_id = 3
+        old_dataset = "olddataset"
+        new_dataset = "newdataset"
+        project_id = "myproject"
+        settings = {"billing_source": {"data_source": {"dataset": new_dataset}}}
+
+        gcp_obj = Sources(
+            source_id=test_source_id,
+            auth_header=self.test_header,
+            offset=3,
+            authentication={"credentials": {"project_id": project_id}},
+            billing_source={"data_source": {"dataset": old_dataset}},
+            source_type=Provider.PROVIDER_GCP,
+            name="Test GCP Source",
+        )
+        gcp_obj.save()
+
+        storage.update_application_settings(test_source_id, settings)
+        db_obj = Sources.objects.get(source_id=test_source_id)
+
+        self.assertEqual(db_obj.authentication.get("credentials").get("project_id"), project_id)
+        self.assertEqual(db_obj.billing_source.get("data_source").get("dataset"), new_dataset)
+        self.assertIsNone(db_obj.billing_source.get("data_source").get("table_id"))
+
     def test_save_status(self):
         """Test to verify source status is saved."""
         test_source_id = 3