betodealmeida commented on code in PR #21943:
URL: https://github.com/apache/superset/pull/21943#discussion_r1019660134
##########
superset/databases/commands/test_connection.py:
##########
@@ -100,31 +99,36 @@ def ping(engine: Engine) -> bool:
with closing(engine.raw_connection()) as conn:
return engine.dialect.do_ping(conn)
- try:
- alive = func_timeout(
-
int(app.config["TEST_DATABASE_CONNECTION_TIMEOUT"].total_seconds()),
- ping,
- args=(engine,),
- )
- except (sqlite3.ProgrammingError, RuntimeError):
- # SQLite can't run on a separate thread, so ``func_timeout``
fails
- # RuntimeError catches the equivalent error from duckdb.
- alive = engine.dialect.do_ping(engine)
- except FunctionTimedOut as ex:
- raise SupersetTimeoutException(
- error_type=SupersetErrorType.CONNECTION_DATABASE_TIMEOUT,
- message=(
- "Please check your connection details and database
settings, "
- "and ensure that your database is accepting
connections, "
- "then try connecting again."
- ),
- level=ErrorLevel.ERROR,
- extra={"sqlalchemy_uri": database.sqlalchemy_uri},
- ) from ex
- except Exception as ex: # pylint: disable=broad-except
- alive = False
- # So we stop losing the original message if any
- ex_str = str(ex)
+ with database.get_sqla_engine_with_context() as engine:
+ try:
+ alive = func_timeout(
+ int(
+ app.config[
+ "TEST_DATABASE_CONNECTION_TIMEOUT"
+ ].total_seconds()
+ ),
Review Comment:
Nit, I'd compute this outside the `with` block just to improve readability:
```python
timeout = app.config["TEST_DATABASE_CONNECTION_TIMEOUT"].total_seconds()
with database.get_sqla_engine_with_context() as engine:
try:
alive = func_timeout(int(timeout), ping, args=(engine,))
```
Also, `func_timeout` should take floats. :-P
##########
superset/db_engine_specs/base.py:
##########
@@ -894,17 +903,17 @@ def df_to_sql(
:param to_sql_kwargs: The kwargs to be passed to
pandas.DataFrame.to_sql` method
"""
- engine = cls.get_engine(database)
to_sql_kwargs["name"] = table.table
if table.schema:
# Only add schema when it is preset and non empty.
to_sql_kwargs["schema"] = table.schema
- if engine.dialect.supports_multivalues_insert:
- to_sql_kwargs["method"] = "multi"
+ with cls.get_engine(database) as engine:
+ if engine.dialect.supports_multivalues_insert:
+ to_sql_kwargs["method"] = "multi"
Review Comment:
I think we need to improve this a little bit... here we're building an
engine just to check an attribute in the dialect, which means we're setting up
and tearing down an SSH connection just to read an attribute. :-(
Maybe we should add a `get_dialect` method to the DB engine spec, that
builds the engine without the context manager:
```python
@classmethod
def get_dialect(database, schema, source):
engine = database.get_sqla_engine(schema=schema, source=source)
return engine.dialect
```
Then when we only need the dialect we can call this method, which is cheaper.
##########
superset/db_engine_specs/bigquery.py:
##########
@@ -340,8 +340,12 @@ def df_to_sql(
if not table.schema:
raise Exception("The table schema must be defined")
- engine = cls.get_engine(database)
- to_gbq_kwargs = {"destination_table": str(table), "project_id":
engine.url.host}
+ to_gbq_kwargs = {}
+ with cls.get_engine(database) as engine:
+ to_gbq_kwargs = {
+ "destination_table": str(table),
+ "project_id": engine.url.host,
Review Comment:
I wonder if we should add an attribute to DB engine specs annotating if they
support SSH tunnel or not? BigQuery, eg, will probably never support it.
##########
superset/datasets/commands/importers/v1/utils.py:
##########
@@ -168,7 +168,8 @@ def load_data(
connection = session.connection()
else:
logger.warning("Loading data outside the import transaction")
- connection = database.get_sqla_engine()
+ with database.get_sqla_engine_with_context() as engine:
+ connection = engine
Review Comment:
Engine will be closed after line 172, making `df.to_sql` fail. You need to
have `df.to_sql` inside the context manager, while the engine is open.
##########
superset/connectors/sqla/utils.py:
##########
@@ -137,13 +136,18 @@ def get_virtual_table_metadata(dataset: SqlaTable) ->
List[ResultSetColumnType]:
# TODO(villebro): refactor to use same code that's used by
# sql_lab.py:execute_sql_statements
try:
- with closing(engine.raw_connection()) as conn:
- cursor = conn.cursor()
- query = dataset.database.apply_limit_to_sql(statements[0], limit=1)
- db_engine_spec.execute(cursor, query)
- result = db_engine_spec.fetch_data(cursor, limit=1)
- result_set = SupersetResultSet(result, cursor.description,
db_engine_spec)
- cols = result_set.columns
+ with dataset.database.get_sqla_engine_with_context(
+ schema=dataset.schema
+ ) as engine:
+ with closing(engine.raw_connection()) as conn:
Review Comment:
We don't need to do this now, but eventually we could have a
`database.get_raw_connection` context manager, to simplify things a little bit.
This way we could rewrite this as:
```python
with dataset.database.get_raw_connection(schema=dataset.schema) as conn:
cursor = conn.cursor()
...
```
And the implementation of `get_raw_connection()` would take care of closing
the connection:
```python
@contextmanager
def get_raw_connection(...):
with get_sqla_engine_with_context(...) as engine:
with closing(engine.raw_connection()) as conn:
yield conn
```
##########
tests/integration_tests/sqllab_tests.py:
##########
@@ -733,7 +733,7 @@ def test_execute_sql_statements(self,
mock_execute_sql_statement, mock_get_query
mock_query = mock.MagicMock()
mock_query.database.allow_run_async = False
mock_cursor = mock.MagicMock()
-
mock_query.database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value
= (
+
mock_query.database.get_sqla_engine_with_context.return_value.__enter__.return_value.raw_connection.return_value.cursor.return_value
= (
Review Comment:
You can replace all intermediary `.return_value` with `()` (but not the
last):
```suggestion
mock_query.database.get_sqla_engine_with_context().__enter__().raw_connection().cursor.return_value
= (
```
##########
superset/examples/paris.py:
##########
@@ -28,29 +28,29 @@
def load_paris_iris_geojson(only_metadata: bool = False, force: bool = False)
-> None:
tbl_name = "paris_iris_mapping"
database = database_utils.get_example_database()
- engine = database.get_sqla_engine()
- schema = inspect(engine).default_schema_name
- table_exists = database.has_table_by_name(tbl_name)
+ with database.get_sqla_engine_with_context() as engine:
Review Comment:
We really should convert these old examples into the new format (YAML +
CSV)...
##########
superset/db_engine_specs/hive.py:
##########
@@ -205,7 +203,8 @@ def df_to_sql(
if table_exists:
raise SupersetException("Table already exists")
elif to_sql_kwargs["if_exists"] == "replace":
- engine.execute(f"DROP TABLE IF EXISTS {str(table)}")
+ with cls.get_engine(database) as engine:
+ engine.execute(f"DROP TABLE IF EXISTS {str(table)}")
Review Comment:
It's interesting that sometimes we use `engine.raw_connection().execute`,
and others we use `engine.execute`. Ideally we should standardize in the latter
wherever possible, since it's more concise.
##########
superset/databases/commands/validate.py:
##########
@@ -101,30 +101,30 @@ def run(self) -> None:
database.set_sqlalchemy_uri(sqlalchemy_uri)
database.db_engine_spec.mutate_db_for_connection_test(database)
- engine = database.get_sqla_engine()
- try:
- with closing(engine.raw_connection()) as conn:
- alive = engine.dialect.do_ping(conn)
- except Exception as ex:
- url = make_url_safe(sqlalchemy_uri)
- context = {
- "hostname": url.host,
- "password": url.password,
- "port": url.port,
- "username": url.username,
- "database": url.database,
- }
- errors = database.db_engine_spec.extract_errors(ex, context)
- raise DatabaseTestConnectionFailedError(errors) from ex
+ with database.get_sqla_engine_with_context() as engine:
+ try:
+ with closing(engine.raw_connection()) as conn:
+ alive = engine.dialect.do_ping(conn)
+ except Exception as ex:
+ url = make_url_safe(sqlalchemy_uri)
+ context = {
+ "hostname": url.host,
+ "password": url.password,
+ "port": url.port,
+ "username": url.username,
+ "database": url.database,
+ }
+ errors = database.db_engine_spec.extract_errors(ex, context)
+ raise DatabaseTestConnectionFailedError(errors) from ex
- if not alive:
- raise DatabaseOfflineError(
- SupersetError(
- message=__("Database is offline."),
- error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
- level=ErrorLevel.ERROR,
- ),
- )
+ if not alive:
+ raise DatabaseOfflineError(
+ SupersetError(
+ message=__("Database is offline."),
+ error_type=SupersetErrorType.GENERIC_DB_ENGINE_ERROR,
+ level=ErrorLevel.ERROR,
+ ),
+ )
Review Comment:
We can dedent this block, no?
##########
superset/models/core.py:
##########
@@ -370,11 +370,11 @@ def get_sqla_engine_with_context(
source: Optional[utils.QuerySource] = None,
) -> Engine:
try:
- yield self.get_sqla_engine(schema=schema, nullpool=nullpool,
source=source)
+ yield self._get_sqla_engine(schema=schema, nullpool=nullpool,
source=source)
except Exception as ex:
- raise self.db_engine_spec.get_dbapi_mapped_exception(ex)
+ raise ex
Review Comment:
You don't need `try/except` here if you're raising all exceptions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
