superset icon indicating copy to clipboard operation
superset copied to clipboard

feat: Expose different hooks dynamically to inject different database

Open wugeer opened this issue 8 months ago • 0 comments

This environment configuration setting hook allows administrators to alter different database connection parameters on the fly based on user information and hook strategy. This can be use for a variety of purposes:

  • rewire a subset of users to use different database user accounts,for example, integrating with ldap
  • pass user related information to the database for logging or QoS purposes
  • custom, per-database-engine, environment-specific impersonation

SUMMARY

Although there is already a DB_CONNECTION_MUTATOR hook available to inject database connection logic on the fly, if we do not modify the code in superset/models/core.py within the Database.get_sqla_engine method, this injection will affect all data source connections. Specifically, if different data sources require different injection logic, then we need to make different modifications to the code within superset/models/core.py to handle the different database types. Therefore, I have added a DB_CONNECTION_MODIFIER hook (clearly, DB_CONNECTION_MUTATOR would be a better name, but it is already in use, so I used a synonym instead), which allows for different data sources to implement different injection logic on the fly, as long as these database injection logic classes implement the run method of the ImpalaURLModifier class. This increases the flexibility of injecting database connection logic on the fly.

BEFORE/AFTER SCREENSHOTS OR ANIMATED GIF

TESTING INSTRUCTIONS

For example, when each user connects to the postgres data source, their use the current username to connect to the database, and the password is new_password_ str plus the username. Add a new class PostgresDBConnectModifier in file superset/utils/database_connect_modifier.py

class PostgresDBConnectModifier(BaseDBConnectModifier):
    # When connecting to a postgres data source,
    # replace the default connection username and password

    @classmethod
    def run(cls, sqlalchemy_url: URL, params: dict[str, Any], username: str, *args: Any,
            **kwargs: Any) -> (URL, dict[str, Any]):
        new_password = cls._get_new_password(username)
        sqlalchemy_url.username = username
        sqlalchemy_url.password = new_password
        return sqlalchemy_url, params

    @staticmethod
    def _get_new_password(username):
        # Implement password generation logic
        return 'new_password_' + username

superset/config.py中

DB_CONNECTION_MODIFIER_ENABLED = True
DB_CONNECTION_MODIFIER: dict[str, type[BaseDBConnectModifier]] = {
    "postgresql": PostgresDBConnectModifier,
}

The user logs in to the superset page and queries the postgres database table in sqllab; or the underlying data table triggered by the chart page is the postgres data source.

ADDITIONAL INFORMATION

  • [ ] Has associated issue:
  • [x] Required feature flags:
  • [ ] Changes UI
  • [ ] Includes DB Migration (follow approval process in SIP-59)
    • [ ] Migration is atomic, supports rollback & is backwards-compatible
    • [ ] Confirm DB migration upgrade and downgrade tested
    • [ ] Runtime estimates and downtime expectations provided
  • [x] Introduces new feature or API
  • [ ] Removes existing feature or API

wugeer avatar Jun 18 '24 03:06 wugeer