diff --git a/UPDATING.md b/UPDATING.md index fdbfd8821497..fca1bf30e1a4 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -873,6 +873,11 @@ error notice. This does not authorize replay of data-bearing notifications. - The Pinot extra requires pinotdb[sqlalchemy]>=8.0.0,<10.0.0. Earlier releases declare SQLAlchemy below 2 in their SQLAlchemy extra. - The Databricks extra requires databricks-sqlalchemy 2.x (at least 2.0.1). The 1.x dialect requires SQLAlchemy below 2. + Saved `databricks+connector://` connections are automatically migrated to + `databricks://` by `superset db upgrade`, preserving the separately encrypted + token and connection settings. Downgrade restores converted connections left + unchanged since the upgrade and imported connections carrying the conversion + marker, not connections created natively with the new scheme. - `superset deletion-retention force-purge` now exits **1** when the target is blocked by a deletion rule or is not found (the messages are unchanged), so a diff --git a/docs/docs/using-superset/creating-your-first-dashboard.mdx b/docs/docs/using-superset/creating-your-first-dashboard.mdx index 5aecfc457c4a..79e216fb74ca 100644 --- a/docs/docs/using-superset/creating-your-first-dashboard.mdx +++ b/docs/docs/using-superset/creating-your-first-dashboard.mdx @@ -83,6 +83,11 @@ for your target database. Click the **Test Connection** button to confirm things work end to end. If the connection looks good, save the configuration by clicking the **Connect** button in the bottom right corner of the modal window: +Closing the connection window starts a fresh session the next time you add a database. +A pending save may still complete on the server, but its response does not populate +or close the new connection window. Check the database list before retrying a save +that was pending when you closed the window. + Congratulations, you've just added a new data source in Superset! ### Sharing a Database Connection diff --git a/docs/src/data/databases.json b/docs/src/data/databases.json index b0f1a236eced..6c7c291a2b1d 100644 --- a/docs/src/data/databases.json +++ b/docs/src/data/databases.json @@ -3695,11 +3695,11 @@ "notes": "Requires ODBC driver. For serverless SQL warehouses." }, { - "name": "databricks-dbapi (Legacy)", - "pypi_package": "databricks-dbapi[sqlalchemy]", - "connection_string": "databricks+connector://token:{access_token}@{host}:{port}/{database}", + "name": "Field-based connection", + "pypi_package": "apache-superset[databricks]", + "connection_string": "databricks://token:{access_token}@{host}:{port}/{database}", "is_recommended": false, - "notes": "Legacy connector. Use Python Connector for new deployments." + "notes": "Use the dynamic form to enter connection fields individually." } ], "category": "Cloud Data Warehouses" @@ -3707,7 +3707,7 @@ "engine": "databricks", "engine_name": "Databricks", "engine_aliases": [], - "default_driver": "databricks-sql-python", + "default_driver": "databricks", "supports_file_upload": true, "supports_dynamic_schema": true, "supports_catalog": true, @@ -3779,7 +3779,7 @@ "max_score": 201, "documentation": { "pypi_packages": [], - "connection_string": "databricks+connector://token:{access_token}@{host}:{port}/{database_name}", + "connection_string": "databricks://token:{access_token}@{host}:{port}/{database_name}", "category": "Cloud Data Warehouses" }, "engine": "databricks", diff --git a/docs/user_docs_versioned_docs/version-6.1.0/_versioned_data/src/data/databases.json b/docs/user_docs_versioned_docs/version-6.1.0/_versioned_data/src/data/databases.json index bc0ac9de7978..65dc355d1f56 100644 --- a/docs/user_docs_versioned_docs/version-6.1.0/_versioned_data/src/data/databases.json +++ b/docs/user_docs_versioned_docs/version-6.1.0/_versioned_data/src/data/databases.json @@ -3672,11 +3672,11 @@ "notes": "Requires ODBC driver. For serverless SQL warehouses." }, { - "name": "databricks-dbapi (Legacy)", - "pypi_package": "databricks-dbapi[sqlalchemy]", - "connection_string": "databricks+connector://token:{access_token}@{host}:{port}/{database}", + "name": "Field-based connection", + "pypi_package": "apache-superset[databricks]", + "connection_string": "databricks://token:{access_token}@{host}:{port}/{database}", "is_recommended": false, - "notes": "Legacy connector. Use Python Connector for new deployments." + "notes": "Use the dynamic form to enter connection fields individually." } ], "category": "Cloud Data Warehouses" @@ -3684,7 +3684,7 @@ "engine": "databricks", "engine_name": "Databricks", "engine_aliases": [], - "default_driver": "databricks-sql-python", + "default_driver": "databricks", "supports_file_upload": true, "supports_dynamic_schema": true, "supports_catalog": true, @@ -3751,7 +3751,7 @@ "max_score": 201, "documentation": { "pypi_packages": [], - "connection_string": "databricks+connector://token:{access_token}@{host}:{port}/{database_name}", + "connection_string": "databricks://token:{access_token}@{host}:{port}/{database_name}", "category": "Cloud Data Warehouses" }, "engine": "databricks", diff --git a/docs/user_docs_versioned_docs/version-6.1.0/databases/supported/databricks.mdx b/docs/user_docs_versioned_docs/version-6.1.0/databases/supported/databricks.mdx index 1415a8d323eb..3a4ac335ad37 100644 --- a/docs/user_docs_versioned_docs/version-6.1.0/databases/supported/databricks.mdx +++ b/docs/user_docs_versioned_docs/version-6.1.0/databases/supported/databricks.mdx @@ -78,12 +78,12 @@ export const databaseInfo = { notes: 'Requires ODBC driver. For serverless SQL warehouses.', }, { - name: 'databricks-dbapi (Legacy)', - pypi_package: 'databricks-dbapi[sqlalchemy]', + name: 'Field-based connection', + pypi_package: 'apache-superset[databricks]', connection_string: - 'databricks+connector://token:{access_token}@{host}:{port}/{database}', + 'databricks://token:{access_token}@{host}:{port}/{database}', is_recommended: false, - notes: 'Legacy connector. Use Python Connector for new deployments.', + notes: 'Use the dynamic form to enter connection fields individually.', }, ], }, @@ -126,3 +126,38 @@ export const databaseInfo = { }; + + +### Upgrading legacy connections + +Workspaces upgrading from 6.0 with a legacy Databricks connection have their +`databricks+connector://` URI automatically migrated to `databricks://` during +`superset db upgrade`. The migration preserves credentials, HTTP path, catalog, +and the rest of the URI. No API access or manual connection recreation is needed. + +Edit the connection to review the saved host, HTTP path, catalog, schema, and +masked access token, or enter a replacement token. When creating a connection +with a SQLAlchemy URI, **Connect this database using the dynamic form instead** +opens the field-based form. + +Downgrading restores the old scheme on locally converted connections only when +their URI, settings, and modification timestamp still match the migrated state. +Imported connections carrying the conversion marker are also eligible for rollback. Connections created natively with `databricks://` are not reverted. Malformed +`extra` settings do not block the upgrade: their text is preserved, and only the URI +scheme changes. The migration keeps fingerprints in a bookkeeping table so these +connections can also be safely downgraded without parsing their settings. Invalid +settings may still need correction before the connection can execute queries. + +The Default Catalog field is optional. Connections converted from the legacy +connector retain the legacy runtime default catalog when none is configured, +so existing dataset catalogs and schema permissions remain consistent. Clearing +a field removes its corresponding legacy `connect_args` override. Changing the +connection destination requires re-entering the access token. + +The new-connection picker offers **Databricks**, not **Databricks (legacy)**. +Existing legacy connections can still be edited with their saved form. + +To create a separate connection after editing one, close the edit dialog and +choose **+ Database**. The database picker starts a new connection without +reusing the edited connection's name or URI, even if the previous connection +was still loading when the edit dialog was closed. diff --git a/superset-frontend/src/features/databases/DatabaseModal/index.test.tsx b/superset-frontend/src/features/databases/DatabaseModal/index.test.tsx index f605e6e328f1..5e79ffea750c 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/index.test.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/index.test.tsx @@ -19,6 +19,7 @@ // TODO: These tests should be made atomic in separate files +import { useState } from 'react'; import fetchMock from 'fetch-mock'; import { render, @@ -27,6 +28,7 @@ import { within, waitFor, fireEvent, + act, } from 'spec/helpers/testing-library'; import { getExtensionsRegistry } from '@superset-ui/core'; import setupCodeOverrides from 'src/setup/setupCodeOverrides'; @@ -44,6 +46,8 @@ jest.mock('@superset-ui/core', () => ({ isFeatureEnabled: () => true, })); +const realUseSingleViewResource = hooks.useSingleViewResource; + const mockHistoryPush = jest.fn(); jest.mock('react-router-dom', () => ({ ...jest.requireActual('react-router-dom'), @@ -336,6 +340,385 @@ describe('DatabaseModal', () => { useRedux: true, }); + /** Check that closing a mounted modal cannot reuse a saved connection id. */ + const expectFreshCreateAfterEdit = async ( + closeAction: 'Close' | 'X' | 'Finish' | 'Late X', + ): Promise => { + const resourceHook = jest.spyOn(hooks, 'useSingleViewResource'); + const previousImplementation = resourceHook.getMockImplementation()!; + resourceHook.mockImplementation(realUseSingleViewResource); + let releaseEditFetch: () => void = () => {}; + const editFetchReady = new Promise(resolve => { + releaseEditFetch = resolve; + }); + const editResponse = { + result: { + id: 10, + database_name: 'Existing connection', + backend: 'sqlite', + configuration_method: ConfigurationMethod.SqlalchemyUri, + sqlalchemy_uri: 'sqlite:///existing.db', + }, + }; + fetchMock.get( + 'glob:*/api/v1/database/10/connection', + closeAction === 'Late X' + ? async () => { + await editFetchReady; + return editResponse; + } + : editResponse, + ); + fetchMock.put(DATABASE_FETCH_ENDPOINT, { + id: 10, + result: { + database_name: 'Existing connection', + backend: 'sqlite', + configuration_method: ConfigurationMethod.SqlalchemyUri, + sqlalchemy_uri: 'sqlite:///existing.db', + }, + }); + fetchMock.post('glob:*/api/v1/database/test_connection/', { + message: 'OK', + }); + const Harness = () => { + const [show, setShow] = useState(true); + const [databaseId, setDatabaseId] = useState(10); + return ( + <> + + + { + setShow(false); + setDatabaseId(undefined); + }} + /> + + ); + }; + try { + render(, { useRedux: true }); + if (closeAction === 'Late X') { + await waitFor(() => + expect( + fetchMock.callHistory.calls('glob:*/api/v1/database/10/connection'), + ).toHaveLength(1), + ); + } else { + expect( + await screen.findByDisplayValue('Existing connection'), + ).toBeInTheDocument(); + } + await userEvent.click(document.querySelector('.ant-modal-wrap')!); + expect(screen.getByRole('dialog')).toBeInTheDocument(); + if (closeAction === 'Finish') { + await userEvent.click(screen.getByRole('button', { name: 'Finish' })); + } else if (closeAction === 'X' || closeAction === 'Late X') { + await userEvent.click(screen.getAllByLabelText('Close')[0]); + } else { + await userEvent.click(screen.getByText('Close')); + } + await waitFor(() => + expect(screen.queryByRole('dialog')).not.toBeInTheDocument(), + ); + if (closeAction !== 'Late X') fetchMock.clearHistory(); + await userEvent.click(screen.getByRole('button', { name: '+ Database' })); + await screen.findByRole('dialog'); + if (closeAction === 'Late X') { + await act(async () => { + releaseEditFetch(); + await fetchMock.callHistory.flush(true); + }); + } + const picker = screen.queryByRole('heading', { + name: /select a database to connect/i, + }); + // Exercise saving even if stale state incorrectly skips the picker. + if (picker) { + await userEvent.click( + await screen.findByRole('button', { name: /sqlite/i }), + ); + } + fireEvent.change(await screen.findByTestId('database-name-input'), { + target: { value: 'My new sandbox' }, + }); + fireEvent.change(screen.getByTestId('sqlalchemy-uri-input'), { + target: { value: 'sqlite:///sandbox.db' }, + }); + await userEvent.click( + screen.getByRole('button', { name: /test connection/i }), + ); + const connect = screen.getByRole('button', { name: 'Connect' }); + await waitFor(() => expect(connect).toBeEnabled()); + await userEvent.click(connect); + await waitFor(() => + expect( + fetchMock.callHistory.calls(DATABASE_CONNECT_ENDPOINT).length + + fetchMock.callHistory.calls(DATABASE_FETCH_ENDPOINT, { + method: 'PUT', + }).length, + ).toBeGreaterThan(0), + ); + expect( + fetchMock.callHistory.calls(DATABASE_FETCH_ENDPOINT, { + method: 'PUT', + }), + ).toHaveLength(0); + expect(picker).not.toBeNull(); + expect( + fetchMock.callHistory.calls(DATABASE_CONNECT_ENDPOINT), + ).toHaveLength(1); + expect( + JSON.parse( + fetchMock.callHistory.lastCall(DATABASE_CONNECT_ENDPOINT)!.options + .body as string, + ), + ).not.toHaveProperty('id'); + await waitFor(() => + expect(screen.queryByRole('dialog')).not.toBeInTheDocument(), + ); + // A successful create also updates the resource cache before closing. + await userEvent.click(screen.getByRole('button', { name: '+ Database' })); + expect( + await screen.findByRole('heading', { + name: /select a database to connect/i, + }), + ).toBeInTheDocument(); + await userEvent.click(screen.getAllByLabelText('Close')[0]); + await waitFor(() => + expect(screen.queryByRole('dialog')).not.toBeInTheDocument(), + ); + await userEvent.click( + screen.getByRole('button', { name: 'Edit existing' }), + ); + expect( + await screen.findByDisplayValue('Existing connection'), + ).toBeInTheDocument(); + } finally { + releaseEditFetch(); + resourceHook.mockImplementation(previousImplementation); + } + }; + + test('closing edit then creating uses POST without the old id', async () => { + await expectFreshCreateAfterEdit('Close'); + }); + + test('X edit then creating uses POST without the old id', async () => { + await expectFreshCreateAfterEdit('X'); + }); + + test('saving edit then creating uses POST without the old id', async () => { + await expectFreshCreateAfterEdit('Finish'); + }); + + test('late edit fetch after X then creating uses POST without the old id', async () => { + await expectFreshCreateAfterEdit('Late X'); + }); + + /** Check that late writes cannot populate or close a subsequent create session. */ + const expectFreshCreateAfterLateWrite = async ( + operation: 'create' | 'update', + ): Promise => { + const resourceHook = jest.spyOn(hooks, 'useSingleViewResource'); + const previousImplementation = resourceHook.getMockImplementation()!; + resourceHook.mockImplementation(realUseSingleViewResource); + let releaseWrite: () => void = () => {}; + const writeReady = new Promise(resolve => { + releaseWrite = resolve; + }); + const connection = { + id: 77, + database_name: 'First connection', + backend: 'postgresql', + configuration_method: ConfigurationMethod.DynamicForm, + }; + let createCount = 0; + fetchMock.modifyRoute('database-connect', { + response: async () => { + createCount += 1; + if (operation === 'create' && createCount === 1) await writeReady; + const id = createCount === 1 ? 77 : 88; + return { id, result: { ...connection, id } }; + }, + }); + fetchMock.put('glob:*/api/v1/database/77', async () => { + await writeReady; + return { id: 77, result: connection }; + }); + const onDatabaseAdd = jest.fn(); + const Harness = () => { + const [show, setShow] = useState(true); + return ( + <> + + setShow(false)} + onDatabaseAdd={onDatabaseAdd} + /> + + ); + }; + /** Fill a new dynamic-form connection without relying on saved state. */ + const connectPostgres = async (name: string, host: string) => { + await userEvent.click( + await screen.findByRole('button', { name: /postgresql/i }), + ); + fireEvent.change( + document.querySelector( + 'input[name="database_name"]', + )!, + { + target: { value: name }, + }, + ); + for (const [field, value] of Object.entries({ + host, + port: '5432', + database: 'postgres', + username: 'testdb', + })) { + fireEvent.change( + document.querySelector(`input[name="${field}"]`)!, + { target: { value } }, + ); + } + fireEvent.blur( + document.querySelector('input[name="username"]')!, + ); + const connect = screen.getByRole('button', { name: 'Connect' }); + await waitFor(() => expect(connect).toBeEnabled()); + await userEvent.click(connect); + }; + try { + render(, { useRedux: true }); + await connectPostgres('First connection', 'first.example'); + await waitFor(() => + expect( + fetchMock.callHistory.calls(DATABASE_CONNECT_ENDPOINT), + ).toHaveLength(1), + ); + if (operation === 'update') { + expect(await screen.findByText(/step 3 of 3/i)).toBeInTheDocument(); + await userEvent.click(screen.getByRole('button', { name: 'Finish' })); + await waitFor(() => + expect( + fetchMock.callHistory.calls('glob:*/api/v1/database/77'), + ).toHaveLength(1), + ); + } + await userEvent.click(screen.getAllByLabelText('Close')[0]); + await waitFor(() => + expect(screen.queryByRole('dialog')).not.toBeInTheDocument(), + ); + await userEvent.click(screen.getByRole('button', { name: '+ Database' })); + await screen.findByRole('dialog'); + onDatabaseAdd.mockClear(); + await act(async () => { + releaseWrite(); + await fetchMock.callHistory.flush(true); + }); + expect( + screen.getByRole('heading', { + name: /select a database to connect/i, + }), + ).toBeInTheDocument(); + expect(onDatabaseAdd).not.toHaveBeenCalled(); + fetchMock.clearHistory(); + await connectPostgres('Second connection', 'second.example'); + expect(await screen.findByText(/step 3 of 3/i)).toBeInTheDocument(); + expect( + fetchMock.callHistory.calls(DATABASE_CONNECT_ENDPOINT), + ).toHaveLength(1); + expect( + fetchMock.callHistory.calls('glob:*/api/v1/database/77', { + method: 'PUT', + }), + ).toHaveLength(0); + const payload = JSON.parse( + fetchMock.callHistory.lastCall(DATABASE_CONNECT_ENDPOINT)!.options + .body as string, + ); + expect(payload).not.toHaveProperty('id'); + expect(payload.database_name).toBe('Second connection'); + expect(payload.parameters.host).toBe('second.example'); + } finally { + releaseWrite(); + resourceHook.mockImplementation(previousImplementation); + } + }; + + test('late create response after X leaves the next connection as a fresh POST', async () => { + await expectFreshCreateAfterLateWrite('create'); + }); + + test('late update response after X leaves the next connection as a fresh POST', async () => { + await expectFreshCreateAfterLateWrite('update'); + }); + + test('switching edit targets loads the active target while the previous fetch is pending', async () => { + const resourceHook = jest.spyOn(hooks, 'useSingleViewResource'); + const previousImplementation = resourceHook.getMockImplementation()!; + resourceHook.mockImplementation(realUseSingleViewResource); + let releaseEditFetch: () => void = () => {}; + const editFetchReady = new Promise(resolve => { + releaseEditFetch = resolve; + }); + const connection = { + backend: 'sqlite', + configuration_method: ConfigurationMethod.SqlalchemyUri, + sqlalchemy_uri: 'sqlite:///existing.db', + }; + fetchMock.get('glob:*/api/v1/database/10/connection', async () => { + await editFetchReady; + return { + result: { ...connection, id: 10, database_name: 'Previous target' }, + }; + }); + fetchMock.get('glob:*/api/v1/database/20/connection', { + result: { ...connection, id: 20, database_name: 'Active target' }, + }); + try { + const { rerender } = setup({ databaseId: 10 }); + await waitFor(() => + expect( + fetchMock.callHistory.calls('glob:*/api/v1/database/10/connection'), + ).toHaveLength(1), + ); + rerender(); + expect( + await screen.findByDisplayValue('Active target'), + ).toBeInTheDocument(); + await act(async () => { + releaseEditFetch(); + await fetchMock.callHistory.flush(true); + }); + expect(screen.getByDisplayValue('Active target')).toBeInTheDocument(); + expect( + screen.queryByDisplayValue('Previous target'), + ).not.toBeInTheDocument(); + } finally { + releaseEditFetch(); + resourceHook.mockImplementation(previousImplementation); + } + }); + // eslint-disable-next-line no-restricted-globals -- TODO: Migrate from describe blocks describe('Visual: New database connection', () => { test('renders the initial load of Step 1 correctly', async () => { @@ -1484,6 +1867,129 @@ describe('DatabaseModal', () => { }); }); + test.each(['Finish', 'Back'] as const)( + 'keeps a created PostgreSQL field connection on step 3 without a saved schema: %s', + async action => { + const resourceHook = jest.spyOn(hooks, 'useSingleViewResource'); + const previousImplementation = resourceHook.getMockImplementation()!; + resourceHook.mockImplementation(realUseSingleViewResource); + let releaseBackFetch: () => void = () => {}; + const backFetchReady = new Promise(resolve => { + releaseBackFetch = resolve; + }); + fetchMock.get('glob:*/api/v1/database/10/connection', async () => { + await backFetchReady; + return { + result: { + id: 10, + database_name: 'PostgreSQL', + backend: 'postgresql', + configuration_method: ConfigurationMethod.DynamicForm, + }, + }; + }); + try { + fetchMock.modifyRoute('database-connect', { + response: { + id: 10, + result: { + id: 10, + database_name: 'PostgreSQL', + backend: 'postgresql', + configuration_method: ConfigurationMethod.DynamicForm, + }, + }, + }); + fetchMock.put(DATABASE_FETCH_ENDPOINT, { result: { id: 10 } }); + const addDangerToast = jest.fn(); + setup({ addDangerToast }); + await userEvent.click( + await screen.findByRole('button', { name: /postgresql/i }), + ); + expect(await screen.findByText(/step 2 of 3/i)).toBeInTheDocument(); + const fieldByName = (name: string) => + document.querySelector( + `input[name="${name}"]`, + )!; + for (const [name, value] of Object.entries({ + host: 'localhost', + port: '5432', + database: 'postgres', + username: 'testdb', + })) { + fireEvent.change(fieldByName(name), { target: { value } }); + } + fireEvent.blur(fieldByName('username')); + const connectButton = screen.getByRole('button', { + name: 'Connect', + }); + await waitFor(() => expect(connectButton).toBeEnabled()); + fireEvent.click(connectButton); + await waitFor(() => + expect( + fetchMock.callHistory.calls(DATABASE_CONNECT_ENDPOINT), + ).toHaveLength(1), + ); + expect(await screen.findByText(/step 3 of 3/i)).toBeInTheDocument(); + expect(screen.getByText('Database connected')).toBeInTheDocument(); + expect( + screen.queryByText('SQLAlchemy URI'), + ).not.toBeInTheDocument(); + if (action === 'Back') { + await userEvent.click( + screen.getByRole('button', { name: 'Back' }), + ); + await waitFor(() => + expect( + fetchMock.callHistory.calls( + 'glob:*/api/v1/database/10/connection', + ), + ).toHaveLength(1), + ); + await userEvent.click( + screen.getByRole('button', { name: 'Back' }), + ); + expect( + await screen.findByRole('heading', { + name: /select a database to connect/i, + }), + ).toBeInTheDocument(); + await act(async () => { + releaseBackFetch(); + await fetchMock.callHistory.flush(true); + }); + expect( + screen.getByRole('heading', { + name: /select a database to connect/i, + }), + ).toBeInTheDocument(); + return; + } + await userEvent.click( + screen.getByRole('button', { name: 'Finish' }), + ); + await waitFor(() => + expect( + fetchMock.callHistory.calls(DATABASE_FETCH_ENDPOINT), + ).toHaveLength(1), + ); + expect( + JSON.parse( + fetchMock.callHistory.lastCall(DATABASE_FETCH_ENDPOINT)!.options + .body as string, + ), + ).toEqual( + expect.objectContaining({ + configuration_method: ConfigurationMethod.DynamicForm, + }), + ); + } finally { + releaseBackFetch(); + resourceHook.mockImplementation(previousImplementation); + } + }, + ); + test('does not fire redundant validation on blur when db has not changed', async () => { setup(); @@ -2054,6 +2560,7 @@ describe('dbReducer', () => { const currentState = dbReducer({}, action); expect(currentState).toEqual({ ...databaseFixture, + configuration_method: ConfigurationMethod.DynamicForm, engine: 'postgres', masked_encrypted_extra: '', parameters: undefined, @@ -2797,3 +3304,120 @@ describe('dbReducer', () => { }); }); }); + +test.each([true, false])( + 'editing an unavailable driver renders a usable form (saved schema: %s)', + async hasSavedSchema => { + const resource: DatabaseObject = { + ...databaseFixture, + id: 10, + backend: 'databricks', + driver: 'connector', + database_name: 'Legacy Databricks', + sqlalchemy_uri: + 'databricks+connector://token:XXXXXXXXXX@example.cloud.databricks.com:443/default', + parameters: { + host: 'example.cloud.databricks.com', + port: 443, + access_token: 'XXXXXXXXXX', + database: 'default', + }, + extra: JSON.stringify({ + engine_params: { + connect_args: { http_path: '/sql/1.0/warehouses/example' }, + }, + }), + parameters_schema: hasSavedSchema + ? { + properties: { + host: { type: 'string' }, + port: { type: 'integer' }, + access_token: { type: 'string' }, + database: { type: 'string' }, + http_path: { type: 'string' }, + }, + } + : {}, + }; + const updateResource = jest.fn().mockResolvedValue(resource); + const hook = jest.spyOn(hooks, 'useSingleViewResource').mockReturnValue({ + state: { loading: false, resource, error: null }, + fetchResource: jest.fn().mockResolvedValue(resource), + createResource: jest.fn(), + updateResource, + clearError: jest.fn(), + setResource: jest.fn(), + }); + // Matches the SQLAlchemy 2 installation: the old driver has no form in + // /available, but its saved schema and masked connection details survive. + fetchMock.get(AVAILABLE_DB_ENDPOINT, { + databases: [ + { + engine: 'databricks', + name: 'Databricks', + default_driver: 'databricks-sql-python', + available_drivers: ['databricks'], + }, + ], + }); + try { + render(, { + useRedux: true, + }); + if (hasSavedSchema) { + expect( + await screen.findByDisplayValue('example.cloud.databricks.com'), + ).toBeInTheDocument(); + expect( + screen.getByDisplayValue('/sql/1.0/warehouses/example'), + ).toBeInTheDocument(); + const token = screen.getByDisplayValue('XXXXXXXXXX'); + fireEvent.change(token, { + target: { value: 'replacement-test-token' }, + }); + expect(token).toHaveValue('replacement-test-token'); + } else { + const uri = await screen.findByTestId('sqlalchemy-uri-input'); + expect(uri).toHaveValue(resource.sqlalchemy_uri); + fireEvent.change(uri, { + target: { value: 'databricks://example.cloud.databricks.com' }, + }); + expect(uri).toHaveValue('databricks://example.cloud.databricks.com'); + await userEvent.click( + screen.getByRole('button', { name: /^finish$/i }), + ); + expect(updateResource).toHaveBeenCalledWith( + 10, + expect.objectContaining({ + configuration_method: ConfigurationMethod.SqlalchemyUri, + sqlalchemy_uri: 'databricks://example.cloud.databricks.com', + }), + false, + true, + expect.any(Function), + ); + } + } finally { + hook.mockRestore(); + fetchMock.clearHistory().removeRoutes(); + } + }, +); + +test.each(['postgresql', 'mysql', 'databricks', 'gsheets'])( + 'only editing %s falls back to URI when the saved schema is absent', + backend => { + const payload = { + backend, + configuration_method: ConfigurationMethod.DynamicForm, + }; + expect( + dbReducer({}, { type: ActionType.Fetched, payload, isEditMode: false }) + ?.configuration_method, + ).toBe(ConfigurationMethod.DynamicForm); + expect( + dbReducer({}, { type: ActionType.Fetched, payload, isEditMode: true }) + ?.configuration_method, + ).toBe(ConfigurationMethod.SqlalchemyUri); + }, +); diff --git a/superset-frontend/src/features/databases/DatabaseModal/index.tsx b/superset-frontend/src/features/databases/DatabaseModal/index.tsx index c8e41e957d09..1fe4ab87a604 100644 --- a/superset-frontend/src/features/databases/DatabaseModal/index.tsx +++ b/superset-frontend/src/features/databases/DatabaseModal/index.tsx @@ -211,6 +211,7 @@ export type DBReducerActionType = } | { type: ActionType.Fetched; + isEditMode?: boolean; payload: Partial; } | { @@ -572,7 +573,14 @@ export function dbReducer( ...trimmedState, [action.payload.name]: action.payload.value, }; - case ActionType.Fetched: + case ActionType.Fetched: { + const configurationMethod = + action.isEditMode && + action.payload.configuration_method === + ConfigurationMethod.DynamicForm && + !Object.keys(action.payload.parameters_schema?.properties || {}).length + ? ConfigurationMethod.SqlalchemyUri + : action.payload.configuration_method; // convert query to a string and store in query_input query = action.payload?.parameters?.query || {}; // eslint-disable-next-line camelcase @@ -598,7 +606,7 @@ export function dbReducer( return { ...action.payload, engine: action.payload.backend || trimmedState.engine, - configuration_method: action.payload.configuration_method, + configuration_method: configurationMethod, catalog: engineRootCatalog, parameters: { ...(action.payload.parameters || trimmedState.parameters), @@ -616,7 +624,7 @@ export function dbReducer( ...action.payload, masked_encrypted_extra: action.payload.masked_encrypted_extra || '', engine: action.payload.backend || trimmedState.engine, - configuration_method: action.payload.configuration_method, + configuration_method: configurationMethod, parameters: action.payload.parameters || trimmedState.parameters, ssh_tunnel: action.payload.ssh_tunnel || trimmedState.ssh_tunnel, // eslint-disable-next-line camelcase @@ -625,7 +633,7 @@ export function dbReducer( // eslint-disable-next-line camelcase query_input, }; - + } case ActionType.DbSelected: // set initial state for blank form return { @@ -660,7 +668,7 @@ const DatabaseModal: FunctionComponent = ({ >(dbReducer, null); // Database fetch logic const { - state: { loading: dbLoading, resource: dbFetched, error: dbErrors }, + state: { resource: dbFetched, error: dbErrors }, fetchResource, createResource, updateResource, @@ -672,6 +680,19 @@ const DatabaseModal: FunctionComponent = ({ 'connection', ); + const databaseRequestGeneration = useRef(0); + const dbFetchedGeneration = useRef(0); + + /** Bind resource publication and save continuations to the requesting session. */ + const getResponseGuard = () => { + const generation = databaseRequestGeneration.current; + return () => { + const active = show && generation === databaseRequestGeneration.current; + if (active) dbFetchedGeneration.current = generation; + return active; + }; + }; + const [tabKey, setTabKey] = useState(DEFAULT_TAB_KEY); const [availableDbs, getAvailableDbs] = useAvailableDatabases(); const [ @@ -768,7 +789,7 @@ const DatabaseModal: FunctionComponent = ({ const showDBError = validationErrors || dbErrors; const history = useHistory(); - const dbModel: DatabaseForm = + const availableDbModel: DatabaseForm = // TODO: we need a centralized engine in one place // first try to match both engine and driver @@ -787,6 +808,16 @@ const DatabaseModal: FunctionComponent = ({ ) || {}; + // A saved driver's schema can remain usable even when that driver is no longer + // installed. Do not replace it with another driver's form for the same engine. + const dbModel: DatabaseForm = + isEditMode && db?.parameters_schema?.properties + ? { + ...availableDbModel, + parameters: db.parameters_schema as DatabaseForm['parameters'], + } + : availableDbModel; + const handleClearValidationErrors = useCallback(() => { setValidationErrors(null); setHasValidated(false); @@ -904,6 +935,7 @@ const DatabaseModal: FunctionComponent = ({ }, [db, getValidation]); const onClose = () => { + databaseRequestGeneration.current += 1; setDB({ type: ActionType.Reset }); setHasConnectedDb(false); handleClearValidationErrors(); // reset validation errors on close @@ -949,6 +981,7 @@ const DatabaseModal: FunctionComponent = ({ }); const onSave = async () => { + const shouldApplyResponse = getResponseGuard(); let dbConfigExtraExtensionOnSaveError; setLoading(true); setHasValidated(false); @@ -981,6 +1014,7 @@ const DatabaseModal: FunctionComponent = ({ } const errors = await getValidation(dbToUpdate, true); + if (!shouldApplyResponse()) return; // ``getValidation`` returns ``[]`` on success, a field-keyed object // for blocking errors (e.g. the duplicate ``database_name`` check), // and ``null`` for stale or unexpected responses. During save we @@ -1075,7 +1109,10 @@ const DatabaseModal: FunctionComponent = ({ db.id as number, dbToUpdate as DatabaseObject, dbToUpdate.configuration_method === ConfigurationMethod.DynamicForm, // onShow toast on SQLA Forms + true, + shouldApplyResponse, ); + if (!shouldApplyResponse()) return; if (result) { if (onDatabaseAdd) onDatabaseAdd(); dbConfigExtraExtension @@ -1100,7 +1137,9 @@ const DatabaseModal: FunctionComponent = ({ const dbId = await createResource( dbToUpdate as DatabaseObject, dbToUpdate.configuration_method === ConfigurationMethod.DynamicForm, // onShow toast on SQLA Forms + shouldApplyResponse, ); + if (!shouldApplyResponse()) return; if (dbId) { setHasConnectedDb(true); if (onDatabaseAdd) onDatabaseAdd(); @@ -1153,19 +1192,20 @@ const DatabaseModal: FunctionComponent = ({ setLoading(false); }; + /** Fetch only into the modal session that requested the database. */ + const fetchDatabase = (id: number) => fetchResource(id, getResponseGuard()); + // Initialize const fetchDB = () => { if (isEditMode && databaseId) { - if (!dbLoading) { - fetchResource(databaseId).catch(e => - addDangerToast( - t( - 'Sorry there was an error fetching database information: %s', - e.message, - ), + fetchDatabase(databaseId).catch(e => + addDangerToast( + t( + 'Sorry there was an error fetching database information: %s', + e.message, ), - ); - } + ), + ); } }; @@ -1324,13 +1364,14 @@ const DatabaseModal: FunctionComponent = ({ const handleBackButtonOnFinish = () => { if (dbFetched) { - fetchResource(dbFetched.id as number); + fetchDatabase(dbFetched.id as number); } setShowCTAbtns(false); setEditNewDb(true); }; const handleBackButtonOnConnect = () => { + databaseRequestGeneration.current += 1; handleClearValidationErrors(); if (editNewDb) setHasConnectedDb(false); if (importingModal) setImportingModal(false); @@ -1516,18 +1557,30 @@ const DatabaseModal: FunctionComponent = ({ if (databaseId && show) { fetchDB(); } + return () => { + // Closing or switching targets invalidates outstanding connection requests. + databaseRequestGeneration.current += 1; + }; }, [show, databaseId]); useEffect(() => { - if (dbFetched) { + if ( + show && + dbFetched && + dbFetchedGeneration.current === databaseRequestGeneration.current && + (!databaseId || dbFetched.id === databaseId) + ) { setDB({ type: ActionType.Fetched, + isEditMode, payload: dbFetched, }); // keep a copy of the name separate for display purposes // because it shouldn't change when the form is updated setDbName(dbFetched.database_name); } + // Apply resource changes only while open, including multi-step creates. + // Opening the modal must not reapply a cached record after onClose resets it. }, [dbFetched]); useEffect(() => { @@ -1882,7 +1935,7 @@ const DatabaseModal: FunctionComponent = ({ const fetchAndSetDB = () => { setLoading(true); - fetchResource(dbFetched?.id as number).then(r => { + fetchDatabase(dbFetched?.id as number).then(r => { setItem(LocalStorageKeys.Database, r); }); }; diff --git a/superset-frontend/src/views/CRUD/hooks.test.tsx b/superset-frontend/src/views/CRUD/hooks.test.tsx index 191033706749..101594cfcc50 100644 --- a/superset-frontend/src/views/CRUD/hooks.test.tsx +++ b/superset-frontend/src/views/CRUD/hooks.test.tsx @@ -939,6 +939,52 @@ test('useSingleViewResource: fetchResource appends pathSuffix', async () => { }); }); +test('useSingleViewResource: discarded fetch does not replace the active resource', async () => { + const response = deferredJsonResponse(); + jest.spyOn(SupersetClient, 'get').mockReturnValue(response.promise); + const { result } = renderHook(() => + useSingleViewResource('database', 'Database', jest.fn()), + ); + let active = true; + let request: Promise = Promise.resolve(); + act(() => { + result.current.setResource({ id: 20 }); + request = result.current.fetchResource(10, () => active); + }); + active = false; + await act(async () => { + response.resolve({ + json: { result: { id: 10 } }, + } as unknown as JsonResponse); + expect(await request).toEqual({ id: 10 }); + }); + expect(result.current.state.resource).toEqual({ id: 20 }); + expect(result.current.state.loading).toBe(false); +}); + +test('useSingleViewResource: discarded fetch errors do not update state or show a toast', async () => { + const response = deferredJsonResponse(); + jest.spyOn(SupersetClient, 'get').mockReturnValue(response.promise); + const handleError = jest.fn(); + const { result } = renderHook(() => + useSingleViewResource('database', 'Database', handleError), + ); + let active = true; + let request: Promise = Promise.resolve(); + act(() => { + request = result.current.fetchResource(10, () => active); + }); + active = false; + await act(async () => { + response.reject({ error: 'Fetch failed' }); + await request; + }); + expect(result.current.state.resource).toBeNull(); + expect(result.current.state.error).toBeNull(); + expect(result.current.state.loading).toBe(false); + expect(handleError).not.toHaveBeenCalled(); +}); + test('useSingleViewResource: createResource posts to correct endpoint', async () => { const postSpy = jest.spyOn(SupersetClient, 'post').mockResolvedValue({ json: { id: 99, result: { name: 'New Chart' } }, @@ -992,6 +1038,89 @@ test('useSingleViewResource: updateResource puts to correct endpoint', async () }); }); +test.each(['create', 'update'] as const)( + 'useSingleViewResource: discarded %s preserves the active resource and returns its response', + async operation => { + const response = deferredJsonResponse(); + jest + .spyOn(SupersetClient, operation === 'create' ? 'post' : 'put') + .mockReturnValue(response.promise); + const { result } = renderHook(() => + useSingleViewResource<{ id?: number; name?: string }>( + 'database', + 'Database', + jest.fn(), + ), + ); + let active = true; + let request: Promise = Promise.resolve(); + act(() => { + request = + operation === 'create' + ? result.current.createResource({ name: 'Old' }, false, () => active) + : result.current.updateResource( + 77, + { name: 'Old' }, + false, + true, + () => active, + ); + }); + active = false; + act(() => result.current.setResource({ id: 88, name: 'Active' })); + await act(async () => { + response.resolve({ + json: { id: 77, result: { name: 'Old' } }, + } as unknown as JsonResponse); + expect(await request).toEqual( + operation === 'create' ? 77 : { name: 'Old' }, + ); + }); + expect(result.current.state.resource).toEqual({ id: 88, name: 'Active' }); + expect(result.current.state.error).toBeNull(); + }, +); + +test.each(['create', 'update'] as const)( + 'useSingleViewResource: discarded %s errors do not update state or show a toast', + async operation => { + const response = deferredJsonResponse(); + jest + .spyOn(SupersetClient, operation === 'create' ? 'post' : 'put') + .mockReturnValue(response.promise); + const handleError = jest.fn(); + const { result } = renderHook(() => + useSingleViewResource<{ name: string }>( + 'database', + 'Database', + handleError, + ), + ); + let active = true; + let request: Promise = Promise.resolve(); + act(() => { + request = + operation === 'create' + ? result.current.createResource({ name: 'Old' }, false, () => active) + : result.current.updateResource( + 77, + { name: 'Old' }, + false, + true, + () => active, + ); + }); + active = false; + await act(async () => { + response.reject({ error: 'Write failed' }); + await request; + }); + expect(result.current.state.resource).toBeNull(); + expect(result.current.state.error).toBeNull(); + expect(handleError).not.toHaveBeenCalled(); + }, +); + test('useSingleViewResource: clearError resets error to null', async () => { // First make a failing request to get an error state jest.spyOn(SupersetClient, 'get').mockRejectedValue('Network error'); diff --git a/superset-frontend/src/views/CRUD/hooks.ts b/superset-frontend/src/views/CRUD/hooks.ts index addc9aeb3d81..49a16d276181 100644 --- a/superset-frontend/src/views/CRUD/hooks.ts +++ b/superset-frontend/src/views/CRUD/hooks.ts @@ -428,7 +428,7 @@ export function useSingleViewResource< ); const fetchResource = useCallback( - (resourceID: number) => { + (resourceID: number, shouldApplyResponse: () => boolean = () => true) => { // Set loading state updateState({ loading: true, @@ -442,13 +442,16 @@ export function useSingleViewResource< }) .then( ({ json = {} }) => { - updateState({ - resource: json.result, - error: null, - }); + if (shouldApplyResponse()) { + updateState({ + resource: json.result, + error: null, + }); + } return json.result; }, createErrorHandler((errMsg: Record) => { + if (!shouldApplyResponse()) return; handleErrorMsg( t( 'An error occurred while fetching %ss: %s', @@ -470,7 +473,11 @@ export function useSingleViewResource< ); const createResource = useCallback( - (resource: P, hideToast = false) => { + ( + resource: P, + hideToast = false, + shouldApplyResponse: () => boolean = () => true, + ) => { // Set loading state updateState({ loading: true, @@ -483,13 +490,16 @@ export function useSingleViewResource< }) .then( ({ json = {} }) => { - updateState({ - resource: { id: json.id, ...json.result }, - error: null, - }); + if (shouldApplyResponse()) { + updateState({ + resource: { id: json.id, ...json.result }, + error: null, + }); + } return json.id; }, createErrorHandler((errMsg: Record) => { + if (!shouldApplyResponse()) return; // we did not want toasts for db-connection-ui but did not want to disable it everywhere if (!hideToast) { handleErrorMsg( @@ -514,7 +524,13 @@ export function useSingleViewResource< ); const updateResource = useCallback( - (resourceID: number, resource: P, hideToast = false, setLoading = true) => { + ( + resourceID: number, + resource: P, + hideToast = false, + setLoading = true, + shouldApplyResponse: () => boolean = () => true, + ) => { // Set loading state if (setLoading) { updateState({ @@ -529,13 +545,16 @@ export function useSingleViewResource< }) .then( ({ json = {} }) => { - updateState({ - resource: { ...json.result, id: json.id }, - error: null, - }); + if (shouldApplyResponse()) { + updateState({ + resource: { ...json.result, id: json.id }, + error: null, + }); + } return json.result; }, createErrorHandler(errMsg => { + if (!shouldApplyResponse()) return errMsg; if (!hideToast) { handleErrorMsg( t( diff --git a/superset/commands/database/update.py b/superset/commands/database/update.py index a0fe624b38a9..9fd0b099b726 100644 --- a/superset/commands/database/update.py +++ b/superset/commands/database/update.py @@ -305,7 +305,10 @@ def _check_no_unsafe_secret_rebind(self) -> None: if "sqlalchemy_uri" in self._properties: submitted_uri = self._properties["sqlalchemy_uri"] or "" connection_identity_changed = uri_identity_changed( - model.sqlalchemy_uri, submitted_uri + model.sqlalchemy_uri, + submitted_uri, + model.extra, + self._properties.get("extra", model.extra), ) try: submitted_password = make_url_safe(submitted_uri).password diff --git a/superset/commands/database/utils.py b/superset/commands/database/utils.py index ef428c110658..0b6358832b8a 100644 --- a/superset/commands/database/utils.py +++ b/superset/commands/database/utils.py @@ -45,11 +45,33 @@ logger = logging.getLogger(__name__) -def uri_identity_changed(existing_uri: str | None, submitted_uri: str | None) -> bool: +def _connect_args(serialized_extra: str | None) -> dict[str, Any]: + try: + connect_args = ( + json.loads(serialized_extra or "{}") + .get("engine_params", {}) + .get("connect_args", {}) + ) + except (TypeError, ValueError, AttributeError): + return {} + return connect_args if isinstance(connect_args, dict) else {} + + +def uri_identity_changed( + existing_uri: str | None, + submitted_uri: str | None, + existing_extra: str | None = None, + submitted_extra: str | None = None, +) -> bool: """ Whether two SQLAlchemy URIs differ once their password is stripped -- i.e. whether the effective connection destination (driver, host, port, database, username, query params) changed. + + The stored engine spec compares URIs that spell the same destination + differently (e.g. a setting in ``connect_args`` versus the URI query) by + their effective destination. ``engine_params`` changes are checked + separately by ``engine_params_changed``. """ try: stored = make_url_safe(existing_uri or "")._replace(password=None) @@ -58,7 +80,12 @@ def uri_identity_changed(existing_uri: str | None, submitted_uri: str | None) -> # An unparseable URI cannot be compared: treat it as a change so a # stored secret never survives onto it. return True - return stored != incoming + if stored == incoming: + return False + engine_spec = Database.get_db_engine_spec(stored) + return engine_spec.get_connection_identity( + stored, _connect_args(existing_extra) + ) != engine_spec.get_connection_identity(incoming, _connect_args(submitted_extra)) def engine_params_changed( diff --git a/superset/commands/database/validate.py b/superset/commands/database/validate.py index 6edf192d8426..0a4081e96729 100644 --- a/superset/commands/database/validate.py +++ b/superset/commands/database/validate.py @@ -22,6 +22,7 @@ from superset import is_feature_enabled from superset.commands.base import BaseCommand from superset.commands.database.exceptions import ( + DatabaseInvalidError, DatabaseOfflineError, DatabaseTestConnectionFailedError, InvalidEngineError, @@ -30,7 +31,9 @@ from superset.commands.database.utils import ( engine_params_changed, ssh_tunnel_endpoint_changed, + uri_identity_changed, ) +from superset.constants import PASSWORD_MASK from superset.daos.database import DatabaseDAO from superset.databases.utils import make_url_safe from superset.db_engine_specs import get_engine_spec @@ -80,6 +83,11 @@ def run(self) -> None: # noqa: C901 ), ) + if hasattr(engine_spec, "get_extra_from_parameters"): + self._properties["extra"] = engine_spec.get_extra_from_parameters( + self._properties.get("parameters", {}), + self._properties.get("extra", "{}"), + ) # perform initial validation (host, port, database, username) errors = engine_spec.validate_parameters(self._properties) # type: ignore @@ -130,7 +138,26 @@ def run(self) -> None: # noqa: C901 self._properties.get("parameters"), encrypted_extra, ) - if self._model and sqlalchemy_uri == self._model.safe_sqlalchemy_uri(): + if ( + self._model + and self._properties["engine"] == "databricks" + and make_url_safe(sqlalchemy_uri).password == PASSWORD_MASK + and not self._reaches_stored_destination(sqlalchemy_uri) + ): + raise InvalidParametersError( + [ + SupersetError( + message=__( + "Re-enter the access token when changing the connection " + "destination." + ), + error_type=SupersetErrorType.CONNECTION_ACCESS_DENIED_ERROR, + level=ErrorLevel.ERROR, + extra={"invalid": ["access_token"]}, + ) + ] + ) + if self._model and self._reaches_stored_destination(sqlalchemy_uri): if identity_changed: raise InvalidParametersError( [ @@ -217,6 +244,28 @@ def _load_model(self) -> None: if (database_id := self._properties.get("id")) is not None: self._model = DatabaseDAO.find_by_id(database_id) + def _reaches_stored_destination(self, sqlalchemy_uri: str) -> bool: + """ + Whether a URI rebuilt from masked parameters names the stored + connection's destination, even when spelled differently. + """ + model = self._model + assert model is not None + safe_uri = model.safe_sqlalchemy_uri() + if sqlalchemy_uri == safe_uri: + return True + try: + password = make_url_safe(sqlalchemy_uri).password + except DatabaseInvalidError: + return False + # A newly typed password is the one to test, not the stored one. + return password == PASSWORD_MASK and not uri_identity_changed( + safe_uri, + sqlalchemy_uri, + model.extra, + self._properties.get("extra", "{}"), + ) + def _get_database_name_error(self) -> Optional[SupersetError]: """Check for duplicate database name and return error if found.""" database_id = self._properties.get("id") diff --git a/superset/databases/api.py b/superset/databases/api.py index 1fd874d0861f..129571de0303 100644 --- a/superset/databases/api.py +++ b/superset/databases/api.py @@ -125,6 +125,7 @@ ) from superset.databases.utils import get_table_metadata from superset.db_engine_specs import get_available_engine_specs +from superset.db_engine_specs.databricks import DatabricksNativeEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.exceptions import ( DatabaseNotFoundException, @@ -1999,7 +2000,7 @@ def available(self) -> Response: preferred_databases: list[str] = app.config.get("PREFERRED_DATABASES", []) available_databases = [] for engine_spec, drivers in get_available_engine_specs().items(): - if not drivers: + if not drivers or engine_spec is DatabricksNativeEngineSpec: continue payload: dict[str, Any] = { @@ -2255,9 +2256,9 @@ def get_list( # pylint: disable=arguments-differ @safe @statsd_metrics @event_logger.log_this_with_context( - action=lambda self, - *args, - **kwargs: f"{self.__class__.__name__}.delete_oauth2_tokens", + action=lambda self, *args, **kwargs: ( + f"{self.__class__.__name__}.delete_oauth2_tokens" + ), log_to_statsd=False, ) @requires_json @@ -2312,9 +2313,9 @@ def delete_oauth2_tokens(self, pk: int) -> FlaskResponse: @safe @statsd_metrics @event_logger.log_this_with_context( - action=lambda self, - *args, - **kwargs: f"{self.__class__.__name__}.delete_all_oauth2_tokens", + action=lambda self, *args, **kwargs: ( + f"{self.__class__.__name__}.delete_all_oauth2_tokens" + ), log_to_statsd=False, ) def delete_all_oauth2_tokens(self, pk: int) -> FlaskResponse: diff --git a/superset/databases/schemas.py b/superset/databases/schemas.py index 13f4f3fcbf3d..28045e32c379 100644 --- a/superset/databases/schemas.py +++ b/superset/databases/schemas.py @@ -437,6 +437,12 @@ def build_sqlalchemy_uri( except json.JSONDecodeError: encrypted_extra = {} + if hasattr(engine_spec, "get_extra_from_parameters"): + extra = engine_spec.get_extra_from_parameters( + parameters, data.get("extra") + ) + if extra is not None: + data["extra"] = extra data["sqlalchemy_uri"] = engine_spec.build_sqlalchemy_uri( parameters, encrypted_extra, @@ -1012,6 +1018,9 @@ def fix_schemas_allowed_for_csv_upload( # pylint: disable=invalid-name per_user_caching = fields.Boolean(required=False) version = fields.String(required=False, allow_none=True) schema_options = fields.Dict(keys=fields.Str(), values=fields.Raw()) + # Set by the b8f3a2c91d64 migration on converted Databricks connections, so + # exports of those connections carry it. + migrated_from_databricks_connector_uri = fields.Boolean(required=False) class ImportV1DatabaseSchema(Schema): diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index b6aa468293d8..d17f80a88977 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -3217,6 +3217,32 @@ def unmask_encrypted_extra(cls, old: str | None, new: str | None) -> str | None: return json.dumps(new_config) + @classmethod + def get_connection_identity(cls, uri: URL, connect_args: dict[str, Any]) -> URL: + """ + Return the destination a stored credential is sent to, for checks that + decide whether the credential may be reused on an edited connection. + """ + return uri + + @classmethod + def get_extra_from_parameters( + cls, parameters: dict[str, Any], extra: str | None + ) -> str | None: + """Reconcile field-form values with settings stored outside the URI.""" + return extra + + @classmethod + def get_parameters_from_database( + cls, + database: Database, + encrypted_extra: dict[str, Any], + ) -> dict[str, Any]: + """Return masked form parameters, including engine-specific stored settings.""" + return cls.get_parameters_from_uri( # type: ignore + make_url_safe(database.sqlalchemy_uri), encrypted_extra=encrypted_extra + ) + @classmethod def get_public_information(cls) -> dict[str, Any]: """ diff --git a/superset/db_engine_specs/databricks.py b/superset/db_engine_specs/databricks.py index d15797eac698..6e6127794e83 100644 --- a/superset/db_engine_specs/databricks.py +++ b/superset/db_engine_specs/databricks.py @@ -209,7 +209,7 @@ class DatabricksPythonConnectorSchema(DatabricksBaseSchema): """ http_path_field = fields.Str(required=True) - default_catalog = fields.Str(required=True) + default_catalog = fields.Str(required=False) default_schema = fields.Str(required=True) @@ -742,8 +742,11 @@ def validate_parameters( # type: ignore class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): engine = "databricks" engine_name = "Databricks" - default_driver = "databricks-sql-python" - drivers = {"databricks-sql-python": "Databricks SQL Python"} + default_driver = "databricks" + drivers = { + "databricks": "Databricks SQL Python", + "databricks-sql-python": "Databricks SQL Python", + } parameters_schema = DatabricksPythonConnectorSchema() @@ -817,13 +820,15 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): "notes": "Requires ODBC driver. For serverless SQL warehouses.", }, { - "name": "databricks-dbapi (Legacy)", - "pypi_package": "databricks-dbapi[sqlalchemy]", + "name": "Field-based connection", + "pypi_package": "apache-superset[databricks]", "connection_string": ( - "databricks+connector://token:{access_token}@{host}:{port}/{database}" + "databricks://token:{access_token}@{host}:{port}/{database}" ), "is_recommended": False, - "notes": "Legacy connector. Use Python Connector for new deployments.", + "notes": ( + "Use the dynamic form to enter connection fields individually." + ), }, ], } @@ -836,7 +841,6 @@ class DatabricksPythonConnectorEngineSpec(DatabricksDynamicBaseEngineSpec): } required_parameters = DatabricksDynamicBaseEngineSpec.required_parameters | { - "default_catalog", "default_schema", "http_path_field", } @@ -897,18 +901,82 @@ def get_parameters_from_uri( # type: ignore "access_token": url.password, "host": url.host, "port": url.port, - "http_path_field": query["http_path"], - "default_catalog": query["catalog"], - "default_schema": query["schema"], + "http_path_field": query.get("http_path", ""), + "default_catalog": query.get("catalog", ""), + "default_schema": query.get("schema", url.database or ""), "encryption": encryption, } + @classmethod + def get_parameters_from_database( + cls, database: Database, encrypted_extra: dict[str, Any] + ) -> dict[str, Any]: + """Expose legacy connect_args in the native connector's field-based form.""" + parameters = super().get_parameters_from_database(database, encrypted_extra) + query = database.url_object.query + for field, key in ( + ("http_path_field", "http_path"), + ("default_catalog", "catalog"), + ("default_schema", "schema"), + ): + if key not in query and key in database.connect_args: + parameters[field] = database.connect_args[key] + return parameters + + @classmethod + def get_extra_from_parameters( + cls, parameters: dict[str, Any], extra: str | None + ) -> str | None: + """Remove legacy overrides when a native form explicitly clears a field.""" + try: + settings = json.loads(extra or "{}") + connect_args = settings.get("engine_params", {}).get("connect_args", {}) + except (ValueError, AttributeError, TypeError): + return extra + if not isinstance(connect_args, dict): + return extra + changed = False + for field, key in ( + ("http_path_field", "http_path"), + ("default_catalog", "catalog"), + ("default_schema", "schema"), + ): + if field in parameters and not parameters[field] and key in connect_args: + del connect_args[key] + changed = True + return json.dumps(settings) if changed else extra + + @classmethod + def get_connection_identity(cls, uri: URL, connect_args: dict[str, Any]) -> URL: + """ + Describe legacy connect_args and native query settings identically. + + Mirrors ``adjust_engine_params``: query settings take precedence over + connect_args, and the URI path is the schema when neither sets one. + """ + query: dict[str, str | tuple[str, ...]] = { + key: str(connect_args[key]) + for key in ("http_path", "catalog", "schema") + if key in connect_args + } + query.update(uri.query) + if "schema" not in query and uri.database: + query["schema"] = uri.database + return uri._replace(database=None).set(query=query) + @classmethod def get_default_catalog( cls, database: Database, ) -> str | None: - return database.url_object.query.get("catalog") + if catalog := database.url_object.query.get( + "catalog" + ) or database.connect_args.get("catalog"): + return catalog + # Converted rows retain the runtime catalog used by legacy permissions. + if database.get_extra().get("migrated_from_databricks_connector_uri"): + return DatabricksNativeEngineSpec.get_default_catalog(database) + return None @classmethod def get_catalog_names( @@ -932,12 +1000,22 @@ def adjust_engine_params( if schema: uri = uri.update_query_dict({"schema": schema}) + elif ( + "schema" not in uri.query and "schema" not in connect_args and uri.database + ): + # The SQLAlchemy 2 dialect reads schema from the query, not the path. + uri = uri.update_query_dict({"schema": uri.database}) + + # Edits made in the native form take precedence over legacy connect_args. + for key in ("http_path", "catalog", "schema"): + if key in uri.query: + connect_args.pop(key, None) return uri, connect_args class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): - """Legacy Databricks connector using databricks-dbapi.""" + """Compatibility form for legacy Databricks connections.""" engine = "databricks" engine_name = "Databricks (legacy)" @@ -952,9 +1030,7 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): ) metadata = { - "description": ( - "Legacy Databricks connector using the databricks-dbapi driver." - ), + "description": "Field-based configuration for legacy Databricks connections.", "logo": "databricks.png", "homepage_url": "https://www.databricks.com/", "categories": [ @@ -962,7 +1038,7 @@ class DatabricksNativeEngineSpec(DatabricksDynamicBaseEngineSpec): DatabaseCategory.ANALYTICAL_DATABASES, DatabaseCategory.PROPRIETARY, ], - "pypi_packages": ["databricks-dbapi[sqlalchemy]"], + "pypi_packages": ["apache-superset[databricks]"], "connection_string": ( "databricks+connector://token:{access_token}@{host}:{port}/{database_name}?http_path={http_path}" ), @@ -1021,7 +1097,7 @@ def build_sqlalchemy_uri( # type: ignore # False) is required here since this URI is stored/used to actually # connect, not just displayed. return URL.create( - f"{cls.engine}+{cls.default_driver}".rstrip("+"), + f"{cls.engine}+{cls.default_driver}", username="token", password=parameters.get("access_token"), host=parameters["host"], diff --git a/superset/migrations/versions/2026-10-29_00-00_b8f3a2c91d64_migrate_databricks_connector_uri.py b/superset/migrations/versions/2026-10-29_00-00_b8f3a2c91d64_migrate_databricks_connector_uri.py new file mode 100644 index 000000000000..446aab2d758d --- /dev/null +++ b/superset/migrations/versions/2026-10-29_00-00_b8f3a2c91d64_migrate_databricks_connector_uri.py @@ -0,0 +1,184 @@ +# 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. + +"""Use the SQLAlchemy 2 Databricks dialect for saved legacy connections. + +Revision ID: b8f3a2c91d64 +Revises: f72b6e9c104a + +Only the URI prefix is rewritten. Credentials live in the separate encrypted +password column, which this migration never reads or writes. The extra marker +and a migration-owned fingerprint table limit rollback to unchanged converted rows. +""" + +from datetime import datetime +from hashlib import sha256 +from typing import Any + +import sqlalchemy as sa +from alembic import op + +from superset.migrations.shared.utils import create_table, drop_table, has_table +from superset.utils import json + +revision = "b8f3a2c91d64" +down_revision = "f72b6e9c104a" + +LEGACY_PREFIX = "databricks+connector://" +NATIVE_PREFIX = "databricks://" +MARKER = "migrated_from_databricks_connector_uri" +BATCH_SIZE = 1000 + +# Deliberately independent of application models and encrypted column types. +databases = sa.table( + "dbs", + sa.column("id", sa.Integer), + sa.column("sqlalchemy_uri", sa.String(1024)), + sa.column("extra", sa.Text), + sa.column("changed_on", sa.DateTime), +) + + +# Fingerprints avoid copying connection settings or credentials into a backup table. +records = sa.table( + "databricks_uri_migration", + sa.column("id", sa.Integer), + sa.column("fingerprint", sa.String(64)), + sa.column("extra_marked", sa.Boolean), +) + + +def _load_extra(value: str | None) -> dict[str, Any] | None: + """Return no object for malformed settings, without replacing their text.""" + try: + extra = json.loads(value) if value is not None else None + except (ValueError, RecursionError): + return None + return extra if isinstance(extra, dict) else None + + +def _fingerprint(uri: str, extra: str | None, changed_on: datetime | None) -> str: + """Recognize the migrated state, including edits to other database fields.""" + state = json.dumps([uri, extra, str(changed_on)]) + return sha256(state.encode("utf-8")).hexdigest() + + +def _query(upgrading: bool, tracked: bool, last_id: int) -> sa.sql.Select: + """Find source URIs or rollback candidates, including exported markers.""" + query = sa.select(databases).where(databases.c.id > last_id) + if upgrading: + query = query.where(databases.c.sqlalchemy_uri.like(f"{LEGACY_PREFIX}%")) + elif tracked: + query = ( + query.add_columns(records.c.fingerprint, records.c.extra_marked) + .select_from(databases.outerjoin(records, databases.c.id == records.c.id)) + .where( + sa.or_(records.c.id.isnot(None), databases.c.extra.like(f"%{MARKER}%")) + ) + ) + else: + query = query.add_columns( + sa.literal(None).label("fingerprint"), + sa.literal(None).label("extra_marked"), + ).where(databases.c.extra.like(f"%{MARKER}%")) + return query + + +def _migrate(upgrading: bool) -> None: + """Use keyset batches so rewriting the filter cannot skip later rows.""" + bind = op.get_bind() + tracked = has_table(records.name) + last_id = -1 + while True: + query = _query(upgrading, tracked, last_id) + rows = bind.execute(query.order_by(databases.c.id).limit(BATCH_SIZE)).fetchall() + if not rows: + break + for row in rows: + last_id = row.id + extra = _load_extra(row.extra) + uri = row.sqlalchemy_uri + extra_text = row.extra + if upgrading: + # LIKE can be case-insensitive on some metadata databases. + if not uri.startswith(LEGACY_PREFIX): + continue + uri = NATIVE_PREFIX + uri[len(LEGACY_PREFIX) :] + extra_marked = extra is not None and ( + MARKER not in extra or extra[MARKER] is True + ) + if extra is not None and extra_marked: + extra[MARKER] = True + extra_text = json.dumps(extra) + + else: + unchanged = row.fingerprint == _fingerprint( + uri, extra_text, row.changed_on + ) + imported = ( + row.fingerprint is None + and uri.startswith(NATIVE_PREFIX) + and extra is not None + and extra.get(MARKER) is True + ) + if (unchanged or imported) and uri.startswith(NATIVE_PREFIX): + uri = LEGACY_PREFIX + uri[len(NATIVE_PREFIX) :] + if ( + (row.extra_marked or imported) + and extra is not None + and extra.get(MARKER) is True + ): + del extra[MARKER] + extra_text = json.dumps(extra) + # Only replace the snapshot that this batch actually read. + updated = bind.execute( + databases.update() + .where( + databases.c.id == row.id, + databases.c.sqlalchemy_uri == row.sqlalchemy_uri, + databases.c.extra == row.extra, + databases.c.changed_on == row.changed_on, + ) + .values(sqlalchemy_uri=uri, extra=extra_text) + ) + if upgrading and updated.rowcount: + # Preserve conflicting user settings instead of aborting the upgrade. + bind.execute(records.delete().where(records.c.id == row.id)) + bind.execute( + records.insert().values( + id=row.id, + fingerprint=_fingerprint(uri, extra_text, row.changed_on), + extra_marked=extra_marked, + ) + ) + + +def upgrade() -> None: + """Convert legacy schemes and track the resulting rows for safe rollback.""" + create_table( + records.name, + sa.Column("id", sa.Integer, primary_key=True, autoincrement=False), + sa.Column("fingerprint", sa.String(64), nullable=False), + sa.Column("extra_marked", sa.Boolean, nullable=False), + ) + _migrate(upgrading=True) + + +def downgrade() -> None: + """Restore unchanged converted URIs, then remove migration bookkeeping.""" + _migrate(upgrading=False) + drop_table(records.name) diff --git a/superset/models/core.py b/superset/models/core.py index c6a34ab2f45e..da20575ed392 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -423,7 +423,11 @@ def backend(self) -> str: @property def driver(self) -> str: - return self.url_object.get_driver_name() + try: + return self.url_object.get_driver_name() + except NoSuchModuleError: + # Driverless connections must still open in the URI repair form. + return "" @property def masked_encrypted_extra(self) -> str | None: @@ -435,14 +439,14 @@ def parameters(self) -> dict[str, Any]: # the sqlalchemy_uri # When returning the parameters we should use the masked SQLAlchemy URI and the # masked ``encrypted_extra`` to prevent exposing sensitive credentials. - masked_uri = make_url_safe(self.sqlalchemy_uri) - encrypted_config = {} - if (masked_encrypted_extra := self.masked_encrypted_extra) is not None: - with suppress(TypeError, json.JSONDecodeError): - encrypted_config = json.loads(masked_encrypted_extra) try: - parameters = self.db_engine_spec.get_parameters_from_uri( # type: ignore - masked_uri, + make_url_safe(self.sqlalchemy_uri) + encrypted_config = {} + if (masked_encrypted_extra := self.masked_encrypted_extra) is not None: + with suppress(TypeError, json.JSONDecodeError): + encrypted_config = json.loads(masked_encrypted_extra) + parameters = self.db_engine_spec.get_parameters_from_database( + self, encrypted_extra=encrypted_config, ) except Exception: # pylint: disable=broad-except diff --git a/superset/translations/messages.pot b/superset/translations/messages.pot index 5f53aa5bc053..9c31621dcaf8 100644 --- a/superset/translations/messages.pot +++ b/superset/translations/messages.pot @@ -11096,6 +11096,9 @@ msgstr "" msgid "Raw records" msgstr "" +msgid "Re-enter the access token when changing the connection destination." +msgstr "" + msgid "Re-pick fields" msgstr "" diff --git a/tests/unit_tests/commands/databases/update_test.py b/tests/unit_tests/commands/databases/update_test.py index 34d04bc87622..e2f2a607baed 100644 --- a/tests/unit_tests/commands/databases/update_test.py +++ b/tests/unit_tests/commands/databases/update_test.py @@ -1285,3 +1285,64 @@ def test_update_oauth2_unresolvable_endpoints_purges_tokens( command._handle_oauth2() purge_oauth2_tokens.assert_called_once() + + +CONVERTED_DATABRICKS_URI = "databricks://token:XXXXXXXXXX@host:443/default" +CONVERTED_DATABRICKS_EXTRA = json.dumps( + {"engine_params": {"connect_args": {"http_path": "/legacy", "catalog": "main"}}} +) + + +def _converted_databricks_update(mocker: MockerFixture, submitted_uri: str) -> Any: + old_database = mocker.MagicMock(allow_multi_catalog=False) + old_database.sqlalchemy_uri = CONVERTED_DATABRICKS_URI + old_database.extra = CONVERTED_DATABRICKS_EXTRA + old_database.encrypted_extra = None + old_database.ssh_tunnel = None + old_database.get_default_catalog.return_value = "main" + old_database.id = 1 + + new_database = mocker.MagicMock(allow_multi_catalog=False) + new_database.get_default_catalog.return_value = "main" + + database_dao = mocker.patch("superset.commands.database.update.DatabaseDAO") + database_dao.find_by_id.return_value = old_database + database_dao.update.return_value = new_database + + mocker.patch("superset.commands.database.update.SyncPermissionsCommand") + mocker.patch.object(UpdateDatabaseCommand, "_update_catalog_attribute") + + UpdateDatabaseCommand( + 1, + {"sqlalchemy_uri": submitted_uri, "extra": CONVERTED_DATABRICKS_EXTRA}, + ).run() + return database_dao + + +def test_update_converted_databricks_connection_keeps_masked_token( + mocker: MockerFixture, +) -> None: + """ + Saving a converted Databricks connection from the native form rebuilds its + legacy settings into the URI query. The destination is unchanged, so the + masked token must be accepted rather than treated as a connection move. + """ + database_dao = _converted_databricks_update( + mocker, + "databricks://token:XXXXXXXXXX@host:443" + "?catalog=main&http_path=%2Flegacy&schema=default", + ) + + database_dao.update.assert_called_once() + + +def test_update_converted_databricks_http_path_change_requires_new_token( + mocker: MockerFixture, +) -> None: + """A different warehouse is still a destination change for a masked token.""" + with pytest.raises(DatabaseInvalidError): + _converted_databricks_update( + mocker, + "databricks://token:XXXXXXXXXX@host:443" + "?catalog=main&http_path=%2Fother&schema=default", + ) diff --git a/tests/unit_tests/commands/databases/validate_test.py b/tests/unit_tests/commands/databases/validate_test.py index 85d4cf9268ff..12ad416aaa2c 100644 --- a/tests/unit_tests/commands/databases/validate_test.py +++ b/tests/unit_tests/commands/databases/validate_test.py @@ -222,6 +222,7 @@ def test_command_duplicate_database_name(mocker: MockerFixture) -> None: "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, ), ) @@ -259,6 +260,7 @@ def test_command_duplicate_database_name_on_update(mocker: MockerFixture) -> Non "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, ), ) @@ -339,6 +341,7 @@ def test_validate_ssh_tunnel_missing_db_port(mocker: MockerFixture) -> None: "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, ), ) @@ -373,6 +376,7 @@ def test_get_ssh_tunnel_errors_missing_required_fields( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, ), ) @@ -429,6 +433,7 @@ def test_get_ssh_tunnel_errors_unencrypted_private_key_is_valid( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, build_sqlalchemy_uri=mocker.MagicMock(return_value="postgresql://"), unmask_encrypted_extra=mocker.MagicMock(return_value="{}"), ), @@ -479,6 +484,7 @@ def test_ssh_tunnel_forwarded_to_connection_test( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, build_sqlalchemy_uri=mocker.MagicMock(return_value="postgresql://"), unmask_encrypted_extra=mocker.MagicMock(return_value="{}"), ), @@ -552,6 +558,7 @@ def test_ssh_tunnel_unmasked_when_only_engine_params_changed( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, # a fresh (non-masked) password, as if the user supplied the # real one for the main connection -- this is the legitimate # path that gets past the main URI-identity check below @@ -615,6 +622,7 @@ def test_get_ssh_tunnel_errors_skipped_when_parameters_ssh_false( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, build_sqlalchemy_uri=mocker.MagicMock(return_value="postgresql://"), unmask_encrypted_extra=mocker.MagicMock(return_value="{}"), ), @@ -658,6 +666,7 @@ def test_get_ssh_tunnel_errors_skipped_when_not_enabled( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, build_sqlalchemy_uri=mocker.MagicMock(return_value="postgresql://"), unmask_encrypted_extra=mocker.MagicMock(return_value="{}"), ), @@ -758,6 +767,7 @@ def test_validate_engine_params_change_requires_new_credentials( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, build_sqlalchemy_uri=mocker.MagicMock( return_value="postgresql://u:XXXXXXXXXX@host1/d" ), @@ -823,6 +833,7 @@ def test_validate_ssh_tunnel_host_change_requires_new_credentials( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, build_sqlalchemy_uri=mocker.MagicMock( return_value="postgresql://u:XXXXXXXXXX@host1/d" ), @@ -868,6 +879,7 @@ def test_ssh_tunnel_missing_message_is_interpolated( "superset.commands.database.validate.get_engine_spec", return_value=mocker.MagicMock( validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, ), ) @@ -894,3 +906,208 @@ def test_ssh_tunnel_missing_message_is_interpolated( assert missing_field_messages assert all("%(missing)s" not in msg for msg in missing_field_messages) assert any("server_port" in msg for msg in missing_field_messages) + + +@pytest.mark.parametrize( + "rebuilt_uri, expected_uri", + [ + # unchanged native-form fields: the stored token is rehydrated + ( + "databricks://token:XXXXXXXXXX@host:443" + "?catalog=main&http_path=%2Flegacy&schema=default", + "databricks://token:real-token@host:443/default", + ), + # a different warehouse requires a freshly entered token + ( + "databricks://token:XXXXXXXXXX@host:443" + "?catalog=main&http_path=%2Fother&schema=default", + "databricks://token:XXXXXXXXXX@host:443" + "?catalog=main&http_path=%2Fother&schema=default", + ), + # a newly typed token is used as typed + ( + "databricks://token:new-token@host:443" + "?catalog=main&http_path=%2Flegacy&schema=default", + "databricks://token:new-token@host:443" + "?catalog=main&http_path=%2Flegacy&schema=default", + ), + ], +) +def test_validate_converted_databricks_connection_token( + mocker: MockerFixture, rebuilt_uri: str, expected_uri: str +) -> None: + """ + A converted Databricks connection keeps legacy settings in ``connect_args`` + while the native form rebuilds them into the URI query. Validation must + still test an unchanged connection with its stored token. + """ + extra = json.dumps( + {"engine_params": {"connect_args": {"http_path": "/legacy", "catalog": "main"}}} + ) + existing = mocker.MagicMock() + existing.safe_sqlalchemy_uri.return_value = ( + "databricks://token:XXXXXXXXXX@host:443/default" + ) + existing.sqlalchemy_uri_decrypted = "databricks://token:real-token@host:443/default" + existing.encrypted_extra = "{}" + existing.extra = extra + existing.ssh_tunnel = None + + database = mocker.MagicMock() + with database.get_sqla_engine() as engine: + engine.dialect.do_ping.return_value = True + + DatabaseDAO = mocker.patch( # noqa: N806 + "superset.commands.database.validate.DatabaseDAO" + ) + DatabaseDAO.find_by_id.return_value = existing + DatabaseDAO.validate_update_uniqueness.return_value = True + DatabaseDAO.build_db_for_connection_test.return_value = database + + mocker.patch( + "superset.commands.database.validate.get_engine_spec", + return_value=mocker.MagicMock( + validate_parameters=mocker.MagicMock(return_value=[]), + get_extra_from_parameters=lambda parameters, extra: extra, + build_sqlalchemy_uri=mocker.MagicMock(return_value=rebuilt_uri), + unmask_encrypted_extra=mocker.MagicMock(return_value="{}"), + ), + ) + + properties = { + "id": 1, + "engine": "databricks", + "driver": "databricks", + "parameters": {"host": "host", "port": 443}, + "extra": extra, + } + if "XXXXXXXXXX" in expected_uri: + with pytest.raises(InvalidParametersError, match="Re-enter the access token"): + ValidateDatabaseParametersCommand(properties).run() + database.set_sqlalchemy_uri.assert_not_called() + else: + ValidateDatabaseParametersCommand(properties).run() + database.set_sqlalchemy_uri.assert_called_once_with(expected_uri) + + +def test_validate_converted_databricks_without_catalog(mocker: MockerFixture) -> None: + """An unchanged converted row uses real field validation and its stored token.""" + from superset.models.core import Database + + extra = json.dumps( + { + "migrated_from_databricks_connector_uri": True, + "engine_params": {"connect_args": {"http_path": "/legacy"}}, + } + ) + existing = Database(extra=extra) + existing.set_sqlalchemy_uri("databricks://token:test-token@host:443/default") + mocker.patch( + "superset.commands.database.validate.DatabaseDAO.find_by_id", + return_value=existing, + ) + mocker.patch( + "superset.db_engine_specs.databricks.is_hostname_valid", return_value=True + ) + mocker.patch("superset.db_engine_specs.databricks.is_port_open", return_value=True) + database = mocker.MagicMock() + with database.get_sqla_engine() as engine: + engine.dialect.do_ping.return_value = True + mocker.patch( + "superset.commands.database.validate.DatabaseDAO.build_db_for_connection_test", + return_value=database, + ) + ValidateDatabaseParametersCommand( + { + "id": 1, + "engine": "databricks", + "driver": "databricks-sql-python", + "parameters": existing.parameters, + "extra": extra, + } + ).run() + database.set_sqlalchemy_uri.assert_called_once_with( + existing.sqlalchemy_uri_decrypted + ) + + +@pytest.mark.parametrize( + "field, value", + [("host", "other-host"), ("port", 444), ("http_path_field", "/other")], +) +def test_validate_databricks_destination_change_requires_fresh_token( + mocker: MockerFixture, + field: str, + value: str | int, +) -> None: + """Never send a password mask to a changed Databricks destination.""" + from superset.models.core import Database + + existing = Database( + extra='{"engine_params":{"connect_args":{"http_path":"/legacy"}}}' + ) + existing.set_sqlalchemy_uri("databricks://token:test-token@host:443/default") + mocker.patch( + "superset.commands.database.validate.DatabaseDAO.find_by_id", + return_value=existing, + ) + mocker.patch( + "superset.db_engine_specs.databricks.is_hostname_valid", return_value=True + ) + mocker.patch("superset.db_engine_specs.databricks.is_port_open", return_value=True) + build = mocker.patch( + "superset.commands.database.validate.DatabaseDAO.build_db_for_connection_test" + ) + parameters = {**existing.parameters, field: value} + with pytest.raises(InvalidParametersError, match="Re-enter the access token"): + ValidateDatabaseParametersCommand( + { + "id": 1, + "engine": "databricks", + "driver": "databricks-sql-python", + "parameters": parameters, + "extra": existing.extra, + } + ).run() + build.assert_not_called() + + +def test_validate_legacy_databricks_form_keeps_masked_token( + mocker: MockerFixture, +) -> None: + """An unconverted saved form preserves its driver and stored credential.""" + from superset.models.core import Database + + existing = Database( + extra='{"engine_params":{"connect_args":{"http_path":"/legacy"}}}' + ) + existing.set_sqlalchemy_uri( + "databricks+connector://token:test-token@host:443/default" + ) + mocker.patch( + "superset.commands.database.validate.DatabaseDAO.find_by_id", + return_value=existing, + ) + mocker.patch( + "superset.db_engine_specs.databricks.is_hostname_valid", return_value=True + ) + mocker.patch("superset.db_engine_specs.databricks.is_port_open", return_value=True) + database = mocker.MagicMock() + with database.get_sqla_engine() as engine: + engine.dialect.do_ping.return_value = True + mocker.patch( + "superset.commands.database.validate.DatabaseDAO.build_db_for_connection_test", + return_value=database, + ) + ValidateDatabaseParametersCommand( + { + "id": 1, + "engine": "databricks", + "driver": "connector", + "parameters": existing.parameters, + "extra": existing.extra, + } + ).run() + database.set_sqlalchemy_uri.assert_called_once_with( + existing.sqlalchemy_uri_decrypted + ) diff --git a/tests/unit_tests/databases/api_test.py b/tests/unit_tests/databases/api_test.py index b5e7b488cd66..e61ce759855c 100644 --- a/tests/unit_tests/databases/api_test.py +++ b/tests/unit_tests/databases/api_test.py @@ -3044,3 +3044,26 @@ def test_related_objects_limits_dataset_details( assert payload["count"] == MAX_RELATED_DATASETS + 2 assert len(payload["result"]) == MAX_RELATED_DATASETS assert can_access.call_count == MAX_RELATED_DATASETS + + +def test_available_hides_legacy_databricks( + mocker: MockerFixture, + client: FlaskClient, + full_api_access: None, +) -> None: + """New connections offer the supported Databricks form, not the legacy one.""" + from superset.db_engine_specs.databricks import ( + DatabricksNativeEngineSpec, + DatabricksPythonConnectorEngineSpec, + ) + + mocker.patch( + "superset.databases.api.get_available_engine_specs", + return_value={ + DatabricksNativeEngineSpec: {"connector"}, + DatabricksPythonConnectorEngineSpec: {"databricks-sql-python"}, + }, + ) + response = client.get("/api/v1/database/available/") + assert response.status_code == 200 + assert [item["name"] for item in response.json["databases"]] == ["Databricks"] diff --git a/tests/unit_tests/databases/schema_tests.py b/tests/unit_tests/databases/schema_tests.py index 3024654f52c7..876bc56971c9 100644 --- a/tests/unit_tests/databases/schema_tests.py +++ b/tests/unit_tests/databases/schema_tests.py @@ -800,3 +800,39 @@ def test_validate_parameters_schema_accepts_server_host_key() -> None: loaded["ssh_tunnel"]["server_host_key"] == "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAA" ) + + +@pytest.mark.parametrize("schema_name", ["DatabasePostSchema", "DatabasePutSchema"]) +def test_databricks_form_clear_catalog_removes_legacy_override( + schema_name: str, +) -> None: + """Both create and update normalize cleared fields before persisting extra.""" + from superset.databases import schemas + + schema = getattr(schemas, schema_name)() + payload = schema.load( + { + "database_name": "Databricks", + "configuration_method": "dynamic_form", + "engine": "databricks", + "driver": "databricks-sql-python", + "parameters": { + "host": "host", + "port": 443, + "access_token": "test-token", + "default_schema": "default", + "http_path_field": "/legacy", + "default_catalog": "", + }, + "extra": json.dumps( + { + "engine_params": { + "connect_args": {"catalog": "old", "http_path": "/legacy"} + } + } + ), + } + ) + args = json.loads(payload["extra"])["engine_params"]["connect_args"] + assert "catalog" not in args + assert args["http_path"] == "/legacy" diff --git a/tests/unit_tests/db_engine_specs/test_databricks.py b/tests/unit_tests/db_engine_specs/test_databricks.py index af8b50def4f3..3dc75f29089e 100644 --- a/tests/unit_tests/db_engine_specs/test_databricks.py +++ b/tests/unit_tests/db_engine_specs/test_databricks.py @@ -19,7 +19,7 @@ from datetime import datetime from importlib import import_module from importlib.metadata import PackageNotFoundError, requires -from typing import Any, Optional +from typing import Any, cast, Literal, Optional from unittest.mock import MagicMock from urllib.parse import parse_qs, urlparse @@ -43,6 +43,7 @@ DatabricksHiveEngineSpec, DatabricksNativeEngineSpec, DatabricksPythonConnectorEngineSpec, + DatabricksPythonConnectorParametersType, ) from superset.db_engine_specs.spark import SparkEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType @@ -136,8 +137,8 @@ def test_build_sqlalchemy_uri() -> None: sqlalchemy_uri = DatabricksNativeEngineSpec.build_sqlalchemy_uri( parameters, encrypted_extra ) - assert sqlalchemy_uri == ( - "databricks+connector://token:abc12345@my_hostname:1234/test" + assert ( + sqlalchemy_uri == "databricks+connector://token:abc12345@my_hostname:1234/test" ) @@ -1646,3 +1647,202 @@ def test_databricks_hive_spec_shares_extended_aggregations() -> None: for aggregate in ("MEDIAN", "STDDEV_SAMP", "VAR_SAMP"): assert DatabricksHiveEngineSpec.get_extended_aggregation_func(aggregate) + + +def test_native_connector_advertises_installed_driver() -> None: + """The SQLAlchemy 2 dialect must expose a field-based connection form.""" + assert DatabricksPythonConnectorEngineSpec.default_driver == "databricks" + assert "databricks" in DatabricksPythonConnectorEngineSpec.drivers + + +def test_converted_connection_parameters_include_legacy_settings() -> None: + """The native form retains masked credentials and settings stored in extra.""" + from superset.models.core import Database + + database = Database( + sqlalchemy_uri="databricks://token:XXXXXXXXXX@example.cloud.databricks.com:443/default", + extra=json.dumps( + { + "engine_params": { + "connect_args": { + "http_path": "/sql/1.0/warehouses/example", + "catalog": "main", + } + } + } + ), + ) + assert database.parameters == { + "access_token": "XXXXXXXXXX", + "host": "example.cloud.databricks.com", + "port": 443, + "http_path_field": "/sql/1.0/warehouses/example", + "default_catalog": "main", + "default_schema": "default", + "encryption": False, + } + assert DatabricksPythonConnectorEngineSpec.get_default_catalog(database) == "main" + + +def test_native_connector_preserves_legacy_schema_and_prioritizes_form_edits() -> None: + """Path schemas reach the dialect; query fields supersede legacy extra.""" + spec = DatabricksPythonConnectorEngineSpec + uri, args = spec.adjust_engine_params( + make_url("databricks://token:masked@host:443/default"), + {"http_path": "/legacy", "catalog": "main"}, + ) + assert uri.query["schema"] == "default" + assert args == {"http_path": "/legacy", "catalog": "main"} + uri, args = spec.adjust_engine_params( + make_url("databricks://host?http_path=/edited&catalog=new&schema=edited"), + {"http_path": "/legacy", "catalog": "main", "schema": "old"}, + ) + assert uri.query == {"http_path": "/edited", "catalog": "new", "schema": "edited"} + assert args == {} + + +def test_native_form_save_preserves_masked_token_and_updates_legacy_settings() -> None: + """Saving masked native fields keeps the token while accepting new settings.""" + from superset.models.core import Database + + database = Database( + extra=json.dumps( + { + "engine_params": { + "connect_args": {"http_path": "/legacy", "catalog": "main"} + } + } + ) + ) + database.set_sqlalchemy_uri("databricks://token:form-test-token@host:443/default") + parameters = database.parameters + parameters["http_path_field"] = "/edited" + parameters["default_catalog"] = "updated" + uri = DatabricksPythonConnectorEngineSpec.build_sqlalchemy_uri( + cast(DatabricksPythonConnectorParametersType, parameters) + ) + database.set_sqlalchemy_uri(uri) + assert database.url_object.password == "form-test-token" # noqa: S105 + assert database.parameters["http_path_field"] == "/edited" + assert database.parameters["default_catalog"] == "updated" + _, args = DatabricksPythonConnectorEngineSpec.adjust_engine_params( + database.url_object, database.connect_args + ) + assert "http_path" not in args + assert "catalog" not in args + + +def test_connection_identity_matches_legacy_settings_and_native_query() -> None: + """Legacy connect_args and native query settings name the same destination.""" + legacy = DatabricksPythonConnectorEngineSpec.get_connection_identity( + make_url("databricks://token@host:443/default"), + {"http_path": "/legacy", "catalog": "main"}, + ) + native = DatabricksPythonConnectorEngineSpec.get_connection_identity( + make_url( + "databricks://token@host:443" + "?catalog=main&http_path=%2Flegacy&schema=default" + ), + {"http_path": "/legacy", "catalog": "main"}, + ) + other_warehouse = DatabricksPythonConnectorEngineSpec.get_connection_identity( + make_url( + "databricks://token@host:443?catalog=main&http_path=%2Fother&schema=default" + ), + {"http_path": "/legacy", "catalog": "main"}, + ) + + assert legacy == native + assert native != other_warehouse + assert BaseEngineSpec.get_connection_identity( + make_url("postgresql://u@host/db"), {"host": "other"} + ) == make_url("postgresql://u@host/db") + + +@pytest.mark.parametrize("catalogs", [[("hive_metastore",)], [("main",), ("other",)]]) +def test_migrated_default_catalog_preserves_legacy_runtime_default( + mocker: MockerFixture, catalogs: list[tuple[str]] +) -> None: + """Converted rows keep the catalog used by their dataset permissions.""" + from superset.models.core import Database + + database = Database( + sqlalchemy_uri="databricks://token:XXXXXXXXXX@host:443/default", + extra=json.dumps( + { + "migrated_from_databricks_connector_uri": True, + "engine_params": {"connect_args": {"http_path": "/legacy"}}, + } + ), + ) + engine = mocker.MagicMock() + connection = ( + engine.__enter__.return_value.connect.return_value.__enter__.return_value + ) + result = mocker.MagicMock() + result.scalar.return_value = "hive_metastore" + connection.execute.side_effect = [catalogs, result] + mocker.patch.object(database, "get_sqla_engine", return_value=engine) + assert ( + DatabricksPythonConnectorEngineSpec.get_default_catalog(database) + == "hive_metastore" + ) + assert connection.execute.call_count == (1 if len(catalogs) == 1 else 2) + database.extra = "{}" + assert DatabricksPythonConnectorEngineSpec.get_default_catalog(database) is None + + +def test_native_catalog_is_optional_in_form_schema() -> None: + """The browser must not require a catalog that legacy forms never collected.""" + schema = DatabricksPythonConnectorEngineSpec.parameters_json_schema() + assert "default_catalog" not in schema["required"] + + +@pytest.mark.parametrize( + "field, key", + [ + ("default_catalog", "catalog"), + ("default_schema", "schema"), + ("http_path_field", "http_path"), + ], +) +def test_native_form_clear_overrides_legacy_connect_args( + field: Literal["default_catalog", "default_schema", "http_path_field"], key: str +) -> None: + """Clearing a field removes its old connect_args value at connection time.""" + parameters = DatabricksPythonConnectorParametersType( + host="host", + port=443, + access_token="test-token", # noqa: S106 + default_catalog="main", + default_schema="default", + http_path_field="/legacy", + encryption=False, + ) + parameters[field] = "" + uri = make_url(DatabricksPythonConnectorEngineSpec.build_sqlalchemy_uri(parameters)) + extra = DatabricksPythonConnectorEngineSpec.get_extra_from_parameters( + dict(parameters), + json.dumps({"engine_params": {"connect_args": {key: "old", "other": "kept"}}}), + ) + assert extra is not None + adjusted, args = DatabricksPythonConnectorEngineSpec.adjust_engine_params( + uri, json.loads(extra)["engine_params"]["connect_args"] + ) + assert args["other"] == "kept" + assert key not in args + assert key not in adjusted.query + + +def test_legacy_form_preserves_connector_scheme() -> None: + """Unconverted forms must not silently change the credential destination.""" + from superset.commands.database.utils import uri_identity_changed + + parameters = DatabricksNativeEngineSpec.get_parameters_from_uri( + "databricks+connector://token:XXXXXXXXXX@host:443/default" + ) + rebuilt = DatabricksNativeEngineSpec.build_sqlalchemy_uri(parameters) + assert rebuilt.startswith("databricks+connector://") + assert not uri_identity_changed( + "databricks+connector://token:XXXXXXXXXX@host:443/default", rebuilt + ) diff --git a/tests/unit_tests/migrations/test_migrate_databricks_connector_uri.py b/tests/unit_tests/migrations/test_migrate_databricks_connector_uri.py new file mode 100644 index 000000000000..79b047a834ee --- /dev/null +++ b/tests/unit_tests/migrations/test_migrate_databricks_connector_uri.py @@ -0,0 +1,417 @@ +# 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. +"""Round-trip coverage for saved legacy Databricks connection URIs.""" + +from collections.abc import Iterator +from importlib import import_module +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +import sqlalchemy as sa +from alembic.migration import MigrationContext +from alembic.operations import Operations +from sqlalchemy.engine import Engine + +from superset.utils import json + +migration = import_module( + "superset.migrations.versions." + "2026-10-29_00-00_b8f3a2c91d64_migrate_databricks_connector_uri" +) +LEGACY = migration.LEGACY_PREFIX +NATIVE = migration.NATIVE_PREFIX +MASKED_SUFFIX = "token:XXXXXXXXXX@example.cloud.databricks.com:443/default" + + +@pytest.fixture +def engine() -> Iterator[Engine]: + """Use actual encrypted column types but only the migration's local table.""" + from superset.models.core import Database + + engine = sa.create_engine("sqlite://") + metadata = sa.MetaData() + sa.Table( + "dbs", + metadata, + sa.Column("id", sa.Integer, primary_key=True), + sa.Column("sqlalchemy_uri", sa.String(1024)), + sa.Column("extra", sa.Text), + sa.Column("changed_on", sa.DateTime), + sa.Column("password", Database.__table__.c.password.type), + sa.Column("encrypted_extra", Database.__table__.c.encrypted_extra.type), + ) + metadata.create_all(engine) + yield engine + engine.dispose() + + +def _run(engine: Engine, upgrading: bool) -> None: + """Execute with real Alembic operations and deliberately small batches.""" + with engine.begin() as conn, patch.object(migration, "BATCH_SIZE", 2): + with Operations.context(MigrationContext.configure(conn)): + (migration.upgrade if upgrading else migration.downgrade)() + + +@pytest.mark.parametrize( + "suffix", + [ + MASKED_SUFFIX, + MASKED_SUFFIX + "?z=%2f&catalog=main&http_path=%2Fsql%2F1.0&a=two+words", + "token:p%40ss%2fword@HOST:443/a%2Fb?duplicate=1&duplicate=2", + ], +) +def test_upgrade_preserves_uri_suffix_and_encrypted_credentials( + engine: Engine, + suffix: str, +) -> None: + """The prefix alone changes; separate encrypted credentials survive.""" + from superset.models.core import Database + + table = sa.Table("dbs", sa.MetaData(), autoload_with=engine) + typed_table = sa.table( + "dbs", + sa.column("id", sa.Integer), + sa.column("sqlalchemy_uri", sa.String), + sa.column("extra", sa.Text), + sa.column("password", Database.__table__.c.password.type), + sa.column("encrypted_extra", Database.__table__.c.encrypted_extra.type), + ) + original_extra = { + "engine_params": { + "connect_args": { + "http_path": "/sql/1.0/warehouses/example", + "catalog": "main", + } + } + } + with engine.begin() as conn: + conn.execute( + typed_table.insert().values( + id=1, + sqlalchemy_uri=LEGACY + suffix, + extra=json.dumps(original_extra), + password="migration-test-password", # noqa: S106 + encrypted_extra='{"other":"test-value"}', + ) + ) + before = conn.execute( + sa.select(table.c.password, table.c.encrypted_extra) + ).one() + _run(engine, True) + _run(engine, True) + with engine.connect() as conn: + row = conn.execute(sa.select(table)).one() + assert row.sqlalchemy_uri == NATIVE + suffix + assert json.loads(row.extra) == {**original_extra, migration.MARKER: True} + assert (row.password, row.encrypted_extra) == tuple(before) + decrypted = conn.execute(sa.select(typed_table)).one() + database = Database( + sqlalchemy_uri=decrypted.sqlalchemy_uri, password=decrypted.password + ) + assert database.url_object.password == "migration-test-password" # noqa: S105 + _run(engine, False) + with engine.connect() as conn: + row = conn.execute(sa.select(table)).one() + assert row.sqlalchemy_uri == LEGACY + suffix + assert json.loads(row.extra) == original_extra + assert (row.password, row.encrypted_extra) == tuple(before) + + +def test_upgrade_batches_and_downgrade_only_converted_rows(engine: Engine) -> None: + """Shrinking the matching set must not skip rows or convert native URIs.""" + table = migration.databases + originals = [ + {"id": i, "sqlalchemy_uri": LEGACY + MASKED_SUFFIX, "extra": "{}"} + for i in range(1, 8) + ] + [ + {"id": 8, "sqlalchemy_uri": NATIVE + MASKED_SUFFIX, "extra": "{}"}, + {"id": 9, "sqlalchemy_uri": "postgresql://localhost/db", "extra": "{}"}, + {"id": 10, "sqlalchemy_uri": "databricks+pyodbc://localhost/db", "extra": "{}"}, + ] + with engine.begin() as conn: + conn.execute(table.insert(), originals) + _run(engine, True) + with engine.connect() as conn: + rows = conn.execute(sa.select(table).order_by(table.c.id)).all() + assert all(row.sqlalchemy_uri == NATIVE + MASKED_SUFFIX for row in rows[:8]) + assert all( + json.loads(row.extra).get(migration.MARKER) is True for row in rows[:7] + ) + assert all(row.extra == "{}" for row in rows[7:]) + _run(engine, False) + _run(engine, False) + with engine.connect() as conn: + assert [ + dict(row._mapping) + for row in conn.execute( + sa.select(table.c.id, table.c.sqlalchemy_uri, table.c.extra).order_by( + table.c.id + ) + ) + ] == originals + + +def test_downgrade_does_not_rewrite_a_subsequently_changed_engine( + engine: Engine, +) -> None: + """Remove the marker without overriding a user's later engine change.""" + table = migration.databases + with engine.begin() as conn: + conn.execute( + table.insert().values( + id=1, sqlalchemy_uri=LEGACY + MASKED_SUFFIX, extra="{}" + ) + ) + _run(engine, True) + with engine.begin() as conn: + assert ( + json.loads(conn.execute(sa.select(table.c.extra)).scalar_one())[ + migration.MARKER + ] + is True + ) + conn.execute(table.update().values(sqlalchemy_uri="postgresql://localhost/db")) + _run(engine, False) + with engine.connect() as conn: + row = conn.execute(sa.select(table)).one() + assert row.sqlalchemy_uri == "postgresql://localhost/db" + assert json.loads(row.extra) == {} + + +@pytest.mark.parametrize( + "extra", + [ + "not json", + "", + None, + "[]", + "null", + '{"' + migration.MARKER, + '{"' + migration.MARKER + '":false}', + json.dumps({migration.MARKER: True}), + "[" * 1100, + ], +) +def test_upgrade_preserves_invalid_extra_and_downgrades_only_unchanged_rows( + engine: Engine, extra: str | None +) -> None: + """A bad row must not block later batches or lose its original text.""" + from datetime import datetime + + table = sa.Table("dbs", sa.MetaData(), autoload_with=engine) + originals = [ + {"id": i, "sqlalchemy_uri": LEGACY + MASKED_SUFFIX, "extra": extra} + for i in range(1, 7) + ] + [ + {"id": 7, "sqlalchemy_uri": NATIVE + MASKED_SUFFIX, "extra": "not json"}, + {"id": 8, "sqlalchemy_uri": LEGACY + MASKED_SUFFIX, "extra": "{}"}, + ] + with engine.begin() as conn: + conn.execute(table.insert(), originals) + _run(engine, True) + _run(engine, True) + with engine.begin() as conn: + rows = conn.execute(sa.select(table).order_by(table.c.id)).all() + assert all(row.sqlalchemy_uri == NATIVE + MASKED_SUFFIX for row in rows) + assert [row.extra for row in rows[:6]] == [extra] * 6 + assert rows[6].extra == "not json" + assert json.loads(rows[7].extra)[migration.MARKER] is True + conn.execute(table.update().where(table.c.id == 2).values(extra="{}")) + conn.execute( + table.update() + .where(table.c.id == 3) + .values(sqlalchemy_uri=NATIVE + MASKED_SUFFIX + "?catalog=edited") + ) + conn.execute( + table.update() + .where(table.c.id == 4) + .values(sqlalchemy_uri="postgresql://localhost/db") + ) + conn.execute( + table.update() + .where(table.c.id == 5) + .values(changed_on=datetime(2026, 1, 1)) + ) + conn.execute(table.delete().where(table.c.id == 6)) + _run(engine, False) + _run(engine, False) + with engine.connect() as conn: + rows = {row.id: row for row in conn.execute(sa.select(table))} + assert rows[1].sqlalchemy_uri == LEGACY + MASKED_SUFFIX + if extra == json.dumps({migration.MARKER: True}): + assert json.loads(rows[1].extra) == {} + else: + assert rows[1].extra == extra + assert rows[2].sqlalchemy_uri == NATIVE + MASKED_SUFFIX + assert rows[2].extra == "{}" + assert rows[3].sqlalchemy_uri == NATIVE + MASKED_SUFFIX + "?catalog=edited" + assert rows[4].sqlalchemy_uri == "postgresql://localhost/db" + assert rows[5].sqlalchemy_uri == NATIVE + MASKED_SUFFIX + assert rows[7].sqlalchemy_uri == NATIVE + MASKED_SUFFIX + assert rows[7].extra == "not json" + assert rows[8].sqlalchemy_uri == LEGACY + MASKED_SUFFIX + assert json.loads(rows[8].extra) == {} + + +def test_converted_connection_extra_can_be_imported(engine: Engine) -> None: + """Exports of converted connections carry the marker; import must accept it.""" + from superset.databases.schemas import ImportV1DatabaseExtraSchema + + table = migration.databases + original_extra = {"engine_params": {"connect_args": {"http_path": "/sql/1.0"}}} + with engine.begin() as conn: + conn.execute( + table.insert().values( + id=1, + sqlalchemy_uri=LEGACY + MASKED_SUFFIX, + extra=json.dumps(original_extra), + ) + ) + _run(engine, True) + with engine.connect() as conn: + exported_extra = json.loads(conn.execute(sa.select(table.c.extra)).scalar_one()) + assert ImportV1DatabaseExtraSchema().load(exported_extra) == { + **original_extra, + migration.MARKER: True, + } + + +@pytest.mark.parametrize("edit", ["uri", "extra", "changed_on", "malformed_extra"]) +def test_downgrade_preserves_edits_to_converted_native_connection( + engine: Engine, edit: str +) -> None: + """An edited Databricks row must not be reverted merely due to its marker.""" + from datetime import datetime + + table = sa.Table("dbs", sa.MetaData(), autoload_with=engine) + with engine.begin() as conn: + conn.execute( + table.insert().values( + id=1, sqlalchemy_uri=LEGACY + MASKED_SUFFIX, extra="{}" + ) + ) + _run(engine, True) + with engine.begin() as conn: + changes = { + "uri": {"sqlalchemy_uri": NATIVE + MASKED_SUFFIX + "?catalog=edited"}, + "extra": {"extra": json.dumps({migration.MARKER: True, "edited": True})}, + "changed_on": {"changed_on": datetime(2026, 1, 1)}, + "malformed_extra": {"extra": "not json"}, + } + conn.execute(table.update().values(**changes[edit])) + expected_uri = conn.execute(sa.select(table.c.sqlalchemy_uri)).scalar_one() + _run(engine, False) + with engine.connect() as conn: + row = conn.execute(sa.select(table)).one() + assert row.sqlalchemy_uri == expected_uri + if edit == "malformed_extra": + assert row.extra == "not json" + return + assert migration.MARKER not in json.loads(row.extra) + if edit == "extra": + assert json.loads(row.extra)["edited"] is True + + +def test_downgrade_imported_marker_and_unmarked_native_rows(engine: Engine) -> None: + """Imported markers retain their meaning, unlike unmarked native rows.""" + table = migration.databases + _run(engine, True) + rows = [ + { + "id": 1, + "sqlalchemy_uri": NATIVE + MASKED_SUFFIX, + "extra": json.dumps({migration.MARKER: True}), + }, + {"id": 2, "sqlalchemy_uri": NATIVE + MASKED_SUFFIX, "extra": "{}"}, + {"id": 3, "sqlalchemy_uri": NATIVE + MASKED_SUFFIX, "extra": "not json"}, + ] + with engine.begin() as conn: + conn.execute(table.insert(), rows) + _run(engine, False) + _run(engine, False) + with engine.connect() as conn: + converted, native, malformed = conn.execute( + sa.select(table).order_by(table.c.id) + ) + assert converted.sqlalchemy_uri == LEGACY + MASKED_SUFFIX + assert json.loads(converted.extra) == {} + assert native.sqlalchemy_uri == NATIVE + MASKED_SUFFIX + assert native.extra == "{}" + assert malformed.sqlalchemy_uri == NATIVE + MASKED_SUFFIX + assert malformed.extra == "not json" + assert not sa.inspect(engine).has_table("databricks_uri_migration") + _run(engine, True) + with engine.connect() as conn: + assert all( + row.sqlalchemy_uri == NATIVE + MASKED_SUFFIX + for row in conn.execute(sa.select(table)) + ) + + +@pytest.mark.parametrize("upgrading", [True, False]) +def test_migration_preserves_edit_after_batch_read( + engine: Engine, upgrading: bool +) -> None: + """A concurrent edit must win over the batch's stale connection settings.""" + table = migration.databases + with engine.begin() as conn: + conn.execute( + table.insert().values( + id=1, sqlalchemy_uri=LEGACY + MASKED_SUFFIX, extra="{}" + ) + ) + if not upgrading: + _run(engine, True) + with engine.begin() as conn: + original_execute = conn.execute + edited = False + + def execute( + statement: sa.sql.Executable, *args: object, **kwargs: object + ) -> Any: + """Materialize a batch, then edit the row before the migration writes.""" + nonlocal edited + result = original_execute(statement, *args, **kwargs) + if isinstance(statement, sa.sql.Select) and not edited: + rows = result.fetchall() + original_execute( + table.update() + .where(table.c.id == 1) + .values( + sqlalchemy_uri="postgresql://edited/db", + extra='{"edited":true}', + ) + ) + edited = True + result = MagicMock() + result.fetchall.return_value = rows + return result + + with ( + Operations.context(MigrationContext.configure(conn)), + patch.object(conn, "execute", side_effect=execute), + ): + if upgrading: + migration.upgrade() + else: + migration.downgrade() + row = original_execute(sa.select(table)).one() + assert row.sqlalchemy_uri == "postgresql://edited/db" + assert row.extra == '{"edited":true}' + if upgrading: + assert original_execute(sa.select(migration.records)).fetchall() == [] diff --git a/tests/unit_tests/models/core_test.py b/tests/unit_tests/models/core_test.py index a14b68d50c8d..3ade2ca17208 100644 --- a/tests/unit_tests/models/core_test.py +++ b/tests/unit_tests/models/core_test.py @@ -287,6 +287,25 @@ def test_dttm_sql_literal( assert SqlaTable(database=database).dttm_sql_literal(dttm, col) == result +@pytest.mark.parametrize( + "uri", + [ + "postgresql://user:pass@host:notaport/db", + "not a SQLAlchemy URI", + "uninstalled_dialect://host/db", + ], +) +def test_connection_details_allow_uri_repair(app_context: None, uri: str) -> None: + """Broken saved connections serialize enough information for the URI editor.""" + from superset.databases.schemas import DatabaseConnectionSchema + + database = Database(database_name="db", sqlalchemy_uri=uri, extra="{}") + result = DatabaseConnectionSchema().dump(database) + assert result["parameters"] == {} + assert result["parameters_schema"] == {} + assert result["sqlalchemy_uri"] == uri + + def test_table_column_database() -> None: database = Database(database_name="db") assert TableColumn(database=database).database is database