Skip to content

Replaced hql references to sql in TrinoHook and PrestoHook #21597 #21630

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eladkal merged 17 commits into from
Feb 22, 2022
Merged

Replaced hql references to sql in TrinoHook and PrestoHook #21597 #21630

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
6f530d7
Added template_ext = ('.json') to databricks operators #18925
utkarsharma2 Feb 11, 2022
b5a2211
Corrected the template_ext value from 'json' to '.json'
utkarsharma2 Feb 12, 2022
9d83cd6
Merge branch 'apache:main' into main
utkarsharma2 Feb 12, 2022
85b2675
Merge branch 'apache:main' into main
utkarsharma2 Feb 13, 2022
c8dd382
Added retries to LivyHook #19384
Feb 13, 2022
a7b58b8
Static code fixes.
Feb 13, 2022
eb6543f
Renamed _retry_args to retry_args.
Feb 14, 2022
3e86860
Passed _retry_args in run_with_advanced_retry().
Feb 14, 2022
104986e
Merge branch 'apache:main' into main
utkarsharma2 Feb 15, 2022
dae69e1
Merge branch 'apache:main' into main
utkarsharma2 Feb 17, 2022
a9b30ee
Added deprecation warning when using hql param.
utkarsharma2 Feb 19, 2022
e0bfed2
Fixed static checks.
utkarsharma2 Feb 19, 2022
bd3f869
Added @overload signatures and sphink's markers to remove 'hql' from …
utkarsharma2 Feb 21, 2022
9db1cb1
Fixed docstings.
utkarsharma2 Feb 21, 2022
532fce2
Added 'sphinx-autoapi-skip' in docstrings of overloaded fucntions.
utkarsharma2 Feb 21, 2022
f7ea13a
Fixed signature of methods in presto hook and trino hook.
utkarsharma2 Feb 21, 2022
c3ce8a1
Merge branch 'apache:main' into main
utkarsharma2 Feb 21, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
115 changes: 103 additions & 12 deletions airflow/providers/presto/hooks/presto.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
# specific language governing permissions and limitations
# under the License.
import os
from typing import Any, Callable, Iterable, Optional
import warnings
from typing import Any, Callable, Iterable, Optional, overload

import prestodb
from prestodb.exceptions import DatabaseError
Expand Down Expand Up @@ -111,27 +112,87 @@ def get_isolation_level(self) -> Any:
def _strip_sql(sql: str) -> str:
return sql.strip().rstrip(';')

def get_records(self, hql, parameters: Optional[dict] = None):
"""Get a set of records from Presto"""
@overload
def get_records(self, sql: str = "", parameters: Optional[dict] = None):
"""Get a set of records from Presto

:param sql: SQL statement to be executed.
:param parameters: The parameters to render the SQL query with.
"""

@overload
def get_records(self, sql: str = "", parameters: Optional[dict] = None, hql: str = ""):
""":sphinx-autoapi-skip:"""

def get_records(self, sql: str = "", parameters: Optional[dict] = None, hql: str = ""):
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

try:
return super().get_records(self._strip_sql(hql), parameters)
return super().get_records(self._strip_sql(sql), parameters)
except DatabaseError as e:
raise PrestoException(e)

def get_first(self, hql: str, parameters: Optional[dict] = None) -> Any:
"""Returns only the first row, regardless of how many rows the query returns."""
@overload
def get_first(self, sql: str = "", parameters: Optional[dict] = None) -> Any:
"""Returns only the first row, regardless of how many rows the query returns.

:param sql: SQL statement to be executed.
:param parameters: The parameters to render the SQL query with.
"""

@overload
def get_first(self, sql: str = "", parameters: Optional[dict] = None, hql: str = "") -> Any:
""":sphinx-autoapi-skip:"""

def get_first(self, sql: str = "", parameters: Optional[dict] = None, hql: str = "") -> Any:
Comment thread
uranusjr marked this conversation as resolved.
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

try:
return super().get_first(self._strip_sql(hql), parameters)
return super().get_first(self._strip_sql(sql), parameters)
except DatabaseError as e:
raise PrestoException(e)

def get_pandas_df(self, hql, parameters=None, **kwargs):
"""Get a pandas dataframe from a sql query."""
@overload
def get_pandas_df(self, sql: str = "", parameters=None, **kwargs):
"""Get a pandas dataframe from a sql query.

:param sql: SQL statement to be executed.
:param parameters: The parameters to render the SQL query with.
"""

@overload
def get_pandas_df(self, sql: str = "", parameters=None, hql: str = "", **kwargs):
""":sphinx-autoapi-skip:"""

def get_pandas_df(self, sql: str = "", parameters=None, hql: str = "", **kwargs):
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

import pandas

cursor = self.get_cursor()
try:
cursor.execute(self._strip_sql(hql), parameters)
cursor.execute(self._strip_sql(sql), parameters)
data = cursor.fetchall()
except DatabaseError as e:
raise PrestoException(e)
Expand All @@ -143,15 +204,45 @@ def get_pandas_df(self, hql, parameters=None, **kwargs):
df = pandas.DataFrame(**kwargs)
return df

@overload
def run(
self,
hql,
sql: str = "",
autocommit: bool = False,
parameters: Optional[dict] = None,
handler: Optional[Callable] = None,
) -> None:
"""Execute the statement against Presto. Can be used to create views."""
return super().run(sql=self._strip_sql(hql), parameters=parameters, handler=handler)

@overload
def run(
self,
sql: str = "",
autocommit: bool = False,
parameters: Optional[dict] = None,
handler: Optional[Callable] = None,
hql: str = "",
) -> None:
""":sphinx-autoapi-skip:"""

def run(
self,
sql: str = "",
autocommit: bool = False,
parameters: Optional[dict] = None,
handler: Optional[Callable] = None,
hql: str = "",
) -> None:
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

return super().run(sql=self._strip_sql(sql), parameters=parameters, handler=handler)

def insert_rows(
self,
Expand Down
121 changes: 109 additions & 12 deletions airflow/providers/trino/hooks/trino.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@
# specific language governing permissions and limitations
# under the License.
import os
from typing import Any, Callable, Iterable, Optional
import warnings
from typing import Any, Callable, Iterable, Optional, overload

import trino
from trino.exceptions import DatabaseError
Expand Down Expand Up @@ -106,27 +107,93 @@ def get_isolation_level(self) -> Any:
def _strip_sql(sql: str) -> str:
return sql.strip().rstrip(';')

def get_records(self, hql: str, parameters: Optional[dict] = None):
"""Get a set of records from Trino"""
@overload
def get_records(self, sql: str = "", parameters: Optional[dict] = None):
"""Get a set of records from Trino

:param sql: SQL statement to be executed.
:param parameters: The parameters to render the SQL query with.
"""

@overload
def get_records(self, sql: str = "", parameters: Optional[dict] = None, hql: str = ""):
""":sphinx-autoapi-skip:"""

def get_records(self, sql: str = "", parameters: Optional[dict] = None, hql: str = ""):
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

try:
return super().get_records(self._strip_sql(hql), parameters)
return super().get_records(self._strip_sql(sql), parameters)
except DatabaseError as e:
raise TrinoException(e)

def get_first(self, hql: str, parameters: Optional[dict] = None) -> Any:
"""Returns only the first row, regardless of how many rows the query returns."""
@overload
def get_first(self, sql: str = "", parameters: Optional[dict] = None) -> Any:
"""Returns only the first row, regardless of how many rows the query returns.

:param sql: SQL statement to be executed.
:param parameters: The parameters to render the SQL query with.
"""

@overload
def get_first(self, sql: str = "", parameters: Optional[dict] = None, hql: str = "") -> Any:
""":sphinx-autoapi-skip:"""

def get_first(self, sql: str = "", parameters: Optional[dict] = None, hql: str = "") -> Any:
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

try:
return super().get_first(self._strip_sql(hql), parameters)
return super().get_first(self._strip_sql(sql), parameters)
except DatabaseError as e:
raise TrinoException(e)

def get_pandas_df(self, hql: str, parameters: Optional[dict] = None, **kwargs): # type: ignore[override]
"""Get a pandas dataframe from a sql query."""
@overload
def get_pandas_df(
self, sql: str = "", parameters: Optional[dict] = None, **kwargs
): # type: ignore[override]
"""Get a pandas dataframe from a sql query.

:param sql: SQL statement to be executed.
:param parameters: The parameters to render the SQL query with.
"""

@overload
def get_pandas_df(
self, sql: str = "", parameters: Optional[dict] = None, hql: str = "", **kwargs
): # type: ignore[override]
""":sphinx-autoapi-skip:"""

def get_pandas_df(
self, sql: str = "", parameters: Optional[dict] = None, hql: str = "", **kwargs
): # type: ignore[override]
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

import pandas

cursor = self.get_cursor()
try:
cursor.execute(self._strip_sql(hql), parameters)
cursor.execute(self._strip_sql(sql), parameters)
data = cursor.fetchall()
except DatabaseError as e:
raise TrinoException(e)
Expand All @@ -138,16 +205,46 @@ def get_pandas_df(self, hql: str, parameters: Optional[dict] = None, **kwargs):
df = pandas.DataFrame(**kwargs)
return df

@overload
def run(
self,
hql: str,
sql: str = "",
autocommit: bool = False,
parameters: Optional[dict] = None,
handler: Optional[Callable] = None,
) -> None:
"""Execute the statement against Trino. Can be used to create views."""

@overload
def run(
self,
sql: str = "",
autocommit: bool = False,
parameters: Optional[dict] = None,
handler: Optional[Callable] = None,
hql: str = "",
) -> None:
""":sphinx-autoapi-skip:"""

def run(
self,
sql: str = "",
autocommit: bool = False,
parameters: Optional[dict] = None,
handler: Optional[Callable] = None,
hql: str = "",
) -> None:
""":sphinx-autoapi-skip:"""
if hql:
warnings.warn(
"The hql parameter has been deprecated. You should pass the sql parameter.",
DeprecationWarning,
stacklevel=2,
)
sql = hql

return super().run(
sql=self._strip_sql(hql), autocommit=autocommit, parameters=parameters, handler=handler
sql=self._strip_sql(sql), autocommit=autocommit, parameters=parameters, handler=handler
)

def insert_rows(
Expand Down