This is an automated email from the ASF dual-hosted git repository. hugh pushed a commit to branch fix-sqllab-int-ssh-tunnel in repository https://gitbox.apache.org/repos/asf/superset.git
commit d38af27b8fac8ef7451702b4e2ddfa5db1f92e98 Author: hughhhh <[email protected]> AuthorDate: Mon Nov 28 21:11:10 2022 -0500 fix merge conflicts --- superset/databases/schemas.py | 2 +- superset/models/core.py | 94 +++++++++++++++++++++++++------------------ superset/views/core.py | 1 + 3 files changed, 56 insertions(+), 41 deletions(-) diff --git a/superset/databases/schemas.py b/superset/databases/schemas.py index fad2da0fee..b874d147a2 100644 --- a/superset/databases/schemas.py +++ b/superset/databases/schemas.py @@ -502,7 +502,7 @@ class DatabaseTestConnectionSchema(Schema, DatabaseParametersSchemaMixin): validate=[Length(1, 1024), sqlalchemy_uri_validator], ) - ssh_tunnel_credentials = fields.Nested(DatabaseSSHTunnel, allow_none=True) + ssh_tunnel = fields.Nested(DatabaseSSHTunnel, allow_none=True) class TableMetadataOptionsResponseSchema(Schema): diff --git a/superset/models/core.py b/superset/models/core.py index fd6ad8be93..e45d30b020 100755 --- a/superset/models/core.py +++ b/superset/models/core.py @@ -372,11 +372,12 @@ class Database( schema: Optional[str] = None, nullpool: bool = True, source: Optional[utils.QuerySource] = None, - ssh_tunnel: Optional["SSHTunnel"] = None, + override_ssh_tunnel: Optional["SSHTunnel"] = None, ) -> Engine: ssh_params = {} - if ssh_tunnel: - # build with override + from superset.databases.dao import DatabaseDAO + if ssh_tunnel := override_ssh_tunnel or DatabaseDAO.get_ssh_tunnel(database_id=self.id): + # if ssh_tunnel is available build engine with information url = make_url_safe(self.sqlalchemy_uri_decrypted) ssh_tunnel.bind_host = url.host ssh_tunnel.bind_port = url.port @@ -393,13 +394,9 @@ class Database( raise ex else: - # do look up in table for using database_id - try: - yield self._get_sqla_engine( - schema=schema, nullpool=nullpool, source=source - ) - except Exception as ex: - raise ex + yield self._get_sqla_engine( + schema=schema, nullpool=nullpool, source=source + ) def _get_sqla_engine( self, @@ -568,8 +565,8 @@ class Database( @property def inspector(self) -> Inspector: - engine = self._get_sqla_engine() - return sqla.inspect(engine) + with self.get_sqla_engine_with_context() as engine: + return sqla.inspect(engine) @cache_util.memoized_func( key="db:{self.id}:schema:{schema}:table_list", @@ -594,14 +591,17 @@ class Database( :return: The table/schema pairs """ try: - return { - (table, schema) - for table in self.db_engine_spec.get_table_names( - database=self, - inspector=self.inspector, - schema=schema, - ) - } + with self.get_sqla_engine_with_context() as engine: + inspector = sqla.inspect(engine) + tables = { + (table, schema) + for table in self.db_engine_spec.get_table_names( + database=self, + inspector=inspector, + schema=schema, + ) + } + return tables except Exception as ex: raise self.db_engine_spec.get_dbapi_mapped_exception(ex) @@ -628,17 +628,25 @@ class Database( :return: set of views """ try: - return { - (view, schema) - for view in self.db_engine_spec.get_view_names( - database=self, - inspector=self.inspector, - schema=schema, - ) - } + with self.get_sqla_engine_with_context() as engine: + inspector = sqla.inspect(engine) + return { + (view, schema) + for view in self.db_engine_spec.get_view_names( + database=self, + inspector=inspector, + schema=schema, + ) + } except Exception as ex: raise self.db_engine_spec.get_dbapi_mapped_exception(ex) + @contextmanager + def get_inspector_with_context(self): + with self.get_sqla_engine_with_context() as engine: + yield sqla.inspect(engine) + + @cache_util.memoized_func( key="db:{self.id}:schema_list", cache=cache_manager.cache, @@ -660,7 +668,9 @@ class Database( :return: schema list """ try: - return self.db_engine_spec.get_schema_names(self.inspector) + with self.get_sqla_engine_with_context() as engine: + inspector = sqla.inspect(engine) + return self.db_engine_spec.get_schema_names(inspector) except Exception as ex: raise self.db_engine_spec.get_dbapi_mapped_exception(ex) from ex @@ -728,7 +738,8 @@ class Database( def get_columns( self, table_name: str, schema: Optional[str] = None ) -> List[Dict[str, Any]]: - return self.db_engine_spec.get_columns(self.inspector, table_name, schema) + with self.get_inspector_with_context() as inspector: + return self.db_engine_spec.get_columns(inspector, table_name, schema) def get_metrics( self, @@ -740,26 +751,29 @@ class Database( def get_indexes( self, table_name: str, schema: Optional[str] = None ) -> List[Dict[str, Any]]: - indexes = self.inspector.get_indexes(table_name, schema) - return self.db_engine_spec.normalize_indexes(indexes) + with self.get_inspector_with_context() as inspector: + indexes = inspector.get_indexes(table_name, schema) + return self.db_engine_spec.normalize_indexes(indexes) def get_pk_constraint( self, table_name: str, schema: Optional[str] = None ) -> Dict[str, Any]: - pk_constraint = self.inspector.get_pk_constraint(table_name, schema) or {} + with self.get_inspector_with_context() as inspector: + pk_constraint = inspector.get_pk_constraint(table_name, schema) or {} - def _convert(value: Any) -> Any: - try: - return utils.base_json_conv(value) - except TypeError: - return None + def _convert(value: Any) -> Any: + try: + return utils.base_json_conv(value) + except TypeError: + return None - return {key: _convert(value) for key, value in pk_constraint.items()} + return {key: _convert(value) for key, value in pk_constraint.items()} def get_foreign_keys( self, table_name: str, schema: Optional[str] = None ) -> List[Dict[str, Any]]: - return self.inspector.get_foreign_keys(table_name, schema) + with self.get_inspector_with_context() as inspector: + return inspector.get_foreign_keys(table_name, schema) def get_schema_access_for_file_upload( # pylint: disable=invalid-name self, diff --git a/superset/views/core.py b/superset/views/core.py index 534f8f667d..29107813d4 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2405,6 +2405,7 @@ class Superset(BaseSupersetView): # pylint: disable=too-many-public-methods } execution_context = SqlJsonExecutionContext(request.json) command = self._create_sql_json_command(execution_context, log_params) + # breakpoint() command_result: CommandResult = command.run() return self._create_response_from_execution_context(command_result) except SqlLabException as ex:
