fix RPC issues

dbt-labs · Sep 21, 2020 · 5906486 · 5906486
1 parent 3388024
commit 5906486
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 19 deletions.
diff --git a/core/dbt/contracts/rpc.py b/core/dbt/contracts/rpc.py
@@ -175,40 +175,40 @@ class RemoteResult(VersionedSchema):
     logs: List[LogMessage]
 
 
-@dataclass
-class TimestampedResult(RemoteResult):
-    generated_at: datetime
-
-
 @dataclass
 @schema_version('remote-deps-result', 1)
-class RemoteDepsResult(TimestampedResult):
-    pass
+class RemoteDepsResult(RemoteResult):
+    generated_at: datetime = field(default_factory=datetime.utcnow)
 
 
 @dataclass
 @schema_version('remote-catalog-result', 1)
-class RemoteCatalogResults(CatalogResults, TimestampedResult):
-    pass
+class RemoteCatalogResults(CatalogResults, RemoteResult):
+    generated_at: datetime = field(default_factory=datetime.utcnow)
 
 
 @dataclass
-@schema_version('remote-compile-result', 1)
-class RemoteCompileResult(TimestampedResult):
+class RemoteCompileResultMixin(RemoteResult):
     raw_sql: str
     compiled_sql: str
     node: CompileResultNode
     timing: List[TimingInfo]
 
+
+@dataclass
+@schema_version('remote-compile-result', 1)
+class RemoteCompileResult(RemoteCompileResultMixin):
+    generated_at: datetime = field(default_factory=datetime.utcnow)
+
     @property
     def error(self):
         return None
 
 
 @dataclass
 @schema_version('remote-execution-result', 1)
-class RemoteExecutionResult(ExecutionResult, TimestampedResult):
-    pass
+class RemoteExecutionResult(ExecutionResult, RemoteResult):
+    generated_at: datetime = field(default_factory=datetime.utcnow)
 
 
 @dataclass
@@ -219,14 +219,16 @@ class ResultTable(JsonSchemaMixin):
 
 @dataclass
 @schema_version('remote-run-operation-result', 1)
-class RemoteRunOperationResult(RemoteExecutionResult):
+class RemoteRunOperationResult(ExecutionResult, RemoteResult):
     success: bool
+    generated_at: datetime = field(default_factory=datetime.utcnow)
 
 
 @dataclass
 @schema_version('remote-run-result', 1)
-class RemoteRunResult(RemoteCompileResult):
+class RemoteRunResult(RemoteCompileResultMixin):
     table: ResultTable
+    generated_at: datetime = field(default_factory=datetime.utcnow)
 
 
 RPCResult = Union[
@@ -391,11 +393,12 @@ class PollResult(RemoteResult, TaskTiming):
 
 @dataclass
 @schema_version('poll-remote-deps-result', 1)
-class PollRemoteEmptyCompleteResult(PollResult, RemoteDepsResult):
+class PollRemoteEmptyCompleteResult(PollResult, RemoteResult):
     state: TaskHandlerState = field(
         metadata=restrict_to(TaskHandlerState.Success,
                              TaskHandlerState.Failed),
     )
+    generated_at: datetime = field(default_factory=datetime.utcnow)
 
     @classmethod
     def from_result(
@@ -430,7 +433,6 @@ class PollExecuteCompleteResult(
     RemoteExecutionResult,
     PollResult,
 ):
-
     state: TaskHandlerState = field(
         metadata=restrict_to(TaskHandlerState.Success,
                              TaskHandlerState.Failed),

diff --git a/core/dbt/rpc/node_runners.py b/core/dbt/rpc/node_runners.py
@@ -4,14 +4,17 @@
 
 import dbt.exceptions
 from dbt.contracts.rpc import (
-    RemoteCompileResult, RemoteRunResult, ResultTable,
+    RemoteCompileResult,
+    RemoteCompileResultMixin,
+    RemoteRunResult,
+    ResultTable,
 )
 from dbt.logger import GLOBAL_LOGGER as logger
 from dbt.task.compile import CompileRunner
 from dbt.rpc.error import dbt_error, RPCException, server_error
 
 
-RPCSQLResult = TypeVar('RPCSQLResult', bound=RemoteCompileResult)
+RPCSQLResult = TypeVar('RPCSQLResult', bound=RemoteCompileResultMixin)
 
 
 class GenericRPCRunner(CompileRunner, Generic[RPCSQLResult]):