Repository: apache/superset
File: superset/db_engine_specs/databricks.py (Python)

Code before:


def monkeypatch_dialect() -> None:
    """
    Monkeypatch dialect to correctly escape single quotes for Databricks.

    The Databricks SQLAlchemy dialect (<3.0) incorrectly escapes single quotes by
    doubling them ('O''Hara') instead of using backslash escaping ('O\'Hara'). The
    fixed version requires SQLAlchemy>=2.0, which is not yet compatible with Superset.

    The challenge: DatabricksDialect inherits from HiveDialect and accesses the
    `colspecs` type mapping through its parent class rather than its own. This means:
    - Direct assignment (`DatabricksDialect.colspecs[types.String] = ...`) works because
    DatabricksDialect.colspecs points to the same dict object as HiveDialect.colspecs,
    but then it affects all other dialects that inherit from HiveDialect, causing
    incorrect escaping for them.
    - Shallow copying then modifying `DatabricksDialect.colspecs` doesn't work
    because it's never accessed.

    The solution: A context-aware approach by creating a custom String type
    that checks which dialect is being used during SQL compilation, allowing
    to apply Databricks-specific escaping only for Databricks connections.
    """
    try:
        from pyhive.sqlalchemy_hive import HiveDialect
        from sqlalchemy.engine.default import DefaultDialect

        class ContextAwareStringType(types.TypeDecorator):
            impl = types.String
            cache_ok = True

            def literal_processor(
                self, dialect: DefaultDialect
            ) -> Callable[[Any], str]:
                if dialect.__class__.__name__ == "DatabricksDialect":
                    return DatabricksStringType().literal_processor(dialect)
                return super().literal_processor(dialect)

        HiveDialect.colspecs[types.String] = ContextAwareStringType

    except ImportError:
        pass


class DatabricksBaseSchema(Schema):
    """
    Fields that are required for both Databricks drivers that uses a
    dynamic form.
    """

    access_token = fields.Str(required=True)

Review comment (suggestion):
Small suggestion, leave the import from line 82 and change this to:

```suggestion
                if dialect.__class__ is DatabricksDialect:
```

Code after:
        return process


def monkeypatch_dialect() -> None:
    """
    Monkeypatch dialect to correctly escape single quotes for Databricks.

    The Databricks SQLAlchemy dialect (<3.0) incorrectly escapes single quotes by
    doubling them ('O''Hara') instead of using backslash escaping ('O\'Hara'). The
    fixed version requires SQLAlchemy>=2.0, which is not yet compatible with Superset.

    Since the DatabricksDialect.colspecs points to the base class (HiveDialect.colspecs)
    we can't patch it without affecting other Hive-based dialects. The solution is to
    introduce a dialect-aware string type so that the change applies only to Databricks.
    """
    try:
        from pyhive.sqlalchemy_hive import HiveDialect

        class ContextAwareStringType(types.TypeDecorator):
            impl = types.String
            cache_ok = True

            def literal_processor(
                self, dialect: DefaultDialect
            ) -> Callable[[Any], str]:
                if dialect.__class__.__name__ == "DatabricksDialect":
                    return DatabricksStringType().literal_processor(dialect)
                return super().literal_processor(dialect)

        HiveDialect.colspecs[types.String] = ContextAwareStringType

    except ImportError:
        pass


class DatabricksBaseSchema(Schema):
    """
    Fields that are required for both Databricks drivers that uses a
    dynamic form.
    """

    access_token = fields.Str(required=True)
    host = fields.Str(required=True)
    port = fields.Integer(
        required=True,
        metadata={"description": __("Database port")},
        validate=Range(min=0, max=2**16, max_inclusive=False),
    )
    encryption = fields.Boolean(
        required=False,
        metadata={"description": __("Use an encrypted connection to the database")},