feat(oauth): adding necessary changes to support bigquery oauth #30674

fisjac · 2024-10-22T16:51:34Z

SUMMARY

Adding needed changes to support future implementation of OAuth2 functionality for BigQuery.

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

N/A

TESTING INSTRUCTIONS

ADDITIONAL INFORMATION

Has associated issue:
Required feature flags:
Changes UI
Includes DB Migration (follow approval process in SIP-59)
- Migration is atomic, supports rollback & is backwards-compatible
- Confirm DB migration upgrade and downgrade tested
- Runtime estimates and downtime expectations provided
Introduces new feature or API
Removes existing feature or API

fisjac · 2024-10-22T16:55:36Z

superset/commands/database/test_connection.py

@@ -225,6 +225,10 @@ def ping(engine: Engine) -> bool:
            # bubble up the exception to return proper status code
            raise
        except Exception as ex:
+            if database.is_oauth2_enabled() and database.db_engine_spec.needs_oauth2(


BigQuery raises an error when running create_engine() necessitating another needs_oauth2(ex) check for the proper oauth2 exceptions to trigger the Oauth dance. Most DB's allow for an engine to be created without valid creds, and instead raises the exception on engine.connect()

fisjac · 2024-10-22T16:56:34Z

superset/db_engine_specs/base.py

@@ -1691,10 +1691,13 @@ def select_star(  # pylint: disable=too-many-arguments
        return sql

    @classmethod
-    def estimate_statement_cost(cls, statement: str, cursor: Any) -> dict[str, Any]:
+    def estimate_statement_cost(


Database is required to create the BigQuery Client when using OAuth2

fisjac · 2024-10-22T16:57:26Z

superset/db_engine_specs/postgres.py

@@ -351,7 +351,16 @@ def get_allow_cost_estimate(cls, extra: dict[str, Any]) -> bool:
        return True

    @classmethod
-    def estimate_statement_cost(cls, statement: str, cursor: Any) -> dict[str, Any]:
+    def estimate_statement_cost(


aligning other specs to have database, and adding missing docstring

fisjac · 2024-10-22T16:57:53Z

superset/db_engine_specs/presto.py

@@ -365,9 +365,12 @@ def get_schema_from_engine_params(
        return parse.unquote(database.split("/")[1])

    @classmethod
-    def estimate_statement_cost(cls, statement: str, cursor: Any) -> dict[str, Any]:
+    def estimate_statement_cost(


adding database param

fisjac · 2024-10-22T16:59:06Z

superset/models/core.py

-                effective_username,
-                access_token,
-            )
+            # Checking if the function signature can accept database as a param


Again, here BigQuery needs the Database object to build the client. Checks if the function signature has database as a param, then passes it.

fisjac · 2024-10-22T16:59:48Z

superset/utils/oauth2.py

@@ -192,3 +192,4 @@ class OAuth2ClientConfigSchema(Schema):
    )
    authorization_request_uri = fields.String(required=True)
    token_request_uri = fields.String(required=True)
+    project_id = fields.String(required=False)


BigQuery takes project_id as a param in its OAuth2 parameters

I think ideally we'd want to store the project id in the URI, to make it compatible with non-oauth2 use cases. And then later when you need you can grab it from database.sqlalchemy_uri.database.

codecov · 2024-10-22T17:09:32Z

Codecov Report

Attention: Patch coverage is 85.18519% with 4 lines in your changes missing coverage. Please review.

Project coverage is 70.82%. Comparing base (76d897e) to head (25cdd3a).
Report is 882 commits behind head on master.

Files with missing lines	Patch %	Lines
superset/db_engine_specs/bigquery.py	50.00%	2 Missing ⚠️
...Modal/DatabaseConnectionForm/OAuth2ClientField.tsx	50.00%	0 Missing and 1 partial ⚠️
superset/commands/database/test_connection.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##           master   #30674       +/-   ##
===========================================
+ Coverage   60.48%   70.82%   +10.33%     
===========================================
  Files        1931     1987       +56     
  Lines       76236    80190     +3954     
  Branches     8568     9170      +602     
===========================================
+ Hits        46114    56792    +10678     
+ Misses      28017    21166     -6851     
- Partials     2105     2232      +127

Flag	Coverage Δ
hive	`48.91% <56.00%> (-0.25%)`	⬇️
javascript	`58.61% <50.00%> (+0.89%)`	⬆️
mysql	`76.74% <80.00%> (?)`
postgres	`76.87% <80.00%> (?)`
presto	`53.38% <56.00%> (-0.42%)`	⬇️
python	`83.91% <88.00%> (+20.42%)`	⬆️
sqlite	`76.32% <80.00%> (?)`
unit	`60.87% <84.00%> (+3.24%)`	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

betodealmeida

Looks good, but wondering if we need to modify _get_client here to take a database object.

superset/models/core.py

geido · 2024-10-25T08:18:03Z

...t-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/OAuth2ClientField.tsx

@@ -106,6 +108,15 @@ export const OAuth2ClientField = ({ changeMethods, db }: FieldPropTypes) => {
            onChange={handleChange('scope')}
          />
        </FormItem>
+        {db.engine === Engines.BigQuery && (
+          <FormItem label="Project ID">


Let's localize the label

geido · 2024-10-25T08:21:13Z

superset/models/core.py

-            )
+            # PR #30674 changed the signature of the method to include database.
+            # This ensures that the change is backwards compatible
+            sig = signature(self.db_engine_spec.update_impersonation_config)


Do we have sufficient test coverage for this case?

betodealmeida

Looks good, but I left a few comments on the handling of project_id. Happy to hop on a meeting to discuss it in more detail.

betodealmeida · 2024-10-25T18:13:33Z

...t-frontend/src/features/databases/DatabaseModal/DatabaseConnectionForm/OAuth2ClientField.tsx


 interface OAuth2ClientInfo {
  id: string;
  secret: string;
  authorization_request_uri: string;
  token_request_uri: string;
  scope: string;
+  project_id?: string;


Since this is BigQuery specific, it's better to not add it here, otherwise the interface loses its purpose. Ideally we want to keep only the common information that every oauth2 client has.

Can we have this in a separate attribute instead?

betodealmeida · 2024-10-25T18:17:48Z

superset/utils/oauth2.py

@@ -192,3 +192,4 @@ class OAuth2ClientConfigSchema(Schema):
    )
    authorization_request_uri = fields.String(required=True)
    token_request_uri = fields.String(required=True)
+    project_id = fields.String(required=False)


I think ideally we'd want to store the project id in the URI, to make it compatible with non-oauth2 use cases. And then later when you need you can grab it from database.sqlalchemy_uri.database.

pull-request-size bot added the size/M label Oct 22, 2024

fisjac commented Oct 22, 2024

View reviewed changes

fisjac requested a review from betodealmeida October 22, 2024 17:05

michael-s-molina added review:draft review:checkpoint Last PR reviewed during the daily review standup labels Oct 22, 2024

fisjac marked this pull request as ready for review October 22, 2024 20:09

dosubot bot added authentication:sso Single Sign On data:connect:googlebigquery Related to BigQuery labels Oct 22, 2024

betodealmeida reviewed Oct 22, 2024

View reviewed changes

superset/models/core.py Outdated Show resolved Hide resolved

superset/models/core.py Outdated Show resolved Hide resolved

michael-s-molina removed review:checkpoint Last PR reviewed during the daily review standup review:draft labels Oct 23, 2024

pull-request-size bot added size/L and removed size/M labels Oct 24, 2024

fisjac force-pushed the bigquery-oauth branch from 11f631e to 96bcb1c Compare October 24, 2024 01:07

bigquery oauth

6e42328

fisjac force-pushed the bigquery-oauth branch from 96bcb1c to 6e42328 Compare October 24, 2024 01:37

fixing double encoding bug for OAuth JWTs

25cdd3a

fisjac force-pushed the bigquery-oauth branch from 5872351 to 25cdd3a Compare October 24, 2024 22:25

geido reviewed Oct 25, 2024

View reviewed changes

betodealmeida requested changes Oct 25, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat(oauth): adding necessary changes to support bigquery oauth #30674

feat(oauth): adding necessary changes to support bigquery oauth #30674

fisjac commented Oct 22, 2024 •

edited

Loading

fisjac Oct 22, 2024 •

edited

Loading

fisjac Oct 22, 2024

fisjac Oct 22, 2024

fisjac Oct 22, 2024

fisjac Oct 22, 2024

fisjac Oct 22, 2024

betodealmeida Oct 25, 2024

codecov bot commented Oct 22, 2024 •

edited

Loading

betodealmeida left a comment

geido Oct 25, 2024

geido Oct 25, 2024

betodealmeida left a comment

betodealmeida Oct 25, 2024

betodealmeida Oct 25, 2024

feat(oauth): adding necessary changes to support bigquery oauth #30674

Are you sure you want to change the base?

feat(oauth): adding necessary changes to support bigquery oauth #30674

Conversation

fisjac commented Oct 22, 2024 • edited Loading

SUMMARY

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

ADDITIONAL INFORMATION

fisjac Oct 22, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Oct 22, 2024 • edited Loading

Codecov Report

betodealmeida left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

betodealmeida left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

fisjac commented Oct 22, 2024 •

edited

Loading

fisjac Oct 22, 2024 •

edited

Loading

codecov bot commented Oct 22, 2024 •

edited

Loading