Skip to content

Commit 3f46922

Browse files
authored
fix: Set spec to b"{}" on empty input (#174)
1 parent 8de4658 commit 3f46922

File tree

3 files changed

+23
-2
lines changed

3 files changed

+23
-2
lines changed

cloudquery/sdk/internal/memdb/memdb.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,22 @@
1+
import json
2+
13
from cloudquery.sdk import plugin
24
from cloudquery.sdk import message
35
from cloudquery.sdk import schema
46
from typing import List, Generator, Dict
57
import pyarrow as pa
68
from cloudquery.sdk.types import JSONType
9+
from dataclasses import dataclass, field
710

811
NAME = "memdb"
912
VERSION = "development"
1013

1114

15+
@dataclass
16+
class Spec:
17+
abc: str = field(default="abc")
18+
19+
1220
class MemDB(plugin.Plugin):
1321
def __init__(self) -> None:
1422
super().__init__(
@@ -73,6 +81,12 @@ def __init__(self) -> None:
7381
),
7482
}
7583

84+
def init(self, spec, no_connection: bool = False):
85+
if no_connection:
86+
return
87+
self._spec_json = json.loads(spec)
88+
self._spec = Spec(**self._spec_json)
89+
7690
def get_tables(self, options: plugin.TableOptions = None) -> List[plugin.Table]:
7791
tables = list(self._tables.values())
7892
return schema.filter_dfs(tables, options.tables, options.skip_tables)

cloudquery/sdk/internal/servers/plugin_v3/plugin.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,9 @@ def GetSpecSchema(self, request, context):
3232
return plugin_pb2.GetSpecSchema.Response(json_schema=self._plugin.json_schema())
3333

3434
def Init(self, request: plugin_pb2.Init.Request, context):
35-
self._plugin.init(request.spec, no_connection=request.no_connection)
35+
self._plugin.init(
36+
sanitize_spec(request.spec), no_connection=request.no_connection
37+
)
3638
return plugin_pb2.Init.Response()
3739

3840
def GetTables(self, request: plugin_pb2.GetTables.Request, context):
@@ -113,3 +115,7 @@ def msg_iterator() -> Generator[WriteMessage, None, None]:
113115
def Close(self, request, context):
114116
self._plugin.close()
115117
return plugin_pb2.Close.Response()
118+
119+
120+
def sanitize_spec(spec=None):
121+
return b"{}" if spec is None or spec == b"" or spec == b"null" else spec

tests/internal/memdb/memdb.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
11
from cloudquery.sdk.internal import memdb
2+
from cloudquery.sdk.internal.servers.plugin_v3 import plugin
23
from cloudquery.sdk.plugin import SyncOptions
34

45

56
def test_memdb():
67
p = memdb.MemDB()
7-
p.init(None)
8+
p.init(plugin.sanitize_spec(b"null"))
89
msgs = []
910
for msg in p.sync(SyncOptions(tables=["*"])):
1011
msgs.append(msg)

0 commit comments

Comments
 (0)