Skip to content

Commit

Permalink
fix: Set spec to b"{}" on empty input (#174)
Browse files Browse the repository at this point in the history
  • Loading branch information
candiduslynx committed May 24, 2024
1 parent 8de4658 commit 3f46922
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 2 deletions.
14 changes: 14 additions & 0 deletions cloudquery/sdk/internal/memdb/memdb.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,22 @@
import json

from cloudquery.sdk import plugin
from cloudquery.sdk import message
from cloudquery.sdk import schema
from typing import List, Generator, Dict
import pyarrow as pa
from cloudquery.sdk.types import JSONType
from dataclasses import dataclass, field

NAME = "memdb"
VERSION = "development"


@dataclass
class Spec:
abc: str = field(default="abc")


class MemDB(plugin.Plugin):
def __init__(self) -> None:
super().__init__(
Expand Down Expand Up @@ -73,6 +81,12 @@ def __init__(self) -> None:
),
}

def init(self, spec, no_connection: bool = False):
if no_connection:
return
self._spec_json = json.loads(spec)
self._spec = Spec(**self._spec_json)

def get_tables(self, options: plugin.TableOptions = None) -> List[plugin.Table]:
tables = list(self._tables.values())
return schema.filter_dfs(tables, options.tables, options.skip_tables)
Expand Down
8 changes: 7 additions & 1 deletion cloudquery/sdk/internal/servers/plugin_v3/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,9 @@ def GetSpecSchema(self, request, context):
return plugin_pb2.GetSpecSchema.Response(json_schema=self._plugin.json_schema())

def Init(self, request: plugin_pb2.Init.Request, context):
self._plugin.init(request.spec, no_connection=request.no_connection)
self._plugin.init(
sanitize_spec(request.spec), no_connection=request.no_connection
)
return plugin_pb2.Init.Response()

def GetTables(self, request: plugin_pb2.GetTables.Request, context):
Expand Down Expand Up @@ -113,3 +115,7 @@ def msg_iterator() -> Generator[WriteMessage, None, None]:
def Close(self, request, context):
self._plugin.close()
return plugin_pb2.Close.Response()


def sanitize_spec(spec=None):
return b"{}" if spec is None or spec == b"" or spec == b"null" else spec
3 changes: 2 additions & 1 deletion tests/internal/memdb/memdb.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,11 @@
from cloudquery.sdk.internal import memdb
from cloudquery.sdk.internal.servers.plugin_v3 import plugin
from cloudquery.sdk.plugin import SyncOptions


def test_memdb():
p = memdb.MemDB()
p.init(None)
p.init(plugin.sanitize_spec(b"null"))
msgs = []
for msg in p.sync(SyncOptions(tables=["*"])):
msgs.append(msg)
Expand Down

0 comments on commit 3f46922

Please sign in to comment.