Source code for airflow.providers.arangodb.hooks.arangodb

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""This module allows connecting to a ArangoDB."""

from __future__ import annotations

from functools import cached_property
from typing import TYPE_CHECKING, Any

from arango import AQLQueryExecuteError, ArangoClient as ArangoDBClient

from airflow.exceptions import AirflowException
from airflow.hooks.base import BaseHook

if TYPE_CHECKING:
    from arango.cursor import Cursor
    from arango.database import StandardDatabase

    from airflow.models import Connection


[docs]class ArangoDBHook(BaseHook):
    """
    Interact with ArangoDB.

    Performs a connection to ArangoDB and retrieves client.

    :param arangodb_conn_id: Reference to :ref:`ArangoDB connection id <howto/connection:arangodb>`.
    """

[docs]    conn_name_attr = "arangodb_conn_id"
[docs]    default_conn_name = "arangodb_default"
[docs]    conn_type = "arangodb"
[docs]    hook_name = "ArangoDB"

    def __init__(self, arangodb_conn_id: str = default_conn_name, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self.arangodb_conn_id = arangodb_conn_id

    @cached_property
[docs]    def client(self) -> ArangoDBClient:
        """Initiates a new ArangoDB connection (cached)."""
        return ArangoDBClient(hosts=self.hosts)

    @cached_property
[docs]    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
[docs]    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
[docs]    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
[docs]    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
[docs]    def password(self) -> str:
        return self._conn.password or ""

[docs]    def get_conn(self) -> ArangoDBClient:
        """Initiate a new ArangoDB connection (cached)."""
        return self.client

[docs]    def query(self, query, **kwargs) -> Cursor:
        """
        Create an ArangoDB session and execute the AQL query in the session.

        :param query: AQL query
        """
        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(
                    f"Failed to execute AQLQuery, error connecting to database: {self.database}"
                )
        except AQLQueryExecuteError as error:
            raise AirflowException(f"Failed to execute AQLQuery, error: {error}")

[docs]    def create_collection(self, name):
        if not self.db_conn.has_collection(name):
            self.db_conn.create_collection(name)
            return True
        else:
            self.log.info("Collection already exists: %s", name)
            return False

[docs]    def create_database(self, name):
        if not self.db_conn.has_database(name):
            self.db_conn.create_database(name)
            return True
        else:
            self.log.info("Database already exists: %s", name)
            return False

[docs]    def create_graph(self, name):
        if not self.db_conn.has_graph(name):
            self.db_conn.create_graph(name)
            return True
        else:
            self.log.info("Graph already exists: %s", name)
            return False

    @classmethod
[docs]    def get_ui_field_behaviour(cls) -> dict[str, Any]:
        return {
            "hidden_fields": ["port", "extra"],
            "relabeling": {
                "host": "ArangoDB Host URL or  comma separated list of URLs (coordinators in a cluster)",
                "schema": "ArangoDB Database",
                "login": "ArangoDB Username",
                "password": "ArangoDB Password",
            },
            "placeholders": {
                "host": 'eg."http://127.0.0.1:8529" or "http://127.0.0.1:8529,http://127.0.0.1:8530"'
                " (coordinators in a cluster)",
                "schema": "_system",
                "login": "root",
                "password": "password",
            },
        }