feat(SIP-85): OAuth2 for databases (apache#27631)

pyproject.toml

@@ -79,7 +79,7 @@ dependencies = [
     "PyJWT>=2.4.0, <3.0",
     "redis>=4.6.0, <5.0",
     "selenium>=3.141.0, <4.10.0",
-    "shillelagh[gsheetsapi]>=1.2.10, <2.0",
+    "shillelagh[gsheetsapi]>=1.2.18, <2.0",
     "shortid",
     "sshtunnel>=0.4.0, <0.5",
     "simplejson>=3.15.0",
@@ -127,13 +127,14 @@ excel = ["xlrd>=1.2.0, <1.3"]
 firebird = ["sqlalchemy-firebird>=0.7.0, <0.8"]
 firebolt = ["firebolt-sqlalchemy>=1.0.0, <2"]
 gevent = ["gevent>=23.9.1"]
-gsheets = ["shillelagh[gsheetsapi]>=1.2.10, <2"]
+gsheets = ["shillelagh[gsheetsapi]>=1.2.18, <2"]
 hana = ["hdbcli==2.4.162", "sqlalchemy_hana==0.4.0"]
 hive = [
     "pyhive[hive]>=0.6.5;python_version<'3.11'",
     "pyhive[hive_pure_sasl]>=0.7.0",
     "tableschema",
     "thrift>=0.14.1, <1.0.0",
+    "thrift_sasl>=0.4.3, < 1.0.0",
 ]
 impala = ["impyla>0.16.2, <0.17"]
 kusto = ["sqlalchemy-kusto>=2.0.0, <3"]
@@ -155,7 +156,7 @@ trino = ["trino>=0.328.0"]
 prophet = ["prophet>=1.1.5, <2"]
 redshift = ["sqlalchemy-redshift>=0.8.1, <0.9"]
 rockset = ["rockset-sqlalchemy>=0.0.1, <1"]
-shillelagh = ["shillelagh[all]>=1.2.10, <2"]
+shillelagh = ["shillelagh[all]>=1.2.18, <2"]
 snowflake = ["snowflake-sqlalchemy>=1.2.4, <2"]
 spark = [
     "pyhive[hive]>=0.6.5;python_version<'3.11'",

requirements/base.txt

@@ -84,6 +84,7 @@ cryptography==42.0.5
     # via
     #   apache-superset
     #   paramiko
+    #   pyopenssl
 deprecated==1.2.13
     # via limits
 deprecation==2.1.0
@@ -147,7 +148,9 @@ geopy==2.4.1
 google-auth==2.27.0
     # via shillelagh
 greenlet==3.0.3
-    # via shillelagh
+    # via
+    #   shillelagh
+    #   sqlalchemy
 gunicorn==21.2.0
     # via apache-superset
 hashids==1.3.1
@@ -278,6 +281,8 @@ pyjwt==2.8.0
     #   flask-jwt-extended
 pynacl==1.5.0
     # via paramiko
+pyopenssl==24.1.0
+    # via shillelagh
 pyparsing==3.1.2
     # via apache-superset
 pyrsistent==0.20.0
@@ -319,7 +324,7 @@ rsa==4.9
     # via google-auth
 selenium==3.141.0
     # via apache-superset
-shillelagh[gsheetsapi]==1.2.10
+shillelagh[gsheetsapi]==1.2.18
     # via apache-superset
 shortid==0.1.2
     # via apache-superset

requirements/development.txt

@@ -202,6 +202,8 @@ ptyprocess==0.7.0
     # via pexpect
 pure-eval==0.2.2
     # via stack-data
+pure-sasl==0.6.2
+    # via thrift-sasl
 pydata-google-auth==1.7.0
     # via pandas-gbq
 pydruid==0.6.6
@@ -252,6 +254,10 @@ statsd==4.0.1
 tableschema==1.20.10
     # via apache-superset
 thrift==0.16.0
+    # via
+    #   apache-superset
+    #   thrift-sasl
+thrift-sasl==0.4.3
     # via apache-superset
 tomli==2.0.1
     # via

superset-frontend/src/components/ErrorMessage/OAuth2RedirectMessage.test.tsx

@@ -0,0 +1,170 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import React from 'react';
+import * as reduxHooks from 'react-redux';
+import { Provider } from 'react-redux';
+import { createStore } from 'redux';
+import { render, fireEvent, waitFor } from '@testing-library/react';
+import '@testing-library/jest-dom';
+import { ThemeProvider, supersetTheme } from '@superset-ui/core';
+import OAuth2RedirectMessage from 'src/components/ErrorMessage/OAuth2RedirectMessage';
+import {
+  ErrorLevel,
+  ErrorSource,
+  ErrorTypeEnum,
+} from 'src/components/ErrorMessage/types';
+import { reRunQuery } from 'src/SqlLab/actions/sqlLab';
+import { triggerQuery } from 'src/components/Chart/chartAction';
+import { onRefresh } from 'src/dashboard/actions/dashboardState';
+
+// Mock the Redux store
+const mockStore = createStore(() => ({
+  sqlLab: {
+    queries: { 'query-id': { sql: 'SELECT * FROM table' } },
+    queryEditors: [{ id: 'editor-id', latestQueryId: 'query-id' }],
+    tabHistory: ['editor-id'],
+  },
+  explore: {
+    slice: { slice_id: 123 },
+  },
+  charts: { '1': {}, '2': {} },
+  dashboardInfo: { id: 'dashboard-id' },
+}));
+
+// Mock actions
+jest.mock('src/SqlLab/actions/sqlLab', () => ({
+  reRunQuery: jest.fn(),
+}));
+
+jest.mock('src/components/Chart/chartAction', () => ({
+  triggerQuery: jest.fn(),
+}));
+
+jest.mock('src/dashboard/actions/dashboardState', () => ({
+  onRefresh: jest.fn(),
+}));
+
+// Mock useDispatch
+const mockDispatch = jest.fn();
+jest.spyOn(reduxHooks, 'useDispatch').mockReturnValue(mockDispatch);
+
+// Mock global window functions
+const mockOpen = jest.spyOn(window, 'open').mockImplementation(() => null);
+const mockAddEventListener = jest.spyOn(window, 'addEventListener');
+const mockRemoveEventListener = jest.spyOn(window, 'removeEventListener');
+
+// Mock window.postMessage
+const originalPostMessage = window.postMessage;
+
+beforeEach(() => {
+  window.postMessage = jest.fn();
+});
+
+afterEach(() => {
+  window.postMessage = originalPostMessage;
+});
+
+function simulateMessageEvent(data: any, origin: string) {
+  const messageEvent = new MessageEvent('message', { data, origin });
+  window.dispatchEvent(messageEvent);
+}
+
+const defaultProps = {
+  error: {
+    error_type: ErrorTypeEnum.OAUTH2_REDIRECT,
+    message: "You don't have permission to access the data.",
+    extra: {
+      url: 'https://example.com',
+      tab_id: 'tabId',
+      redirect_uri: 'https://redirect.example.com',
+    },
+    level: 'warning' as ErrorLevel,
+  },
+  source: 'sqllab' as ErrorSource,
+};
+
+const setup = (overrides = {}) => (
+  <ThemeProvider theme={supersetTheme}>
+    <Provider store={mockStore}>
+      <OAuth2RedirectMessage {...defaultProps} {...overrides} />;
+    </Provider>
+  </ThemeProvider>
+);
+
+describe('OAuth2RedirectMessage Component', () => {
+  it('renders without crashing and displays the correct initial UI elements', () => {
+    const { getByText } = render(setup());
+
+    expect(getByText(/Authorization needed/i)).toBeInTheDocument();
+    expect(getByText(/provide authorization/i)).toBeInTheDocument();
+  });
+
+  it('opens a new window with the correct URL when the link is clicked', () => {
+    const { getByText } = render(setup());
+
+    const linkElement = getByText(/provide authorization/i);
+    fireEvent.click(linkElement);
+
+    expect(mockOpen).toHaveBeenCalledWith('https://example.com', '_blank');
+  });
+
+  it('cleans up the message event listener on unmount', () => {
+    const { unmount } = render(setup());
+
+    expect(mockAddEventListener).toHaveBeenCalled();
+    unmount();
+    expect(mockRemoveEventListener).toHaveBeenCalled();
+  });
+
+  it('dispatches reRunQuery action when a message with correct tab ID is received for SQL Lab', async () => {
+    render(setup());
+
+    simulateMessageEvent({ tabId: 'tabId' }, 'https://redirect.example.com');
+
+    await waitFor(() => {
+      expect(reRunQuery).toHaveBeenCalledWith({ sql: 'SELECT * FROM table' });
+    });
+  });
+
+  it('dispatches triggerQuery action for explore source upon receiving a correct message', async () => {
+    render(setup({ source: 'explore' }));
+
+    simulateMessageEvent({ tabId: 'tabId' }, 'https://redirect.example.com');
+
+    await waitFor(() => {
+      expect(triggerQuery).toHaveBeenCalledWith(true, 123);
+    });
+  });
+
+  it('dispatches onRefresh action for dashboard source upon receiving a correct message', async () => {
+    render(setup({ source: 'dashboard' }));
+
+    simulateMessageEvent({ tabId: 'tabId' }, 'https://redirect.example.com');
+
+    await waitFor(() => {
+      expect(onRefresh).toHaveBeenCalledWith(
+        ['1', '2'],
+        true,
+        0,
+        'dashboard-id',
+      );
+    });
+  });
+});