You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by el...@apache.org on 2023/01/10 05:51:14 UTC
[airflow] branch main updated: Fix ArangoDB static checks (#28820)
This is an automated email from the ASF dual-hosted git repository.
eladkal pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 1e9c8e52fd Fix ArangoDB static checks (#28820)
1e9c8e52fd is described below
commit 1e9c8e52fda95a0a30b3ae298d5d3adc1971ed45
Author: Andrey Anshin <An...@taragol.is>
AuthorDate: Tue Jan 10 09:51:04 2023 +0400
Fix ArangoDB static checks (#28820)
---
airflow/providers/arangodb/hooks/arangodb.py | 73 ++++++++++++++++--------
airflow/providers/arangodb/operators/arangodb.py | 1 -
airflow/providers/arangodb/sensors/arangodb.py | 1 -
tests/providers/arangodb/hooks/test_arangodb.py | 13 ++---
4 files changed, 55 insertions(+), 33 deletions(-)
diff --git a/airflow/providers/arangodb/hooks/arangodb.py b/airflow/providers/arangodb/hooks/arangodb.py
index a30583e837..2f51ac0618 100644
--- a/airflow/providers/arangodb/hooks/arangodb.py
+++ b/airflow/providers/arangodb/hooks/arangodb.py
@@ -18,14 +18,20 @@
"""This module allows connecting to a ArangoDB."""
from __future__ import annotations
-from typing import Any
+from typing import TYPE_CHECKING, Any
from arango import AQLQueryExecuteError, ArangoClient as ArangoDBClient
-from arango.result import Result
from airflow import AirflowException
+from airflow.compat.functools import cached_property
from airflow.hooks.base import BaseHook
+if TYPE_CHECKING:
+ from arango.cursor import Cursor
+ from arango.database import StandardDatabase
+
+ from airflow.models import Connection
+
class ArangoDBHook(BaseHook):
"""
@@ -43,41 +49,60 @@ class ArangoDBHook(BaseHook):
def __init__(self, arangodb_conn_id: str = default_conn_name, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
- self.hosts = None
- self.database = None
- self.username = None
- self.password = None
- self.db_conn = None
self.arangodb_conn_id = arangodb_conn_id
- self.client: ArangoDBClient | None = None
- self.get_conn()
+
+ @cached_property
+ def client(self) -> ArangoDBClient:
+ """Initiates a new ArangoDB connection (cached)."""
+ return ArangoDBClient(hosts=self.hosts)
+
+ @cached_property
+ def db_conn(self) -> StandardDatabase:
+ """Connect to an ArangoDB database and return the database API wrapper."""
+ return self.client.db(name=self.database, username=self.username, password=self.password)
+
+ @cached_property
+ def _conn(self) -> Connection:
+ return self.get_connection(self.arangodb_conn_id)
+
+ @property
+ def hosts(self) -> list[str]:
+ if not self._conn.host:
+ raise AirflowException(f"No ArangoDB Host(s) provided in connection: {self.arangodb_conn_id!r}.")
+ return self._conn.host.split(",")
+
+ @property
+ def database(self) -> str:
+ if not self._conn.schema:
+ raise AirflowException(f"No ArangoDB Database provided in connection: {self.arangodb_conn_id!r}.")
+ return self._conn.schema
+
+ @property
+ def username(self) -> str:
+ if not self._conn.login:
+ raise AirflowException(f"No ArangoDB Username provided in connection: {self.arangodb_conn_id!r}.")
+ return self._conn.login
+
+ @property
+ def password(self) -> str:
+ return self._conn.password or ""
def get_conn(self) -> ArangoDBClient:
- """Function that initiates a new ArangoDB connection"""
- if self.client is not None:
- return self.client
-
- conn = self.get_connection(self.arangodb_conn_id)
- self.hosts = conn.host.split(",")
- self.database = conn.schema
- self.username = conn.login
- self.password = conn.password
-
- self.client = ArangoDBClient(hosts=self.hosts)
- self.db_conn = self.client.db(name=self.database, username=self.username, password=self.password)
+ """Function that initiates a new ArangoDB connection (cached)."""
return self.client
- def query(self, query, **kwargs) -> Result:
+ def query(self, query, **kwargs) -> Cursor:
"""
- Function to create a arangodb session
+ Function to create an ArangoDB session
and execute the AQL query in the session.
:param query: AQL query
- :return: Result
"""
try:
if self.db_conn:
result = self.db_conn.aql.execute(query, **kwargs)
+ if TYPE_CHECKING:
+ assert isinstance(result, Cursor)
return result
else:
raise AirflowException(
diff --git a/airflow/providers/arangodb/operators/arangodb.py b/airflow/providers/arangodb/operators/arangodb.py
index 8f8c357713..8d2f4932f5 100644
--- a/airflow/providers/arangodb/operators/arangodb.py
+++ b/airflow/providers/arangodb/operators/arangodb.py
@@ -41,7 +41,6 @@ class AQLOperator(BaseOperator):
"""
template_fields: Sequence[str] = ("query",)
-
template_ext: Sequence[str] = (".sql",)
template_fields_renderers = {"query": "sql"}
diff --git a/airflow/providers/arangodb/sensors/arangodb.py b/airflow/providers/arangodb/sensors/arangodb.py
index 541cb9b38f..2d13afcce2 100644
--- a/airflow/providers/arangodb/sensors/arangodb.py
+++ b/airflow/providers/arangodb/sensors/arangodb.py
@@ -39,7 +39,6 @@ class AQLSensor(BaseSensorOperator):
"""
template_fields: Sequence[str] = ("query",)
-
template_ext: Sequence[str] = (".sql",)
template_fields_renderers = {"query": "sql"}
diff --git a/tests/providers/arangodb/hooks/test_arangodb.py b/tests/providers/arangodb/hooks/test_arangodb.py
index 707e8a59d2..27676dffbd 100644
--- a/tests/providers/arangodb/hooks/test_arangodb.py
+++ b/tests/providers/arangodb/hooks/test_arangodb.py
@@ -61,14 +61,13 @@ class TestArangoDBHook:
)
def test_query(self, arango_mock):
arangodb_hook = ArangoDBHook()
- arangodb_hook.db_conn = Mock(name="arangodb_database_for_test")
+ with patch.object(arangodb_hook, "db_conn"):
+ arangodb_query = "FOR doc IN students RETURN doc"
+ arangodb_hook.query(arangodb_query)
- arangodb_query = "FOR doc IN students RETURN doc"
- arangodb_hook.query(arangodb_query)
-
- assert arango_mock.called
- assert isinstance(arangodb_hook.client, Mock)
- assert arango_mock.return_value.db.called
+ assert arango_mock.called
+ assert isinstance(arangodb_hook.client, Mock)
+ assert arango_mock.return_value.db.called
@patch(
"airflow.providers.arangodb.hooks.arangodb.ArangoDBClient",