Dev-iL commented on code in PR #69526:
URL: https://github.com/apache/airflow/pull/69526#discussion_r3601560165
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -334,7 +345,23 @@ def get_sqlalchemy_engine(self, engine_kwargs=None) ->
Engine:
self.log.debug("url: %s", url)
self.log.debug("engine_kwargs: %s", engine_kwargs)
- return create_engine(url=url, **engine_kwargs)
+ try:
+ return create_engine(url=url, **engine_kwargs)
+ except (ImportError, ModuleNotFoundError):
+ # SQLAlchemy resolves a bare "postgresql" scheme to the psycopg2
DB-API by default.
+ # Retry with a driver this hook can actually resolve, since
psycopg2 may not be
+ # installed now that it's an optional extra of
apache-airflow-providers-postgres.
+ parsed_url = make_url(url) if make_url is not None else None
+ if parsed_url is None or parsed_url.drivername != "postgresql":
+ raise
+ # psycopg (v3) may be importable where SQLAlchemy is pinned below
2.0, which has no
+ # native "postgresql+psycopg" dialect and would raise
NoSuchModuleError — so only
+ # prefer it on SQLAlchemy 2.0+; otherwise fall back to psycopg2.
+ if find_spec("psycopg") is not None and _is_sqlalchemy_2():
+ return
create_engine(url=parsed_url.set(drivername="postgresql+psycopg"),
**engine_kwargs)
+ if find_spec("psycopg2") is not None:
+ return
create_engine(url=parsed_url.set(drivername="postgresql+psycopg2"),
**engine_kwargs)
Review Comment:
Done
##########
providers/common/sql/src/airflow/providers/common/sql/hooks/sql.py:
##########
@@ -334,7 +345,23 @@ def get_sqlalchemy_engine(self, engine_kwargs=None) ->
Engine:
self.log.debug("url: %s", url)
self.log.debug("engine_kwargs: %s", engine_kwargs)
- return create_engine(url=url, **engine_kwargs)
+ try:
+ return create_engine(url=url, **engine_kwargs)
+ except (ImportError, ModuleNotFoundError):
+ # SQLAlchemy resolves a bare "postgresql" scheme to the psycopg2
DB-API by default.
+ # Retry with a driver this hook can actually resolve, since
psycopg2 may not be
+ # installed now that it's an optional extra of
apache-airflow-providers-postgres.
+ parsed_url = make_url(url) if make_url is not None else None
+ if parsed_url is None or parsed_url.drivername != "postgresql":
+ raise
+ # psycopg (v3) may be importable where SQLAlchemy is pinned below
2.0, which has no
+ # native "postgresql+psycopg" dialect and would raise
NoSuchModuleError — so only
+ # prefer it on SQLAlchemy 2.0+; otherwise fall back to psycopg2.
+ if find_spec("psycopg") is not None and _is_sqlalchemy_2():
+ return
create_engine(url=parsed_url.set(drivername="postgresql+psycopg"),
**engine_kwargs)
+ if find_spec("psycopg2") is not None:
+ return
create_engine(url=parsed_url.set(drivername="postgresql+psycopg2"),
**engine_kwargs)
Review Comment:
Done
--
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]
