Source code for airflow.providers.google.leveldb.operators.leveldb

# 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.
from __future__ import annotations

from typing import TYPE_CHECKING, Any

from airflow.models import BaseOperator
from airflow.providers.google.leveldb.hooks.leveldb import LevelDBHook

if TYPE_CHECKING:
    from airflow.utils.context import Context


[docs]class LevelDBOperator(BaseOperator):
    """
    Execute command in LevelDB.

    .. seealso::
        For more information on how to use this operator, take a look at the guide:
        :ref:`howto/operator:LevelDBOperator`

        :param command: command of plyvel(python wrap for leveldb) for DB object e.g.
            ``"put"``, ``"get"``, ``"delete"``, ``"write_batch"``.
        :param key: key for command(put,get,delete) execution(, e.g. ``b'key'``, ``b'another-key'``)
        :param value: value for command(put) execution(bytes, e.g. ``b'value'``, ``b'another-value'``)
        :param keys: keys for command(write_batch) execution(list[bytes], e.g. ``[b'key', b'another-key'])``
        :param values: values for command(write_batch) execution e.g. ``[b'value'``, ``b'another-value']``
        :param leveldb_conn_id:
        :param create_if_missing: whether a new database should be created if needed
        :param create_db_extra_options: extra options of creation LevelDBOperator. See more in the link below
            `Plyvel DB <https://plyvel.readthedocs.io/en/latest/api.html#DB>`__
    """

    def __init__(
        self,
        *,
        command: str,
        key: bytes,
        value: bytes | None = None,
        keys: list[bytes] | None = None,
        values: list[bytes] | None = None,
        leveldb_conn_id: str = "leveldb_default",
        name: str = "/tmp/testdb/",
        create_if_missing: bool = True,
        create_db_extra_options: dict[str, Any] | None = None,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        self.command = command
        self.key = key
        self.value = value
        self.keys = keys
        self.values = values
        self.leveldb_conn_id = leveldb_conn_id
        self.name = name
        self.create_if_missing = create_if_missing
        self.create_db_extra_options = create_db_extra_options or {}

[docs]    def execute(self, context: Context) -> str | None:
        """
        Execute command in LevelDB.

        :returns: value from get(str, not bytes, to prevent error in json.dumps in serialize_value in xcom.py)
            or str | None
        """
        leveldb_hook = LevelDBHook(leveldb_conn_id=self.leveldb_conn_id)
        leveldb_hook.get_conn(
            name=self.name, create_if_missing=self.create_if_missing, **self.create_db_extra_options
        )
        value = leveldb_hook.run(
            command=self.command,
            key=self.key,
            value=self.value,
            keys=self.keys,
            values=self.values,
        )
        self.log.info("Done. Returned value was: %s", value)
        leveldb_hook.close_conn()
        str_value = value if value is None else value.decode()
        return str_value