Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Prod 2435: bq connection test #5138

Merged
merged 12 commits into from
Jul 31, 2024
Merged

Prod 2435: bq connection test #5138

Changes from 2 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
d02d73c
janky connection test POC for bigquery
thingscouldbeworse Jul 30, 2024
08e3263
add back in
thingscouldbeworse Jul 30, 2024
4e7a8b2
move to overriding the BigQueryConnector test connection method
thingscouldbeworse Jul 30, 2024
292da45
revert unneeded
thingscouldbeworse Jul 30, 2024
338f772
revert unneeded
thingscouldbeworse Jul 30, 2024
f694506
projects, not datasets
thingscouldbeworse Jul 30, 2024
17f36df
Merge branch 'main' into PROD-2435_bq-connection-test
thingscouldbeworse Jul 30, 2024
591456b
isort
thingscouldbeworse Jul 30, 2024
16d65db
docstring and log update
thingscouldbeworse Jul 31, 2024
71f7613
mypy & pylint fixes
thingscouldbeworse Jul 31, 2024
7256029
run connection test
thingscouldbeworse Jul 31, 2024
2bd07b6
Merge branch 'main' into PROD-2435_bq-connection-test
thingscouldbeworse Jul 31, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 21 additions & 0 deletions src/fides/api/util/connection_util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
from fides.api.schemas.connection_configuration.connection_secrets import (
TestStatusMessage,
)
from fides.api.schemas.connection_configuration.connection_secrets_bigquery import BigQuerySchema
from fides.api.schemas.connection_configuration.connection_secrets_saas import (
validate_saas_secrets_external_references,
)
Expand Down Expand Up @@ -323,6 +324,26 @@
) -> TestStatusMessage:
"""Connect, verify with a trivial query or API request, and report the status."""

# test if this is a Bigquery connection
if connection_config.connection_type == ConnectionType.bigquery:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

hmm. this will probably work, but i think there may be a cleaner way to inject the "custom" bigquery logic.

what i'd lean toward is implementing an overridden test_connection method in the BigQueryConnector class that's a subclass of the SQLConnector (this is the bigquery "DSR connector"): https://github.com/ethyca/fides/blob/main/src/fides/api/service/connectors/sql_connector.py#L476

that allows us to keep all the generic wrapping logic in tact, and it puts in the "custom" bigquery logic into a class specifically scoped toward bigquery. this will pave the road nicely for when we move the BigQueryConnector to leverage the python bigquery client, as is currently used in the d&d workflow.

# run a different connection test where we pull all projects
from google.cloud.bigquery import Client as BigQueryClient
connector = get_connector(connection_config)
secrets = connector.configuration.secrets or {}
keyfile_creds = secrets.get("keyfile_creds", {})
client = BigQueryClient.from_service_account_info(keyfile_creds)

Check warning on line 334 in src/fides/api/util/connection_util.py

View check run for this annotation

Codecov / codecov/patch

src/fides/api/util/connection_util.py#L330-L334 

Added lines #L330 - L334 were not covered by tests
dbs = [project for project in client.list_projects()]
if dbs:
status = ConnectionTestStatus.succeeded

Check warning on line 337 in src/fides/api/util/connection_util.py

View check run for this annotation

Codecov / codecov/patch

src/fides/api/util/connection_util.py#L337 

Added line #L337 was not covered by tests
else:
status = ConnectionTestStatus.failed
msg = "No projects found in BigQuery, connection test failed."

Check warning on line 340 in src/fides/api/util/connection_util.py

View check run for this annotation

Codecov / codecov/patch

src/fides/api/util/connection_util.py#L339-L340 

Added lines #L339 - L340 were not covered by tests

return TestStatusMessage(

Check warning on line 342 in src/fides/api/util/connection_util.py

View check run for this annotation

Codecov / codecov/patch

src/fides/api/util/connection_util.py#L342 

Added line #L342 was not covered by tests
msg=msg,
test_status=status,
)

connector = get_connector(connection_config)

try:
Expand Down
Loading