Skip to content

Commit 35c9a90

Browse files
radim94RosterInKamil Bregulaephraimbuddy
authored
Add Google leveldb hook and operator (#13109) (#14105)
* Add Google leveldb hook (#13109) * Add write_batch, options for DB creation(comparator and other) plus fixes (#13109) * Fix some static checks, add docs (#13109) * Apply suggestions from code review Co-authored-by: RosterIn <48057736+RosterIn@users.noreply.github.com> * Fix some otger static checks and docs (#13109) * Fix tests and some build-docs checks (#13109) * Apply suggestions from code review Co-authored-by: RosterIn <48057736+RosterIn@users.noreply.github.com> * Fix build-docs checks (#13109) * Fix build-docs checks (#13109) * fixup! Fix build-docs checks (#13109) * Rewrite example dag as in google package (#13109) * Add extra options in operator, fix docstrings (#13109) * Update airflow/providers/google/leveldb/operators/leveldb.py Co-authored-by: Ephraim Anierobi <splendidzigy24@gmail.com> * Add system testing and docstrings (#13109) * Fix comparator place in spelling wordlist(#13109) Co-authored-by: RosterIn <48057736+RosterIn@users.noreply.github.com> Co-authored-by: Kamil Bregula <kamilbregula@Kamils-MacBook-Pro.local> Co-authored-by: Ephraim Anierobi <splendidzigy24@gmail.com>
1 parent e273366 commit 35c9a90

File tree

22 files changed

+723
-0
lines changed

22 files changed

+723
-0
lines changed

β€Žairflow/providers/google/leveldb/__init__.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.

β€Žairflow/providers/google/leveldb/example_dags/__init__.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.

β€Žairflow/providers/google/leveldb/example_dags/example_leveldb.py

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,52 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
"""
19+
Example use of LevelDB operators.
20+
"""
21+
22+
from airflow import models
23+
from airflow.providers.google.leveldb.operators.leveldb import LevelDBOperator
24+
from airflow.utils.dates import days_ago
25+
26+
default_args = {
27+
'owner': 'airflow',
28+
}
29+
30+
with models.DAG(
31+
'example_leveldb',
32+
default_args=default_args,
33+
start_date=days_ago(2),
34+
schedule_interval=None,
35+
tags=['example'],
36+
) as dag:
37+
# [START howto_operator_leveldb_get_key]
38+
get_key_leveldb_task = LevelDBOperator(
39+
task_id='get_key_leveldb', leveldb_conn_id='leveldb_default', command='get', key=b'key', dag=dag
40+
)
41+
# [END howto_operator_leveldb_get_key]
42+
# [START howto_operator_leveldb_put_key]
43+
put_key_leveldb_task = LevelDBOperator(
44+
task_id='put_key_leveldb',
45+
leveldb_conn_id='leveldb_default',
46+
command='put',
47+
key=b'another_key',
48+
value=b'another_value',
49+
dag=dag,
50+
)
51+
# [END howto_operator_leveldb_put_key]
52+
get_key_leveldb_task >> put_key_leveldb_task

β€Žairflow/providers/google/leveldb/hooks/__init__.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.

β€Žairflow/providers/google/leveldb/hooks/leveldb.py

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""Hook for Level DB"""
18+
from typing import List, Optional
19+
20+
import plyvel
21+
from plyvel import DB
22+
23+
from airflow.exceptions import AirflowException
24+
from airflow.hooks.base import BaseHook
25+
26+
27+
class LevelDBHookException(AirflowException):
28+
"""Exception specific for LevelDB"""
29+
30+
31+
class LevelDBHook(BaseHook):
32+
"""
33+
Plyvel Wrapper to Interact With LevelDB Database
34+
`LevelDB Connection Documentation <https://plyvel.readthedocs.io/en/latest/>`__
35+
"""
36+
37+
conn_name_attr = 'leveldb_conn_id'
38+
default_conn_name = 'leveldb_default'
39+
conn_type = 'leveldb'
40+
hook_name = 'LevelDB'
41+
42+
def __init__(self, leveldb_conn_id: str = default_conn_name):
43+
super().__init__()
44+
self.leveldb_conn_id = leveldb_conn_id
45+
self.connection = self.get_connection(leveldb_conn_id)
46+
self.db = None
47+
48+
def get_conn(self, name: str = '/tmp/testdb/', create_if_missing: bool = False, **kwargs) -> DB:
49+
"""
50+
Creates `Plyvel DB <https://plyvel.readthedocs.io/en/latest/api.html#DB>`__
51+
52+
:param name: path to create database e.g. `/tmp/testdb/`)
53+
:type name: str
54+
:param create_if_missing: whether a new database should be created if needed
55+
:type create_if_missing: bool
56+
:param kwargs: other options of creation plyvel.DB. See more in the link above.
57+
:type kwargs: Dict[str, Any]
58+
:returns: DB
59+
:rtype: plyvel.DB
60+
"""
61+
if self.db is not None:
62+
return self.db
63+
self.db = plyvel.DB(name=name, create_if_missing=create_if_missing, **kwargs)
64+
return self.db
65+
66+
def close_conn(self) -> None:
67+
"""Closes connection"""
68+
db = self.db
69+
if db is not None:
70+
db.close()
71+
self.db = None
72+
73+
def run(
74+
self,
75+
command: str,
76+
key: bytes,
77+
value: bytes = None,
78+
keys: List[bytes] = None,
79+
values: List[bytes] = None,
80+
) -> Optional[bytes]:
81+
"""
82+
Execute operation with leveldb
83+
84+
:param command: command of plyvel(python wrap for leveldb) for DB object e.g.
85+
``"put"``, ``"get"``, ``"delete"``, ``"write_batch"``.
86+
:type command: str
87+
:param key: key for command(put,get,delete) execution(, e.g. ``b'key'``, ``b'another-key'``)
88+
:type key: bytes
89+
:param value: value for command(put) execution(bytes, e.g. ``b'value'``, ``b'another-value'``)
90+
:type value: bytes
91+
:param keys: keys for command(write_batch) execution(List[bytes], e.g. ``[b'key', b'another-key'])``
92+
:type keys: List[bytes]
93+
:param values: values for command(write_batch) execution e.g. ``[b'value'``, ``b'another-value']``
94+
:type values: List[bytes]
95+
:returns: value from get or None
96+
:rtype: Optional[bytes]
97+
"""
98+
if command == 'put':
99+
return self.put(key, value)
100+
elif command == 'get':
101+
return self.get(key)
102+
elif command == 'delete':
103+
return self.delete(key)
104+
elif command == 'write_batch':
105+
return self.write_batch(keys, values)
106+
else:
107+
raise LevelDBHookException("Unknown command for LevelDB hook")
108+
109+
def put(self, key: bytes, value: bytes):
110+
"""
111+
Put a single value into a leveldb db by key
112+
113+
:param key: key for put execution, e.g. ``b'key'``, ``b'another-key'``
114+
:type key: bytes
115+
:param value: value for put execution e.g. ``b'value'``, ``b'another-value'``
116+
:type value: bytes
117+
"""
118+
self.db.put(key, value)
119+
120+
def get(self, key: bytes) -> bytes:
121+
"""
122+
Get a single value into a leveldb db by key
123+
124+
:param key: key for get execution, e.g. ``b'key'``, ``b'another-key'``
125+
:type key: bytes
126+
:returns: value of key from db.get
127+
:rtype: bytes
128+
"""
129+
return self.db.get(key)
130+
131+
def delete(self, key: bytes):
132+
"""
133+
Delete a single value in a leveldb db by key.
134+
135+
:param key: key for delete execution, e.g. ``b'key'``, ``b'another-key'``
136+
:type key: bytes
137+
"""
138+
self.db.delete(key)
139+
140+
def write_batch(self, keys: List[bytes], values: List[bytes]):
141+
"""
142+
Write batch of values in a leveldb db by keys
143+
144+
:param keys: keys for write_batch execution e.g. ``[b'key', b'another-key']``
145+
:type keys: List[bytes]
146+
:param values: values for write_batch execution e.g. ``[b'value', b'another-value']``
147+
:type values: List[bytes]
148+
"""
149+
with self.db.write_batch() as batch:
150+
for i, key in enumerate(keys):
151+
batch.put(key, values[i])

β€Žairflow/providers/google/leveldb/operators/__init__.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.

β€Žairflow/providers/google/leveldb/operators/leveldb.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,100 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
from typing import Any, Dict, List, Optional
18+
19+
from airflow.models import BaseOperator
20+
from airflow.providers.google.leveldb.hooks.leveldb import LevelDBHook
21+
from airflow.utils.decorators import apply_defaults
22+
23+
24+
class LevelDBOperator(BaseOperator):
25+
"""
26+
Execute command in LevelDB
27+
28+
.. seealso::
29+
For more information on how to use this operator, take a look at the guide:
30+
:ref:`howto/operator:LevelDBOperator`
31+
32+
:param command: command of plyvel(python wrap for leveldb) for DB object e.g.
33+
``"put"``, ``"get"``, ``"delete"``, ``"write_batch"``.
34+
:type command: str
35+
:param key: key for command(put,get,delete) execution(, e.g. ``b'key'``, ``b'another-key'``)
36+
:type key: bytes
37+
:param value: value for command(put) execution(bytes, e.g. ``b'value'``, ``b'another-value'``)
38+
:type value: bytes
39+
:param keys: keys for command(write_batch) execution(List[bytes], e.g. ``[b'key', b'another-key'])``
40+
:type keys: List[bytes]
41+
:param values: values for command(write_batch) execution e.g. ``[b'value'``, ``b'another-value']``
42+
:type values: List[bytes]
43+
:param leveldb_conn_id:
44+
:type leveldb_conn_id: str
45+
:param create_if_missing: whether a new database should be created if needed
46+
:type create_if_missing: bool
47+
:param create_db_extra_options: extra options of creation LevelDBOperator. See more in the link below
48+
`Plyvel DB <https://plyvel.readthedocs.io/en/latest/api.html#DB>`__
49+
:type create_db_extra_options: Optional[Dict[str, Any]]
50+
"""
51+
52+
@apply_defaults
53+
def __init__(
54+
self,
55+
*,
56+
command: str,
57+
key: bytes,
58+
value: bytes = None,
59+
keys: List[bytes] = None,
60+
values: List[bytes] = None,
61+
leveldb_conn_id: str = 'leveldb_default',
62+
name: str = '/tmp/testdb/',
63+
create_if_missing: bool = True,
64+
create_db_extra_options: Optional[Dict[str, Any]] = None,
65+
**kwargs,
66+
) -> None:
67+
super().__init__(**kwargs)
68+
self.command = command
69+
self.key = key
70+
self.value = value
71+
self.keys = keys
72+
self.values = values
73+
self.leveldb_conn_id = leveldb_conn_id
74+
self.name = name
75+
self.create_if_missing = create_if_missing
76+
self.create_db_extra_options = create_db_extra_options or {}
77+
78+
def execute(self, context) -> Optional[str]:
79+
"""
80+
Execute command in LevelDB
81+
82+
:returns: value from get(str, not bytes, to prevent error in json.dumps in serialize_value in xcom.py)
83+
or None(Optional[str])
84+
:rtype: Optional[str]
85+
"""
86+
leveldb_hook = LevelDBHook(leveldb_conn_id=self.leveldb_conn_id)
87+
leveldb_hook.get_conn(
88+
name=self.name, create_if_missing=self.create_if_missing, **self.create_db_extra_options
89+
)
90+
value = leveldb_hook.run(
91+
command=self.command,
92+
key=self.key,
93+
value=self.value,
94+
keys=self.keys,
95+
values=self.values,
96+
)
97+
self.log.info("Done. Returned value was: %s", str(value))
98+
leveldb_hook.close_conn()
99+
value = value if value is None else value.decode()
100+
return value

0 commit comments

Comments
 (0)