mirror of
https://github.com/apache/superset.git
synced 2026-07-02 21:05:36 +00:00
Compare commits
3 Commits
fix-sql-la
...
fix-snowfl
| Author | SHA1 | Date | |
|---|---|---|---|
|
9f57e16a34 | ||
|
|
143b1268d2 | ||
|
|
59991bf083 |
@@ -1,254 +0,0 @@
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Tests for SQL Lab immediate backend persistence functionality
|
||||
*
|
||||
* This file tests the addQueryEditorWithBackendSync function that
|
||||
* immediately creates backend tabs when SqllabBackendPersistence is enabled,
|
||||
* ensuring consistent state management and preventing race conditions
|
||||
* between frontend tab creation and backend synchronization.
|
||||
*
|
||||
* Related to issue #34997: Prevents metadata update failures by ensuring
|
||||
* tabs have valid backend IDs before users can modify them.
|
||||
*/
|
||||
|
||||
import configureMockStore from 'redux-mock-store';
|
||||
import thunk from 'redux-thunk';
|
||||
import fetchMock from 'fetch-mock';
|
||||
|
||||
import { isFeatureEnabled } from '@superset-ui/core';
|
||||
import * as actions from './sqlLab';
|
||||
|
||||
// Mock feature flags
|
||||
jest.mock('@superset-ui/core', () => ({
|
||||
...jest.requireActual('@superset-ui/core'),
|
||||
isFeatureEnabled: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockStore = configureMockStore([thunk]);
|
||||
|
||||
describe('SQL Lab immediate backend persistence', () => {
|
||||
beforeEach(() => {
|
||||
fetchMock.reset();
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
fetchMock.reset();
|
||||
isFeatureEnabled.mockRestore();
|
||||
});
|
||||
|
||||
describe('with backend persistence enabled', () => {
|
||||
beforeEach(() => {
|
||||
isFeatureEnabled.mockImplementation(
|
||||
feature => feature === 'SQLLAB_BACKEND_PERSISTENCE',
|
||||
);
|
||||
});
|
||||
|
||||
test('creates backend tab immediately and returns tabViewId', async () => {
|
||||
const store = mockStore({});
|
||||
const queryEditor = { name: 'Test Tab', sql: 'SELECT 1' };
|
||||
|
||||
// Mock successful backend creation
|
||||
fetchMock.post('glob:*/tabstateview/', { id: 42 });
|
||||
|
||||
const result = await store.dispatch(
|
||||
actions.addQueryEditorWithBackendSync(queryEditor),
|
||||
);
|
||||
|
||||
// Verify backend call was made
|
||||
const postCalls = fetchMock
|
||||
.calls()
|
||||
.filter(
|
||||
call =>
|
||||
call[0].includes('/tabstateview/') && call[1].method === 'POST',
|
||||
);
|
||||
expect(postCalls).toHaveLength(1);
|
||||
|
||||
// Verify action was dispatched with backend ID
|
||||
const addActions = store
|
||||
.getActions()
|
||||
.filter(action => action.type === 'ADD_QUERY_EDITOR');
|
||||
expect(addActions).toHaveLength(1);
|
||||
expect(addActions[0].queryEditor.tabViewId).toBe('42');
|
||||
expect(addActions[0].queryEditor.inLocalStorage).toBe(false);
|
||||
expect(addActions[0].queryEditor.loaded).toBe(true);
|
||||
|
||||
// Verify returned query editor
|
||||
expect(result.tabViewId).toBe('42');
|
||||
expect(result.inLocalStorage).toBe(false);
|
||||
});
|
||||
|
||||
test('falls back to local storage on backend failure', async () => {
|
||||
const store = mockStore({});
|
||||
const queryEditor = { name: 'Test Tab', sql: 'SELECT 1' };
|
||||
|
||||
// Mock backend failure
|
||||
fetchMock.post('glob:*/tabstateview/', { status: 500 });
|
||||
|
||||
const result = await store.dispatch(
|
||||
actions.addQueryEditorWithBackendSync(queryEditor),
|
||||
);
|
||||
|
||||
// Verify fallback action was dispatched
|
||||
const actionsDispatched = store.getActions();
|
||||
const addActions = actionsDispatched.filter(
|
||||
action => action.type === 'ADD_QUERY_EDITOR',
|
||||
);
|
||||
const toastActions = actionsDispatched.filter(
|
||||
action => action.type === 'ADD_TOAST',
|
||||
);
|
||||
|
||||
expect(addActions).toHaveLength(1);
|
||||
expect(addActions[0].queryEditor.tabViewId).toBeUndefined();
|
||||
expect(addActions[0].queryEditor.inLocalStorage).toBe(true);
|
||||
|
||||
// Should have error toast
|
||||
expect(toastActions).toHaveLength(1);
|
||||
|
||||
// Verify returned query editor (fallback)
|
||||
expect(result.tabViewId).toBeUndefined();
|
||||
expect(result.inLocalStorage).toBe(true);
|
||||
});
|
||||
|
||||
test('SQL updates work immediately after tab creation', async () => {
|
||||
const store = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [],
|
||||
unsavedQueryEditor: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Mock backend tab creation
|
||||
fetchMock.post('glob:*/tabstateview/', { id: 123 });
|
||||
// Mock SQL update
|
||||
fetchMock.put('glob:*/tabstateview/123', { status: 200 });
|
||||
|
||||
// Create tab with backend sync
|
||||
const queryEditor = await store.dispatch(
|
||||
actions.addQueryEditorWithBackendSync({
|
||||
name: 'Test',
|
||||
sql: 'SELECT 1',
|
||||
}),
|
||||
);
|
||||
|
||||
// Clear previous actions
|
||||
store.clearActions();
|
||||
|
||||
// Update state for the SQL update call
|
||||
const updatedStore = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [queryEditor],
|
||||
unsavedQueryEditor: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Now SQL updates should work immediately
|
||||
await updatedStore.dispatch(
|
||||
actions.queryEditorSetAndSaveSql(queryEditor, 'SELECT 2'),
|
||||
);
|
||||
|
||||
// Verify PUT call was made with correct integer ID
|
||||
const putCalls = fetchMock
|
||||
.calls()
|
||||
.filter(
|
||||
call =>
|
||||
call[0].includes('/tabstateview/') && call[1].method === 'PUT',
|
||||
);
|
||||
expect(putCalls).toHaveLength(1);
|
||||
expect(putCalls[0][0]).toContain('/tabstateview/123');
|
||||
});
|
||||
});
|
||||
|
||||
describe('with backend persistence disabled', () => {
|
||||
beforeEach(() => {
|
||||
isFeatureEnabled.mockImplementation(() => false);
|
||||
});
|
||||
|
||||
test('creates local tab without backend call', async () => {
|
||||
const store = mockStore({});
|
||||
const queryEditor = { name: 'Test Tab', sql: 'SELECT 1' };
|
||||
|
||||
const result = await store.dispatch(
|
||||
actions.addQueryEditorWithBackendSync(queryEditor),
|
||||
);
|
||||
|
||||
// Verify no backend calls were made
|
||||
const calls = fetchMock.calls();
|
||||
expect(calls).toHaveLength(0);
|
||||
|
||||
// Verify local action was dispatched
|
||||
const addActions = store
|
||||
.getActions()
|
||||
.filter(action => action.type === 'ADD_QUERY_EDITOR');
|
||||
expect(addActions).toHaveLength(1);
|
||||
expect(addActions[0].queryEditor.tabViewId).toBeUndefined();
|
||||
expect(addActions[0].queryEditor.inLocalStorage).toBe(true);
|
||||
|
||||
// Verify returned query editor
|
||||
expect(result.tabViewId).toBeUndefined();
|
||||
expect(result.inLocalStorage).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('integration with addNewQueryEditor', () => {
|
||||
test('addNewQueryEditor uses backend sync', async () => {
|
||||
isFeatureEnabled.mockImplementation(
|
||||
feature => feature === 'SQLLAB_BACKEND_PERSISTENCE',
|
||||
);
|
||||
|
||||
const store = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [],
|
||||
unsavedQueryEditor: {},
|
||||
tabHistory: [],
|
||||
databases: {},
|
||||
},
|
||||
common: {
|
||||
conf: {
|
||||
SQLLAB_DEFAULT_DBID: 1,
|
||||
DEFAULT_SQLLAB_LIMIT: 1000,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Mock backend creation
|
||||
fetchMock.post('glob:*/tabstateview/', { id: 999 });
|
||||
|
||||
await store.dispatch(actions.addNewQueryEditor());
|
||||
|
||||
// Verify backend call was made
|
||||
const postCalls = fetchMock
|
||||
.calls()
|
||||
.filter(
|
||||
call =>
|
||||
call[0].includes('/tabstateview/') && call[1].method === 'POST',
|
||||
);
|
||||
expect(postCalls).toHaveLength(1);
|
||||
|
||||
// Verify resulting tab has backend ID
|
||||
const addActions = store
|
||||
.getActions()
|
||||
.filter(action => action.type === 'ADD_QUERY_EDITOR');
|
||||
expect(addActions).toHaveLength(1);
|
||||
expect(addActions[0].queryEditor.tabViewId).toBe('999');
|
||||
expect(addActions[0].queryEditor.inLocalStorage).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -1,305 +0,0 @@
|
||||
/**
|
||||
* 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.
|
||||
*/
|
||||
|
||||
/**
|
||||
* Tests for SQL Lab tab state management and validation
|
||||
*
|
||||
* This file validates the behavior of SQL Lab tab state management,
|
||||
* particularly the interaction between frontend-generated IDs and backend
|
||||
* persistence. Tests edge cases and race conditions that can occur when
|
||||
* the SqllabBackendPersistence feature flag is enabled.
|
||||
*
|
||||
* Related to issue #34997: Ensures metadata updates work correctly
|
||||
* when users modify queries in newly created tabs.
|
||||
*/
|
||||
|
||||
import configureMockStore from 'redux-mock-store';
|
||||
import thunk from 'redux-thunk';
|
||||
import fetchMock from 'fetch-mock';
|
||||
|
||||
import { isFeatureEnabled } from '@superset-ui/core';
|
||||
import * as actions from './sqlLab';
|
||||
import { defaultQueryEditor } from '../fixtures';
|
||||
|
||||
// Mock feature flags
|
||||
jest.mock('@superset-ui/core', () => ({
|
||||
...jest.requireActual('@superset-ui/core'),
|
||||
isFeatureEnabled: jest.fn(),
|
||||
}));
|
||||
|
||||
const mockStore = configureMockStore([thunk]);
|
||||
|
||||
describe('SQL Lab tab state validation', () => {
|
||||
const updateTabStateEndpoint = 'glob:*/tabstateview/*';
|
||||
|
||||
beforeEach(() => {
|
||||
fetchMock.reset();
|
||||
// Enable backend persistence for these tests
|
||||
isFeatureEnabled.mockImplementation(
|
||||
feature => feature === 'SQLLAB_BACKEND_PERSISTENCE',
|
||||
);
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
fetchMock.reset();
|
||||
isFeatureEnabled.mockRestore();
|
||||
});
|
||||
|
||||
describe('New tab creation with string IDs', () => {
|
||||
test('addQueryEditor creates tabs with string nanoid IDs', () => {
|
||||
const queryEditor = { name: 'Test Tab', sql: 'SELECT 1' };
|
||||
const action = actions.addQueryEditor(queryEditor);
|
||||
|
||||
// Verify string ID was generated
|
||||
expect(typeof action.queryEditor.id).toBe('string');
|
||||
expect(action.queryEditor.id).toMatch(/^[a-zA-Z0-9_-]{11}$/);
|
||||
expect(action.queryEditor.inLocalStorage).toBe(true);
|
||||
expect(action.queryEditor.tabViewId).toBeUndefined();
|
||||
});
|
||||
|
||||
test('multiple new tabs get unique string IDs', () => {
|
||||
const editor1 = actions.addQueryEditor({ name: 'Tab 1' });
|
||||
const editor2 = actions.addQueryEditor({ name: 'Tab 2' });
|
||||
const editor3 = actions.addQueryEditor({ name: 'Tab 3' });
|
||||
|
||||
// All should have different string IDs
|
||||
expect(editor1.queryEditor.id).not.toBe(editor2.queryEditor.id);
|
||||
expect(editor2.queryEditor.id).not.toBe(editor3.queryEditor.id);
|
||||
expect(editor1.queryEditor.id).not.toBe(editor3.queryEditor.id);
|
||||
|
||||
// All should be strings
|
||||
expect(typeof editor1.queryEditor.id).toBe('string');
|
||||
expect(typeof editor2.queryEditor.id).toBe('string');
|
||||
expect(typeof editor3.queryEditor.id).toBe('string');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Backend persistence edge cases', () => {
|
||||
test('queryEditorSetAndSaveSql sends string ID to backend', async () => {
|
||||
// Create new tab with string ID (simulating real nanoid output)
|
||||
const newQueryEditor = {
|
||||
...defaultQueryEditor,
|
||||
id: 'FRRULMQgWHa', // Realistic nanoid(11) output
|
||||
inLocalStorage: true,
|
||||
// No tabViewId - simulates new tab before sync
|
||||
};
|
||||
|
||||
const store = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [newQueryEditor],
|
||||
unsavedQueryEditor: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Mock the endpoint to track what ID is sent
|
||||
fetchMock.put(updateTabStateEndpoint, { status: 404 });
|
||||
|
||||
// This should attempt to use the string ID
|
||||
const sql = 'SELECT * FROM test';
|
||||
await store.dispatch(
|
||||
actions.queryEditorSetAndSaveSql(newQueryEditor, sql),
|
||||
);
|
||||
|
||||
// Verify the string ID was sent to backend
|
||||
const putCalls = fetchMock
|
||||
.calls()
|
||||
.filter(
|
||||
call =>
|
||||
call[0].includes('/tabstateview/') && call[1].method === 'PUT',
|
||||
);
|
||||
|
||||
expect(putCalls).toHaveLength(1);
|
||||
expect(putCalls[0][0]).toContain('/tabstateview/FRRULMQgWHa');
|
||||
});
|
||||
|
||||
test('string IDs cause backend endpoint mismatch errors', async () => {
|
||||
const newQueryEditor = {
|
||||
...defaultQueryEditor,
|
||||
id: 'abc123def456', // String ID
|
||||
inLocalStorage: true,
|
||||
};
|
||||
|
||||
const store = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [newQueryEditor],
|
||||
unsavedQueryEditor: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Mock 404 response for string IDs (simulating real backend behavior)
|
||||
fetchMock.put('glob:*/tabstateview/abc123def456', { status: 404 });
|
||||
|
||||
// The function should handle the 404 but dispatch a danger toast
|
||||
await store.dispatch(
|
||||
actions.queryEditorSetAndSaveSql(newQueryEditor, 'SELECT 1'),
|
||||
);
|
||||
|
||||
// Verify error handling occurred (the function catches and handles 404 silently)
|
||||
const actionsDispatched = store.getActions();
|
||||
console.log(
|
||||
'Actions dispatched:',
|
||||
actionsDispatched.map(a => a.type),
|
||||
);
|
||||
|
||||
// Check that the PUT call was attempted with string ID
|
||||
const putCalls = fetchMock
|
||||
.calls()
|
||||
.filter(
|
||||
call =>
|
||||
call[0].includes('/tabstateview/') && call[1].method === 'PUT',
|
||||
);
|
||||
expect(putCalls).toHaveLength(1);
|
||||
expect(putCalls[0][0]).toContain('/abc123def456');
|
||||
});
|
||||
|
||||
test('existing tabs with tabViewId work correctly', async () => {
|
||||
// Existing tab with proper backend ID
|
||||
const existingQueryEditor = {
|
||||
...defaultQueryEditor,
|
||||
id: 'frontend-id-123',
|
||||
tabViewId: '42', // Integer ID from backend (as string)
|
||||
inLocalStorage: false,
|
||||
};
|
||||
|
||||
const store = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [existingQueryEditor],
|
||||
unsavedQueryEditor: {},
|
||||
},
|
||||
});
|
||||
|
||||
// Mock successful response for integer IDs
|
||||
fetchMock.put('glob:*/tabstateview/42', { status: 200 });
|
||||
|
||||
await store.dispatch(
|
||||
actions.queryEditorSetAndSaveSql(existingQueryEditor, 'SELECT 1'),
|
||||
);
|
||||
|
||||
// Verify correct endpoint was called with integer ID
|
||||
const putCalls = fetchMock
|
||||
.calls()
|
||||
.filter(
|
||||
call =>
|
||||
call[0].includes('/tabstateview/') && call[1].method === 'PUT',
|
||||
);
|
||||
|
||||
expect(putCalls).toHaveLength(1);
|
||||
expect(putCalls[0][0]).toContain('/tabstateview/42');
|
||||
});
|
||||
});
|
||||
|
||||
describe('Race condition scenarios', () => {
|
||||
test('rapid tab creation and SQL editing triggers string ID usage', async () => {
|
||||
const store = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [],
|
||||
unsavedQueryEditor: {},
|
||||
tabHistory: [],
|
||||
databases: {},
|
||||
},
|
||||
common: {
|
||||
conf: {
|
||||
SQLLAB_DEFAULT_DBID: 1,
|
||||
DEFAULT_SQLLAB_LIMIT: 1000,
|
||||
},
|
||||
},
|
||||
});
|
||||
|
||||
// Mock endpoints for this specific test
|
||||
fetchMock.put('glob:*/tabstateview/*', { status: 404 });
|
||||
fetchMock.post('glob:*/tabstateview/', { id: 999 });
|
||||
|
||||
// Simulate rapid user actions: create tab then immediately edit SQL
|
||||
await store.dispatch(actions.addNewQueryEditor());
|
||||
const addAction = store
|
||||
.getActions()
|
||||
.find(a => a.type === 'ADD_QUERY_EDITOR');
|
||||
const newTab = addAction.queryEditor;
|
||||
|
||||
// User immediately starts typing (before EditorAutoSync runs)
|
||||
await store.dispatch(
|
||||
actions.queryEditorSetAndSaveSql(newTab, 'SELECT NOW()'),
|
||||
);
|
||||
|
||||
// Should attempt to use string ID
|
||||
const putCalls = fetchMock
|
||||
.calls()
|
||||
.filter(
|
||||
call =>
|
||||
call[0].includes('/tabstateview/') && call[1].method === 'PUT',
|
||||
);
|
||||
expect(putCalls).toHaveLength(1);
|
||||
expect(putCalls[0][0]).toMatch(/\/tabstateview\/[a-zA-Z0-9_-]{11}$/);
|
||||
});
|
||||
|
||||
test('feature flag disabled should not attempt backend updates', async () => {
|
||||
// Disable backend persistence
|
||||
isFeatureEnabled.mockImplementation(() => false);
|
||||
|
||||
const newQueryEditor = {
|
||||
...defaultQueryEditor,
|
||||
id: 'string-id-123',
|
||||
inLocalStorage: true,
|
||||
};
|
||||
|
||||
const store = mockStore({
|
||||
sqlLab: {
|
||||
queryEditors: [newQueryEditor],
|
||||
unsavedQueryEditor: {},
|
||||
},
|
||||
});
|
||||
|
||||
await store.dispatch(
|
||||
actions.queryEditorSetAndSaveSql(newQueryEditor, 'SELECT 1'),
|
||||
);
|
||||
|
||||
// Should not make any backend calls
|
||||
const putCalls = fetchMock
|
||||
.calls()
|
||||
.filter(call => call[0].includes('/tabstateview/'));
|
||||
expect(putCalls).toHaveLength(0);
|
||||
});
|
||||
});
|
||||
|
||||
describe('Comparison: old vs new behavior', () => {
|
||||
test('old addQueryEditor still creates local tabs', () => {
|
||||
// The original function still works for backward compatibility
|
||||
const action = actions.addQueryEditor({ name: 'New Tab' });
|
||||
|
||||
expect(action.queryEditor.tabViewId).toBeUndefined();
|
||||
expect(action.queryEditor.inLocalStorage).toBe(true);
|
||||
});
|
||||
|
||||
test('FIXED: addQueryEditorWithBackendSync creates backend tabs immediately', async () => {
|
||||
// The new function fixes the race condition
|
||||
const store = mockStore({});
|
||||
|
||||
// Mock successful tab creation
|
||||
fetchMock.post('glob:*/tabstateview/', { id: 123 });
|
||||
|
||||
const queryEditor = await store.dispatch(
|
||||
actions.addQueryEditorWithBackendSync({ name: 'New Tab' }),
|
||||
);
|
||||
|
||||
// After fix: should have tabViewId set immediately
|
||||
expect(queryEditor.tabViewId).toBe('123');
|
||||
expect(queryEditor.inLocalStorage).toBe(false);
|
||||
});
|
||||
});
|
||||
});
|
||||
@@ -2027,9 +2027,9 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
|
||||
"""
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def get_extra_params(
|
||||
database: Database, source: QuerySource | None = None
|
||||
cls, database: Database, source: QuerySource | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Some databases require adding elements to connection parameters,
|
||||
|
||||
@@ -261,9 +261,9 @@ class DatabricksDynamicBaseEngineSpec(BasicParametersMixin, DatabricksBaseEngine
|
||||
"port": "port",
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def get_extra_params(
|
||||
database: Database, source: QuerySource | None = None
|
||||
cls, database: Database, source: QuerySource | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Add a user agent to be used in the requests.
|
||||
|
||||
@@ -78,9 +78,9 @@ class DruidEngineSpec(BaseEngineSpec):
|
||||
if orm_col.column_name == "__time":
|
||||
orm_col.is_dttm = True
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def get_extra_params(
|
||||
database: Database, source: QuerySource | None = None
|
||||
cls, database: Database, source: QuerySource | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
For Druid, the path to a SSL certificate is placed in `connect_args`.
|
||||
|
||||
@@ -300,9 +300,9 @@ class DuckDBEngineSpec(DuckDBParametersMixin, BaseEngineSpec):
|
||||
) -> set[str]:
|
||||
return set(inspector.get_table_names(schema))
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def get_extra_params(
|
||||
database: Database, source: QuerySource | None = None
|
||||
cls, database: Database, source: QuerySource | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Add a user agent to be used in the requests.
|
||||
|
||||
@@ -412,9 +412,9 @@ WHERE datistemplate = false;
|
||||
inspector.get_foreign_table_names(schema)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def get_extra_params(
|
||||
database: Database, source: QuerySource | None = None
|
||||
cls, database: Database, source: QuerySource | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
For Postgres, the path to a SSL certificate is placed in `connect_args`.
|
||||
|
||||
@@ -135,9 +135,9 @@ class SnowflakeEngineSpec(PostgresBaseEngineSpec):
|
||||
),
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def get_extra_params(
|
||||
database: Database, source: QuerySource | None = None
|
||||
cls, database: Database, source: QuerySource | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Add a user agent to be used in the requests.
|
||||
|
||||
@@ -301,9 +301,9 @@ class TrinoEngineSpec(PrestoBaseEngineSpec):
|
||||
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
@classmethod
|
||||
def get_extra_params(
|
||||
database: Database, source: QuerySource | None = None
|
||||
cls, database: Database, source: QuerySource | None = None
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Some databases require adding elements to connection parameters,
|
||||
|
||||
217
tests/unit_tests/models/test_core_database_methods.py
Normal file
217
tests/unit_tests/models/test_core_database_methods.py
Normal file
@@ -0,0 +1,217 @@
|
||||
# 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.
|
||||
|
||||
"""
|
||||
Tests for Database model methods that interact with DB engine specs.
|
||||
|
||||
These tests verify the contract between the Database model and engine specs,
|
||||
particularly around method calling patterns and signatures.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from superset.models.core import Database
|
||||
from superset.utils.core import QuerySource
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sqlalchemy_uri",
|
||||
[
|
||||
"snowflake://user:pass@account/db?role=role&warehouse=warehouse",
|
||||
"postgresql://localhost/test",
|
||||
"trino://localhost:8080/hive/default",
|
||||
"databricks://token@workspace/http_path",
|
||||
"duckdb:///path/to/db",
|
||||
],
|
||||
)
|
||||
def test_database_get_extra_with_multiple_engines(sqlalchemy_uri):
|
||||
"""
|
||||
Test that Database.get_extra() works correctly with all database engine specs.
|
||||
|
||||
This test verifies the normal production call pattern:
|
||||
Database.get_extra() -> self.db_engine_spec.get_extra_params(self, source)
|
||||
|
||||
NOTE: This tests the standard flow where db_engine_spec returns a CLASS.
|
||||
Both @staticmethod and @classmethod work in this scenario because Python
|
||||
calls the method directly on the class. The bug requiring @classmethod
|
||||
manifests in dynamic class creation scenarios (see
|
||||
test_dynamic_class_method_copying).
|
||||
"""
|
||||
database = Database(database_name="test_db", sqlalchemy_uri=sqlalchemy_uri)
|
||||
|
||||
# This is the actual production call pattern that was broken
|
||||
# It should work regardless of which engine spec is used
|
||||
extra = database.get_extra(source=QuerySource.SQL_LAB)
|
||||
assert isinstance(extra, dict)
|
||||
|
||||
# Also test without source parameter (original behavior)
|
||||
extra = database.get_extra()
|
||||
assert isinstance(extra, dict)
|
||||
|
||||
|
||||
def test_database_engine_spec_contract():
|
||||
"""
|
||||
Verify the contract between Database model and all engine specs.
|
||||
|
||||
This test ensures that all Database methods that call engine spec methods
|
||||
work correctly with the actual method signatures and decorators.
|
||||
|
||||
Tests the integration between layers, not just isolated unit tests.
|
||||
Like test_database_get_extra_with_multiple_engines, this tests the normal
|
||||
flow where methods are called directly on classes.
|
||||
"""
|
||||
# Test engines that override get_extra_params
|
||||
test_engines = [
|
||||
("snowflake://user:pass@account/db", "SnowflakeEngineSpec"),
|
||||
("postgresql://localhost/test", "PostgresEngineSpec"),
|
||||
("trino://localhost:8080/hive", "TrinoEngineSpec"),
|
||||
("databricks://token@workspace", "DatabricksNativeEngineSpec"),
|
||||
("duckdb:///test.db", "DuckDBEngineSpec"),
|
||||
]
|
||||
|
||||
for sqlalchemy_uri, engine_name in test_engines:
|
||||
database = Database(
|
||||
database_name=f"test_{engine_name.lower()}", sqlalchemy_uri=sqlalchemy_uri
|
||||
)
|
||||
|
||||
# Test that the calling pattern matches the method signature
|
||||
try:
|
||||
# This calls: self.db_engine_spec.get_extra_params(self, source)
|
||||
# If get_extra_params is @staticmethod, this will fail with:
|
||||
# "takes 2 positional arguments but 3 were given"
|
||||
result = database.get_extra(source=QuerySource.SQL_LAB)
|
||||
assert isinstance(result, dict), (
|
||||
f"{engine_name} should return dict from get_extra_params"
|
||||
)
|
||||
|
||||
except TypeError as e:
|
||||
if "positional arguments" in str(e):
|
||||
pytest.fail(
|
||||
f"{engine_name}.get_extra_params has incompatible signature. "
|
||||
f"Error: {e}. "
|
||||
"This suggests @staticmethod should be @classmethod."
|
||||
)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def test_regression_snowflake_oauth2_get_extra_params():
|
||||
"""
|
||||
Regression test for Snowflake OAuth2 get_extra_params bug.
|
||||
This test reproduces the standard scenario to ensure it works.
|
||||
The actual "SnowflakeOAuth2Override" bug is tested in
|
||||
test_dynamic_class_method_copying.
|
||||
"""
|
||||
database = Database(
|
||||
database_name="snowflake_oauth_test",
|
||||
sqlalchemy_uri="snowflake://user:pass@account/db?role=role&warehouse=warehouse",
|
||||
)
|
||||
|
||||
# This specific call pattern caused the original error
|
||||
# The error occurred when OAuth2 was enabled, creating a dynamic override class
|
||||
try:
|
||||
extra = database.get_extra(source=QuerySource.SQL_LAB)
|
||||
assert isinstance(extra, dict)
|
||||
assert "engine_params" in extra
|
||||
assert "connect_args" in extra["engine_params"]
|
||||
|
||||
except TypeError as e:
|
||||
if "takes 2 positional arguments but 3 were given" in str(e):
|
||||
pytest.fail(
|
||||
f"Snowflake OAuth2 method signature issue not fixed: {e}. "
|
||||
"get_extra_params should use @classmethod, not @staticmethod."
|
||||
)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def test_dynamic_class_method_copying():
|
||||
"""
|
||||
Test that engine spec methods work when copied to dynamic classes.
|
||||
"""
|
||||
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
|
||||
|
||||
# Simulate creating a dynamic override class (like SnowflakeOAuth2Override)
|
||||
# This is what likely happens in OAuth2 configurations or extensions
|
||||
dynamic_override = type(
|
||||
"SnowflakeOAuth2Override",
|
||||
(),
|
||||
{"get_extra_params": SnowflakeEngineSpec.get_extra_params},
|
||||
)
|
||||
|
||||
# Create an instance of the dynamic class
|
||||
override_instance = dynamic_override()
|
||||
|
||||
# Create a database for testing
|
||||
database = Database(
|
||||
database_name="test_dynamic_override",
|
||||
sqlalchemy_uri="snowflake://user:pass@account/db",
|
||||
)
|
||||
|
||||
# Test calling the method through the instance (this is where the bug occurred)
|
||||
# With @staticmethod: fails with "takes 2 positional arguments but 3 were given"
|
||||
# With @classmethod: works correctly
|
||||
try:
|
||||
result = override_instance.get_extra_params(database, QuerySource.SQL_LAB)
|
||||
assert isinstance(result, dict)
|
||||
assert "engine_params" in result
|
||||
assert "connect_args" in result["engine_params"]
|
||||
assert "application" in result["engine_params"]["connect_args"]
|
||||
except TypeError as e:
|
||||
if "takes 2 positional arguments but 3 were given" in str(e):
|
||||
pytest.fail(
|
||||
f"Dynamic class method copying failed: {e}. "
|
||||
"This indicates get_extra_params uses @staticmethod when it should "
|
||||
"use @classmethod. The bug manifests when methods are copied to "
|
||||
"dynamic classes and called through instances."
|
||||
)
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
def test_method_reassignment_compatibility():
|
||||
"""
|
||||
Test that engine spec methods can be reassigned and remain compatible.
|
||||
"""
|
||||
from superset.db_engine_specs.snowflake import SnowflakeEngineSpec
|
||||
|
||||
# Test method reassignment (another pattern that can cause issues)
|
||||
class CustomEngineSpec:
|
||||
pass
|
||||
|
||||
# Assign the method to a different class
|
||||
CustomEngineSpec.get_extra_params = SnowflakeEngineSpec.get_extra_params
|
||||
|
||||
database = Database(
|
||||
database_name="test_reassignment",
|
||||
sqlalchemy_uri="snowflake://user:pass@account/db",
|
||||
)
|
||||
|
||||
# Test calling through the reassigned class
|
||||
try:
|
||||
result = CustomEngineSpec.get_extra_params(database, QuerySource.SQL_LAB)
|
||||
assert isinstance(result, dict)
|
||||
assert "engine_params" in result
|
||||
except TypeError as e:
|
||||
if "positional arguments" in str(e):
|
||||
pytest.fail(
|
||||
f"Method reassignment failed: {e}. "
|
||||
"This suggests the method decorator is incompatible with method "
|
||||
"reassignment patterns."
|
||||
)
|
||||
else:
|
||||
raise
|
||||
Reference in New Issue
Block a user