Add discovery server

Author: hermanschaaf

cloudquery/sdk/internal/servers/discovery_v1/__init__.py

@@ -0,0 +1,11 @@
+from typing import List
+
+from cloudquery.discovery_v1 import discovery_pb2, discovery_pb2_grpc
+
+
+class DiscoveryServicer(discovery_pb2_grpc.DiscoveryServicer):
+    def __init__(self, versions: List[int]):
+        self._versions = versions
+
+    def GetVersions(self, request, context):
+        return discovery_pb2.GetVersions.Response(versions=self._versions)

cloudquery/sdk/internal/servers/discovery_v1/discovery.py

@@ -1,17 +1,13 @@
+import pyarrow as pa
 
-import grpc
-from concurrent import futures
-import logging
 from cloudquery.plugin_v3 import plugin_pb2, plugin_pb2_grpc
 from cloudquery.sdk.plugin.plugin import Plugin
 from cloudquery.sdk.schema import tables_to_arrow_schemas
-import pyarrow as pa
 
 
 class PluginServicer(plugin_pb2_grpc.PluginServicer):
     def __init__(self, plugin: Plugin):
         self._plugin = plugin
-        pass
 
     def GetName(self, request, context):
         return plugin_pb2.GetName.Response(name=self._plugin.name())
@@ -28,13 +24,13 @@ def GetTables(self, request: plugin_pb2.GetTables.Request, context):
         schema = tables_to_arrow_schemas(tables)
         tablesBytes = []
         for s in schema:
-          sink = pa.BufferOutputStream()
-          writer = pa.ipc.new_stream(sink, s)
-          writer.close()
-          buf = sink.getvalue().to_pybytes()
-          tablesBytes.append(buf)
+            sink = pa.BufferOutputStream()
+            writer = pa.ipc.new_stream(sink, s)
+            writer.close()
+            buf = sink.getvalue().to_pybytes()
+            tablesBytes.append(buf)
 
-        return plugin_pb2.GetTables.Response(tables=[])
+        return plugin_pb2.GetTables.Response(tables=tablesBytes)
 
     def Sync(self, request, context):
         plugin_pb2.Sync.Response()

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

@@ -4,8 +4,10 @@
 import grpc
 import sys
 
+from cloudquery.discovery_v1 import discovery_pb2_grpc
 from cloudquery.plugin_v3 import plugin_pb2_grpc
 from cloudquery.sdk.docs.generator import Generator
+from cloudquery.sdk.internal.servers.discovery_v1.discovery import DiscoveryServicer
 from cloudquery.sdk.internal.servers.plugin_v3 import PluginServicer
 from cloudquery.sdk.plugin.plugin import Plugin
 
@@ -56,6 +58,8 @@ def run(self, args):
 
     def _serve(self, args):
         server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
+        discovery_pb2_grpc.add_DiscoveryServicer_to_server(
+            DiscoveryServicer([3]), server)
         plugin_pb2_grpc.add_PluginServicer_to_server(
             PluginServicer(self._plugin), server)
         server.add_insecure_port(args.address)
@@ -65,5 +69,5 @@ def _serve(self, args):
 
     def _generate_docs(self, args):
         print("Generating docs in format: " + args.format)
-        generator = Generator(self._plugin.name(), self._plugin.get_tables([], []))
+        generator = Generator(self._plugin.name(), self._plugin.get_tables(tables=["*"], skip_tables=[]))
         generator.generate(args.directory, args.format)