You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by be...@apache.org on 2024/09/13 14:14:15 UTC
(superset) 01/01: feat: allow configuring an engine context manager
This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch engine-context
in repository https://gitbox.apache.org/repos/asf/superset.git
commit fbd9adfaaf83db830b712f4bd3e79f56b7206f13
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Fri Sep 13 10:13:59 2024 -0400
feat: allow configuring an engine context manager
---
superset/config.py | 32 +++++++++++++++++++++-------
superset/models/core.py | 55 ++++++++++++++++++++++++++-----------------------
2 files changed, 54 insertions(+), 33 deletions(-)
diff --git a/superset/config.py b/superset/config.py
index d43028fe64..bccd5faa60 100644
--- a/superset/config.py
+++ b/superset/config.py
@@ -33,10 +33,11 @@ import os
import re
import sys
from collections import OrderedDict
+from contextlib import contextmanager
from datetime import timedelta
from email.mime.multipart import MIMEMultipart
from importlib.resources import files
-from typing import Any, Callable, Literal, TYPE_CHECKING, TypedDict
+from typing import Any, Callable, Iterator, Literal, TYPE_CHECKING, TypedDict
import click
import pkg_resources
@@ -1142,16 +1143,18 @@ def CSV_TO_HIVE_UPLOAD_DIRECTORY_FUNC( # pylint: disable=invalid-name
# uploading CSVs will be stored.
UPLOADED_CSV_HIVE_NAMESPACE: str | None = None
+
# Function that computes the allowed schemas for the CSV uploads.
# Allowed schemas will be a union of schemas_allowed_for_file_upload
# db configuration and a result of this function.
+def allowed_schemas_for_csv_upload( # pylint: disable=unused-argument
+ database: Database,
+ user: models.User,
+) -> list[str]:
+ return [UPLOADED_CSV_HIVE_NAMESPACE] if UPLOADED_CSV_HIVE_NAMESPACE else []
-# mypy doesn't catch that if case ensures list content being always str
-ALLOWED_USER_CSV_SCHEMA_FUNC: Callable[[Database, models.User], list[str]] = ( # noqa: E731
- lambda database, user: [UPLOADED_CSV_HIVE_NAMESPACE]
- if UPLOADED_CSV_HIVE_NAMESPACE
- else []
-)
+
+ALLOWED_USER_CSV_SCHEMA_FUNC = allowed_schemas_for_csv_upload
# Values that should be treated as nulls for the csv uploads.
CSV_DEFAULT_NA_NAMES = list(STR_NA_VALUES)
@@ -1262,6 +1265,21 @@ ALLOWED_EXTRA_AUTHENTICATIONS: dict[str, dict[str, Callable[..., Any]]] = {}
# The id of a template dashboard that should be copied to every new user
DASHBOARD_TEMPLATE_ID = None
+
+# A context manager that wraps the call to `create_engine`. This can be used for many
+# things, such as chrooting to prevent 3rd party drivers to access the filesystem, or
+# setting up custom configuration for database drivers.
+@contextmanager
+def engine_context_manager( # pylint: disable=unused-argument
+ database: Database,
+ catalog: str | None,
+ schema: str | None,
+) -> Iterator[None]:
+ yield None
+
+
+ENGINE_CONTEXT_MANAGER = engine_context_manager
+
# A callable that allows altering the database connection URL and params
# on the fly, at runtime. This allows for things like impersonation or
# arbitrary logic. For instance you can wire different users to
diff --git a/superset/models/core.py b/superset/models/core.py
index 9d432f8114..a8dd15e3a6 100755
--- a/superset/models/core.py
+++ b/superset/models/core.py
@@ -418,38 +418,41 @@ class Database(Model, AuditMixinNullable, ImportExportMixin): # pylint: disable
)
sqlalchemy_uri = self.sqlalchemy_uri_decrypted
- engine_context = nullcontext()
- ssh_tunnel = override_ssh_tunnel or DatabaseDAO.get_ssh_tunnel(
- database_id=self.id
- )
- if ssh_tunnel:
- # if ssh_tunnel is available build engine with information
- engine_context = ssh_manager_factory.instance.create_tunnel(
+ ssh_tunnel = override_ssh_tunnel or DatabaseDAO.get_ssh_tunnel(self.id)
+ ssh_context_manager = (
+ ssh_manager_factory.instance.create_tunnel(
ssh_tunnel=ssh_tunnel,
sqlalchemy_database_uri=sqlalchemy_uri,
)
+ if ssh_tunnel
+ else nullcontext()
+ )
- with engine_context as server_context:
- if ssh_tunnel and server_context:
- logger.info(
- "[SSH] Successfully created tunnel w/ %s tunnel_timeout + %s ssh_timeout at %s",
- sshtunnel.TUNNEL_TIMEOUT,
- sshtunnel.SSH_TIMEOUT,
- server_context.local_bind_address,
- )
- sqlalchemy_uri = ssh_manager_factory.instance.build_sqla_url(
- sqlalchemy_uri,
- server_context,
- )
+ engine_context_manager = config["ENGINE_CONTEXT_MANAGER"]
+
+ with ssh_context_manager as ssh_context:
+ with engine_context_manager(self, catalog, schema):
+ if ssh_context:
+ logger.info(
+ "[SSH] Successfully created tunnel w/ %s tunnel_timeout + %s "
+ "ssh_timeout at %s",
+ sshtunnel.TUNNEL_TIMEOUT,
+ sshtunnel.SSH_TIMEOUT,
+ ssh_context.local_bind_address,
+ )
+ sqlalchemy_uri = ssh_manager_factory.instance.build_sqla_url(
+ sqlalchemy_uri,
+ ssh_context,
+ )
- yield self._get_sqla_engine(
- catalog=catalog,
- schema=schema,
- nullpool=nullpool,
- source=source,
- sqlalchemy_uri=sqlalchemy_uri,
- )
+ yield self._get_sqla_engine(
+ catalog=catalog,
+ schema=schema,
+ nullpool=nullpool,
+ source=source,
+ sqlalchemy_uri=sqlalchemy_uri,
+ )
def _get_sqla_engine( # pylint: disable=too-many-locals
self,