Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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