Source code for

# 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
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Hook for Level DB"""
from typing import List, Optional

import plyvel
from plyvel import DB

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

[docs]class LevelDBHookException(AirflowException):
"""Exception specific for LevelDB"""
[docs]class LevelDBHook(BaseHook): """ Plyvel Wrapper to Interact With LevelDB Database `LevelDB Connection Documentation <>`__ """
[docs] conn_name_attr = 'leveldb_conn_id'
[docs] default_conn_name = 'leveldb_default'
[docs] conn_type = 'leveldb'
[docs] hook_name = 'LevelDB'
def __init__(self, leveldb_conn_id: str = default_conn_name): super().__init__() self.leveldb_conn_id = leveldb_conn_id self.connection = self.get_connection(leveldb_conn_id) self.db = None
[docs] def get_conn(self, name: str = '/tmp/testdb/', create_if_missing: bool = False, **kwargs) -> DB: """ Creates `Plyvel DB <>`__ :param name: path to create database e.g. `/tmp/testdb/`) :type name: str :param create_if_missing: whether a new database should be created if needed :type create_if_missing: bool :param kwargs: other options of creation plyvel.DB. See more in the link above. :type kwargs: Dict[str, Any] :returns: DB :rtype: plyvel.DB """ if self.db is not None: return self.db self.db = plyvel.DB(name=name, create_if_missing=create_if_missing, **kwargs) return self.db
[docs] def close_conn(self) -> None: """Closes connection""" db = self.db if db is not None: db.close() self.db = None
[docs] def run( self, command: str, key: bytes, value: bytes = None, keys: List[bytes] = None, values: List[bytes] = None, ) -> Optional[bytes]: """ Execute operation with leveldb :param command: command of plyvel(python wrap for leveldb) for DB object e.g. ``"put"``, ``"get"``, ``"delete"``, ``"write_batch"``. :type command: str :param key: key for command(put,get,delete) execution(, e.g. ``b'key'``, ``b'another-key'``) :type key: bytes :param value: value for command(put) execution(bytes, e.g. ``b'value'``, ``b'another-value'``) :type value: bytes :param keys: keys for command(write_batch) execution(List[bytes], e.g. ``[b'key', b'another-key'])`` :type keys: List[bytes] :param values: values for command(write_batch) execution e.g. ``[b'value'``, ``b'another-value']`` :type values: List[bytes] :returns: value from get or None :rtype: Optional[bytes] """ if command == 'put': return self.put(key, value) elif command == 'get': return self.get(key) elif command == 'delete': return self.delete(key) elif command == 'write_batch': return self.write_batch(keys, values) else: raise LevelDBHookException("Unknown command for LevelDB hook")
[docs] def put(self, key: bytes, value: bytes): """ Put a single value into a leveldb db by key :param key: key for put execution, e.g. ``b'key'``, ``b'another-key'`` :type key: bytes :param value: value for put execution e.g. ``b'value'``, ``b'another-value'`` :type value: bytes """ self.db.put(key, value)
[docs] def get(self, key: bytes) -> bytes: """ Get a single value into a leveldb db by key :param key: key for get execution, e.g. ``b'key'``, ``b'another-key'`` :type key: bytes :returns: value of key from db.get :rtype: bytes """ return self.db.get(key)
[docs] def delete(self, key: bytes): """ Delete a single value in a leveldb db by key. :param key: key for delete execution, e.g. ``b'key'``, ``b'another-key'`` :type key: bytes """ self.db.delete(key)
[docs] def write_batch(self, keys: List[bytes], values: List[bytes]): """ Write batch of values in a leveldb db by keys :param keys: keys for write_batch execution e.g. ``[b'key', b'another-key']`` :type keys: List[bytes] :param values: values for write_batch execution e.g. ``[b'value', b'another-value']`` :type values: List[bytes] """ with self.db.write_batch() as batch: for i, key in enumerate(keys): batch.put(key, values[i])

Was this entry helpful?