You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by el...@apache.org on 2026/09/15 05:15:43 UTC
(superset) 01/01: feat(databricks): migrate ODBC connections to the connector driver
This is an automated email from the ASF dual-hosted git repository.
eschutho pushed a commit to branch sc-120996-databricks-odbc-migration
in repository https://gitbox.apache.org/repos/asf/superset.git
commit a09f0b16f0fe9ae5b4fb8b5531d12fb3acedbe12
Author: Elizabeth Thompson <es...@gmail.com>
AuthorDate: Tue Sep 15 05:15:30 2026 +0000
feat(databricks): migrate ODBC connections to the connector driver
Databricks connections on the legacy `databricks+pyodbc` driver resolve to
`DatabricksODBCEngineSpec`, which has no `parameters_schema` and so only ever
renders the raw SQLAlchemy URI text field. The `databricks+connector` driver
resolves to `DatabricksNativeEngineSpec`, which defines a `parameters_schema`
and renders the dynamic connection form (epic sc-120993).
The two drivers share the same URL shape and both read http_path/catalog/schema
from `extra.engine_params.connect_args`, so migrating is a scheme-only driver
swap that preserves every credential and setting byte-for-byte. `upgrade`
records a marker in `extra` so `downgrade` reverts only the rows it migrated and
never touches natively-created connector connections; the round trip is exact
and the migration is idempotent.
Co-Authored-By: Claude Opus 4.8 <no...@anthropic.com>
---
...9e21b84_migrate_databricks_odbc_to_connector.py | 147 +++++++++++++
.../test_migrate_databricks_odbc_to_connector.py | 243 +++++++++++++++++++++
2 files changed, 390 insertions(+)
diff --git a/superset/migrations/versions/2026-09-15_00-00_a3f5c9e21b84_migrate_databricks_odbc_to_connector.py b/superset/migrations/versions/2026-09-15_00-00_a3f5c9e21b84_migrate_databricks_odbc_to_connector.py
new file mode 100644
index 00000000000..09cfa15a627
--- /dev/null
+++ b/superset/migrations/versions/2026-09-15_00-00_a3f5c9e21b84_migrate_databricks_odbc_to_connector.py
@@ -0,0 +1,147 @@
+# 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.
+"""migrate Databricks ODBC connections to the connector driver
+
+Part of the effort to move every Databricks connection onto the dynamic
+connection form. Superset renders that form when the resolved engine spec
+defines a ``parameters_schema``. The legacy ``databricks+pyodbc`` driver resolves
+to ``DatabricksODBCEngineSpec``, which has no ``parameters_schema`` and therefore
+only ever exposes the raw "SQLAlchemy URI" text field. The ``databricks+connector``
+driver resolves to ``DatabricksNativeEngineSpec``, which does define a
+``parameters_schema`` and so renders the dynamic form.
+
+The two drivers share the same URL shape --
+
+ databricks+<driver>://token:<access_token>@<host>:<port>/<database>
+
+-- and both read ``http_path``, ``catalog`` and ``schema`` from
+``extra.engine_params.connect_args`` rather than from the URL. Migrating a
+connection is therefore a driver swap in the URL scheme; every credential and
+setting (access token, host, port, database, the encryption ``ssl`` query
+parameter, and the entire ``extra`` / ``encrypted_extra`` payload) is preserved
+byte-for-byte. Only the scheme prefix is rewritten, so unrelated occurrences of
+the driver string inside credentials or query values are left untouched.
+
+To keep ``downgrade`` from clobbering connections a user created directly on
+``databricks+connector`` (which must stay on that driver), ``upgrade`` records a
+marker (``MIGRATION_MARKER``) in ``extra``. ``downgrade`` reverts only the rows
+carrying that marker and removes it, so the round trip is exact and native
+connector connections are never touched. The marker is an unknown key to the
+``extra`` validator, which tolerates it, and is cosmetic until the connection is
+next saved through the API.
+
+Revision ID: a3f5c9e21b84
+Revises: 88a01c781622
+Create Date: 2026-09-15 00:00:00.000000
+
+"""
+
+from __future__ import annotations
+
+from alembic import op
+from sqlalchemy import Column, Integer, Text
+from sqlalchemy.orm import declarative_base
+
+from superset import db
+from superset.migrations.shared.utils import paginated_update
+from superset.utils import json
+
+# revision identifiers, used by Alembic.
+revision = "a3f5c9e21b84"
+down_revision = "88a01c781622"
+
+Base = declarative_base()
+
+ODBC_PREFIX = "databricks+pyodbc"
+CONNECTOR_PREFIX = "databricks+connector"
+
+# Marker written to ``extra`` on upgrade so downgrade can revert exactly the
+# rows this migration touched and leave native connector connections alone.
+MIGRATION_MARKER = "migrated_from_databricks_odbc"
+
+
+class Database(Base): # type: ignore
+ __tablename__ = "dbs"
+
+ id = Column(Integer, primary_key=True)
+ sqlalchemy_uri = Column(Text, nullable=False)
+ extra = Column(Text)
+
+
+def _swap_scheme(uri: str, old: str, new: str) -> str | None:
+ """
+ Rewrite only the driver in a URI's scheme, leaving the rest byte-identical.
+
+ Returns ``None`` when the URI is not on ``old``, so callers can skip the
+ write (which keeps the migration idempotent).
+ """
+ if not uri.startswith(f"{old}://"):
+ return None
+ return f"{new}{uri[len(old) :]}"
+
+
+def _load_extra(extra: str | None) -> dict:
+ """Parse a database's ``extra`` JSON, treating empty/absent as ``{}``."""
+ if not extra:
+ return {}
+ loaded = json.loads(extra)
+ if not isinstance(loaded, dict):
+ # ``extra`` is always a JSON object in practice; anything else is
+ # malformed and should not be silently wrapped or dropped.
+ raise ValueError("Database.extra is not a JSON object")
+ return loaded
+
+
+def upgrade() -> None:
+ bind = op.get_bind()
+ session = db.Session(bind=bind)
+
+ query = session.query(Database).filter(
+ Database.sqlalchemy_uri.like(f"{ODBC_PREFIX}%")
+ )
+ for database in paginated_update(query):
+ updated = _swap_scheme(database.sqlalchemy_uri, ODBC_PREFIX, CONNECTOR_PREFIX)
+ if updated is None:
+ continue
+ database.sqlalchemy_uri = updated
+ extra = _load_extra(database.extra)
+ extra[MIGRATION_MARKER] = True
+ database.extra = json.dumps(extra)
+
+ session.commit()
+
+
+def downgrade() -> None:
+ bind = op.get_bind()
+ session = db.Session(bind=bind)
+
+ query = session.query(Database).filter(
+ Database.sqlalchemy_uri.like(f"{CONNECTOR_PREFIX}%")
+ )
+ for database in paginated_update(query):
+ extra = _load_extra(database.extra)
+ if not extra.get(MIGRATION_MARKER):
+ # Not migrated by us -- a natively created connector connection.
+ continue
+ reverted = _swap_scheme(database.sqlalchemy_uri, CONNECTOR_PREFIX, ODBC_PREFIX)
+ if reverted is None:
+ continue
+ database.sqlalchemy_uri = reverted
+ extra.pop(MIGRATION_MARKER, None)
+ database.extra = json.dumps(extra) if extra else None
+
+ session.commit()
diff --git a/tests/unit_tests/migrations/test_migrate_databricks_odbc_to_connector.py b/tests/unit_tests/migrations/test_migrate_databricks_odbc_to_connector.py
new file mode 100644
index 00000000000..ff70ed736d5
--- /dev/null
+++ b/tests/unit_tests/migrations/test_migrate_databricks_odbc_to_connector.py
@@ -0,0 +1,243 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Tests for migration ``a3f5c9e21b84_migrate_databricks_odbc_to_connector``.
+
+Covers the scheme-swap helper, the upgrade() path over a mixture of Databricks
+ODBC, native connector and unrelated connections, downgrade() protecting
+natively-created connector connections, and the full round trip.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from unittest.mock import patch
+
+import pytest
+from sqlalchemy import create_engine
+from sqlalchemy.orm import Session
+
+from superset.utils import json
+
+migration = import_module(
+ "superset.migrations.versions."
+ "2026-09-15_00-00_a3f5c9e21b84_migrate_databricks_odbc_to_connector"
+)
+
+Database = migration.Database
+_swap_scheme = migration._swap_scheme
+ODBC_PREFIX = migration.ODBC_PREFIX
+CONNECTOR_PREFIX = migration.CONNECTOR_PREFIX
+MARKER = migration.MIGRATION_MARKER
+
+# Superset stores the password as this mask rather than the real credential
+MASK = "X" * 10
+
+# Databricks stores http_path (and optionally catalog/schema) here for both
+# the ODBC and the connector driver, so it rides along untouched.
+EXTRA = json.dumps(
+ {"engine_params": {"connect_args": {"http_path": "/sql/1.0/warehouses/abc"}}}
+)
+
+
+@pytest.fixture
+def engine():
+ engine = create_engine("sqlite:///:memory:")
+ migration.Base.metadata.create_all(engine)
+ return engine
+
+
+def _run(migrate, conn) -> None:
+ session = Session(bind=conn)
+ with (
+ patch.object(migration, "op") as mock_op,
+ patch.object(migration, "db") as mock_db,
+ ):
+ mock_op.get_bind.return_value = conn
+ mock_db.Session.return_value = session
+ migrate()
+
+
+@pytest.mark.parametrize(
+ "uri,expected",
+ [
+ (
+ f"databricks+pyodbc://token:{MASK}@host:443/db",
+ f"databricks+connector://token:{MASK}@host:443/db",
+ ),
+ # unrelated query parameters (e.g. the encryption flag) are preserved
+ (
+ f"databricks+pyodbc://token:{MASK}@host:443/db?ssl=1",
+ f"databricks+connector://token:{MASK}@host:443/db?ssl=1",
+ ),
+ ],
+)
+def test_swap_scheme_rewrites_only_the_driver(uri: str, expected: str) -> None:
+ assert _swap_scheme(uri, ODBC_PREFIX, CONNECTOR_PREFIX) == expected
+
+
+@pytest.mark.parametrize(
+ "uri",
+ [
+ # already migrated / different driver
+ f"databricks+connector://token:{MASK}@host:443/db",
+ f"databricks://token:{MASK}@host:443?http_path=/x",
+ # the driver string only appears in credentials, not the scheme
+ f"postgresql://databricks+pyodbc:{MASK}@host:5432/db",
+ ],
+)
+def test_swap_scheme_leaves_non_matching_uris_alone(uri: str) -> None:
+ assert _swap_scheme(uri, ODBC_PREFIX, CONNECTOR_PREFIX) is None
+
+
+def test_upgrade_rewrites_only_odbc_connections(engine) -> None:
+ with Session(engine) as seed:
+ seed.add_all(
+ [
+ Database(
+ id=1,
+ sqlalchemy_uri=f"databricks+pyodbc://token:{MASK}@host:443/main",
+ extra=EXTRA,
+ ),
+ # a Databricks connection already on the connector driver
+ Database(
+ id=2,
+ sqlalchemy_uri=f"databricks+connector://token:{MASK}@host:443/main",
+ extra=EXTRA,
+ ),
+ # a different engine, untouched
+ Database(
+ id=3,
+ sqlalchemy_uri=f"postgresql://user:{MASK}@host:5432/db",
+ extra=None,
+ ),
+ ]
+ )
+ seed.commit()
+
+ with engine.begin() as conn:
+ _run(migration.upgrade, conn)
+
+ with Session(engine) as verify:
+ migrated = verify.get(Database, 1)
+ assert (
+ migrated.sqlalchemy_uri
+ == f"databricks+connector://token:{MASK}@host:443/main"
+ )
+ # credentials-bearing extra is preserved and the marker is added
+ extra = json.loads(migrated.extra)
+ assert extra["engine_params"]["connect_args"]["http_path"] == (
+ "/sql/1.0/warehouses/abc"
+ )
+ assert extra[MARKER] is True
+
+ # native connector connection is left exactly as it was (no marker)
+ native = verify.get(Database, 2)
+ assert (
+ native.sqlalchemy_uri
+ == f"databricks+connector://token:{MASK}@host:443/main"
+ )
+ assert MARKER not in json.loads(native.extra)
+
+ assert (
+ verify.get(Database, 3).sqlalchemy_uri
+ == f"postgresql://user:{MASK}@host:5432/db"
+ )
+
+
+def test_upgrade_is_idempotent(engine) -> None:
+ with Session(engine) as seed:
+ seed.add(
+ Database(
+ id=1,
+ sqlalchemy_uri=f"databricks+pyodbc://token:{MASK}@host:443/main",
+ extra=EXTRA,
+ )
+ )
+ seed.commit()
+
+ for _ in range(2):
+ with engine.begin() as conn:
+ _run(migration.upgrade, conn)
+
+ with Session(engine) as verify:
+ assert (
+ verify.get(Database, 1).sqlalchemy_uri
+ == f"databricks+connector://token:{MASK}@host:443/main"
+ )
+
+
+def test_downgrade_only_reverts_migrated_connections(engine) -> None:
+ """
+ downgrade must revert the rows upgrade migrated (marker present) and leave
+ natively-created connector connections on the connector driver.
+ """
+ marked_extra = json.dumps(
+ {
+ "engine_params": {"connect_args": {"http_path": "/sql/1.0/warehouses/abc"}},
+ MARKER: True,
+ }
+ )
+ with Session(engine) as seed:
+ seed.add_all(
+ [
+ Database(
+ id=1,
+ sqlalchemy_uri=f"databricks+connector://token:{MASK}@host:443/main",
+ extra=marked_extra,
+ ),
+ Database(
+ id=2,
+ sqlalchemy_uri=f"databricks+connector://token:{MASK}@host:443/main",
+ extra=EXTRA,
+ ),
+ ]
+ )
+ seed.commit()
+
+ with engine.begin() as conn:
+ _run(migration.downgrade, conn)
+
+ with Session(engine) as verify:
+ reverted = verify.get(Database, 1)
+ assert (
+ reverted.sqlalchemy_uri == f"databricks+pyodbc://token:{MASK}@host:443/main"
+ )
+ # the marker is stripped so the row is indistinguishable from the original
+ assert MARKER not in json.loads(reverted.extra)
+
+ native = verify.get(Database, 2)
+ assert (
+ native.sqlalchemy_uri
+ == f"databricks+connector://token:{MASK}@host:443/main"
+ )
+
+
+def test_round_trip_through_upgrade_and_downgrade(engine) -> None:
+ original_uri = f"databricks+pyodbc://token:{MASK}@host:443/main"
+ with Session(engine) as seed:
+ seed.add(Database(id=1, sqlalchemy_uri=original_uri, extra=EXTRA))
+ seed.commit()
+
+ with engine.begin() as conn:
+ _run(migration.upgrade, conn)
+ with engine.begin() as conn:
+ _run(migration.downgrade, conn)
+
+ with Session(engine) as verify:
+ database = verify.get(Database, 1)
+ assert database.sqlalchemy_uri == original_uri
+ assert json.loads(database.extra) == json.loads(EXTRA)