Skip to content

Commit

Permalink
Sources GCP BigQuery Table ID should be maintained while updating app…
Browse files Browse the repository at this point in the history
…-settings (#2646)
  • Loading branch information
Douglas Curtis authored Feb 9, 2021
1 parent ed37904 commit 373ba63
Show file tree
Hide file tree
Showing 5 changed files with 64 additions and 5 deletions.
2 changes: 1 addition & 1 deletion koku/koku/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions koku/providers/gcp/provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
8 changes: 6 additions & 2 deletions koku/providers/test/gcp/tests_provider.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
3 changes: 3 additions & 0 deletions koku/sources/storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
53 changes: 53 additions & 0 deletions koku/sources/test/test_storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 373ba63

Please sign in to comment.