Skip to content

Commit

Permalink
Merge pull request #1652 from fishtown-analytics/feature/dbt-rpc-run
Browse files Browse the repository at this point in the history
dbt rpc tasks
  • Loading branch information
beckjake authored Aug 2, 2019
2 parents b0217ba + 115ef7d commit 07aedc0
Show file tree
Hide file tree
Showing 10 changed files with 309 additions and 11 deletions.
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,9 @@ coverage.xml
.hypothesis/
test.env

# Mypy
.mypy_cache/

# Translations
*.mo
*.pot
Expand Down
2 changes: 1 addition & 1 deletion core/dbt/contracts/project.py
Original file line number Diff line number Diff line change
Expand Up @@ -142,7 +142,7 @@ class ProfileConfig(HyphenatedJsonSchemaMixin, Replaceable):
config: UserConfig
threads: int
# TODO: make this a dynamic union of some kind?
credentials: Optional[Any]
credentials: Optional[Dict[str, Any]]


@dataclass
Expand Down
24 changes: 24 additions & 0 deletions core/dbt/task/compile.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import os
import signal
import threading
from typing import Union, List, Dict, Any

from dbt.adapters.factory import get_adapter
from dbt.clients.jinja import extract_toplevel_blocks
Expand Down Expand Up @@ -162,3 +163,26 @@ def handle_request(self, name, sql, macros=None):

self._raise_set_error()
return self.node_results[0].to_dict()


class RemoteCompileProjectTask(CompileTask, RemoteCallable):
METHOD_NAME = 'compile_project'

def __init__(self, args, config, manifest):
super().__init__(args, config)
self.manifest = manifest.deepcopy(config=config)

def load_manifest(self):
# we started out with a manifest!
pass

def handle_request(
self,
models: Union[None, str, List[str]] = None,
exclude: Union[None, str, List[str]] = None,
) -> Dict[str, List[Any]]:
self.args.models = self._listify(models)
self.args.exclude = self._listify(exclude)

results = self.run()
return {'results': [r.to_dict() for r in results]}
14 changes: 11 additions & 3 deletions core/dbt/task/rpc_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,8 +7,12 @@

from dbt.logger import RPC_LOGGER as logger
from dbt.task.base import ConfiguredTask
from dbt.task.compile import CompileTask, RemoteCompileTask
from dbt.task.run import RemoteRunTask
from dbt.task.compile import (
CompileTask, RemoteCompileTask, RemoteCompileProjectTask
)
from dbt.task.run import RemoteRunTask, RemoteRunProjectTask
from dbt.task.seed import RemoteSeedProjectTask
from dbt.task.test import RemoteTestProjectTask
from dbt.utils import JSONEncoder
from dbt import rpc

Expand All @@ -19,7 +23,11 @@ def __init__(self, args, config, tasks=None):
# compile locally
self.manifest = self._compile_manifest()
self.task_manager = rpc.TaskManager()
tasks = tasks or [RemoteCompileTask, RemoteRunTask]
tasks = tasks or [
RemoteCompileTask, RemoteCompileProjectTask,
RemoteRunTask, RemoteRunProjectTask,
RemoteSeedProjectTask, RemoteTestProjectTask
]
for cls in tasks:
task = cls(args, config, self.manifest)
self.task_manager.add_task_handler(task)
Expand Down
27 changes: 25 additions & 2 deletions core/dbt/task/run.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
from __future__ import print_function

import functools
import time
from typing import Union, List, Dict, Any

from dbt.logger import GLOBAL_LOGGER as logger
from dbt.node_types import NodeType, RunHookType
Expand All @@ -19,6 +18,7 @@

from dbt.compilation import compile_node
from dbt.task.compile import CompileTask, RemoteCompileTask
from dbt.task.runnable import RemoteCallable
from dbt.utils import get_nodes_by_tags


Expand Down Expand Up @@ -183,3 +183,26 @@ class RemoteRunTask(RemoteCompileTask, RunTask):

def get_runner_type(self):
return RPCExecuteRunner


class RemoteRunProjectTask(RunTask, RemoteCallable):
METHOD_NAME = 'run_project'

def __init__(self, args, config, manifest):
super().__init__(args, config)
self.manifest = manifest.deepcopy(config=config)

def load_manifest(self):
# we started out with a manifest!
pass

def handle_request(
self,
models: Union[None, str, List[str]] = None,
exclude: Union[None, str, List[str]] = None,
) -> Dict[str, List[Any]]:
self.args.models = self._listify(models)
self.args.exclude = self._listify(exclude)

results = self.run()
return {'results': [r.to_dict() for r in results]}
20 changes: 19 additions & 1 deletion core/dbt/task/runnable.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from abc import abstractmethod
from datetime import datetime
from multiprocessing.dummy import Pool as ThreadPool
from typing import Union, List, Optional

from dbt import rpc
from dbt.task.base import ConfiguredTask
Expand Down Expand Up @@ -40,10 +41,16 @@ def __init__(self, args, config):
self.manifest = None
self.linker = None

def _runtime_initialize(self):
def load_manifest(self):
self.manifest = load_manifest(self.config)

def compile_manifest(self):
self.linker = compile_manifest(self.config, self.manifest)

def _runtime_initialize(self):
self.load_manifest()
self.compile_manifest()


class GraphRunnableTask(ManifestTask):
def __init__(self, args, config):
Expand Down Expand Up @@ -344,6 +351,17 @@ def handle_request(self, **kwargs):
'from_kwargs not implemented'
)

@staticmethod
def _listify(
value: Optional[Union[str, List[str]]]
) -> Optional[List[str]]:
if value is None:
return None
elif isinstance(value, str):
return [value]
else:
return value

def decode_sql(self, sql):
"""Base64 decode a string. This should only be used for sql in calls.
Expand Down
21 changes: 21 additions & 0 deletions core/dbt/task/seed.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,11 @@
import random
from typing import List, Dict, Any

from dbt.logger import GLOBAL_LOGGER as logger
from dbt.node_runners import SeedRunner
from dbt.node_types import NodeType
from dbt.task.run import RunTask
from dbt.task.runnable import RemoteCallable
import dbt.ui.printer


Expand Down Expand Up @@ -44,3 +47,21 @@ def show_tables(self, results):
for result in results:
if result.error is None:
self.show_table(result)


class RemoteSeedProjectTask(SeedTask, RemoteCallable):
METHOD_NAME = 'seed_project'

def __init__(self, args, config, manifest):
super().__init__(args, config)
self.manifest = manifest.deepcopy(config=config)

def load_manifest(self):
# we started out with a manifest!
pass

def handle_request(self, show: bool = False) -> Dict[str, List[Any]]:
self.args.show = show

results = self.run()
return {'results': [r.to_dict() for r in results]}
35 changes: 31 additions & 4 deletions core/dbt/task/test.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,9 @@
from typing import Union, List, Dict, Any

from dbt.node_runners import TestRunner
from dbt.node_types import NodeType
from dbt.task.run import RunTask
from dbt.task.runnable import RemoteCallable


class TestTask(RunTask):
Expand All @@ -12,11 +15,8 @@ class TestTask(RunTask):
def raise_on_first_error(self):
return False

def before_run(self, adapter, selected_uids):
def safe_run_hooks(self, adapter, hook_type, extra_context):
# Don't execute on-run-* hooks for tests
self.populate_adapter_cache(adapter)

def after_run(self, adapter, results):
pass

def build_query(self):
Expand All @@ -41,3 +41,30 @@ def build_query(self):

def get_runner_type(self):
return TestRunner


class RemoteTestProjectTask(TestTask, RemoteCallable):
METHOD_NAME = 'test_project'

def __init__(self, args, config, manifest):
super().__init__(args, config)
self.manifest = manifest.deepcopy(config=config)

def load_manifest(self):
# we started out with a manifest!
pass

def handle_request(
self,
models: Union[None, str, List[str]] = None,
exclude: Union[None, str, List[str]] = None,
data: bool = False,
schema: bool = False,
) -> Dict[str, List[Any]]:
self.args.models = self._listify(models)
self.args.exclude = self._listify(exclude)
self.args.data = data
self.args.schema = schema

results = self.run()
return {'results': [r.to_dict() for r in results]}
113 changes: 113 additions & 0 deletions test/integration/042_sources_test/seed.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,113 @@
create table {schema}.seed_expected (
favorite_color TEXT,
id INTEGER,
first_name TEXT,
email TEXT,
ip_address TEXT,
updated_at TIMESTAMP WITHOUT TIME ZONE
);


INSERT INTO {schema}.seed_expected
("favorite_color","id","first_name","email","ip_address","updated_at")
VALUES
('blue',1,'Larry','[email protected]','''69.135.206.194''','2008-09-12 19:08:31'),
('blue',2,'Larry','[email protected]','''64.210.133.162''','1978-05-09 04:15:14'),
('blue',3,'Anna','[email protected]','''168.104.64.114''','2011-10-16 04:07:57'),
('blue',4,'Sandra','[email protected]','''229.235.252.98''','1973-07-19 10:52:43'),
('blue',5,'Fred','[email protected]','''78.229.170.124''','2012-09-30 16:38:29'),
('blue',6,'Stephen','[email protected]','''182.227.157.105''','1995-11-07 21:40:50'),
('blue',7,'William','[email protected]','''135.139.249.50''','1982-09-05 03:11:59'),
('blue',8,'Jessica','[email protected]','''203.62.178.210''','1991-10-16 11:03:15'),
('blue',9,'Douglas','[email protected]','''178.187.247.1''','1979-10-01 09:49:48'),
('blue',10,'Lisa','[email protected]','''168.234.128.249''','2011-05-26 07:45:49'),
('blue',11,'Ralph','[email protected]','''55.152.163.149''','1972-11-18 19:06:11'),
('blue',12,'Louise','[email protected]','''141.116.153.154''','2014-11-25 20:56:14'),
('blue',13,'Clarence','[email protected]','''81.171.31.133''','2011-11-17 07:02:36'),
('blue',14,'Daniel','[email protected]','''8.204.211.37''','1980-09-13 00:09:04'),
('blue',15,'Katherine','[email protected]','''176.96.134.59''','1997-08-22 19:36:56'),
('blue',16,'Billy','[email protected]','''214.108.78.85''','2003-10-19 02:14:47'),
('blue',17,'Annie','[email protected]','''190.108.42.70''','1988-10-28 15:12:35'),
('blue',18,'Shirley','[email protected]','''109.251.164.84''','1988-08-24 10:50:57'),
('blue',19,'Roger','[email protected]','''38.145.218.108''','1985-12-31 15:17:15'),
('blue',20,'Lillian','[email protected]','''47.57.236.17''','1970-06-08 02:09:05'),
('blue',21,'Aaron','[email protected]','''205.245.118.221''','1985-10-11 23:07:49'),
('blue',22,'Patrick','[email protected]','''19.8.100.182''','2006-03-29 12:53:56'),
('blue',23,'Phillip','[email protected]','''41.38.254.103''','2011-11-07 15:35:43'),
('blue',24,'Henry','[email protected]','''1.191.216.252''','2008-08-28 08:30:44'),
('blue',25,'Irene','[email protected]','''50.17.60.190''','1994-04-01 07:15:02'),
('blue',26,'Andrew','[email protected]','''123.52.253.176''','2000-11-01 06:03:25'),
('blue',27,'David','[email protected]','''238.23.203.42''','1988-01-25 07:29:18'),
('blue',28,'Henry','[email protected]','''248.102.2.185''','1983-01-01 13:36:37'),
('blue',29,'Evelyn','[email protected]','''32.80.46.119''','1979-07-16 17:24:12'),
('blue',30,'Tammy','[email protected]','''249.246.167.88''','2001-04-03 10:00:23'),
('blue',31,'Jacqueline','[email protected]','''127.181.97.47''','1986-02-11 21:35:50'),
('blue',32,'Earl','[email protected]','''166.47.248.240''','1996-07-06 08:16:27'),
('blue',33,'Juan','[email protected]','''71.77.2.200''','1987-01-31 03:46:44'),
('blue',34,'Diane','[email protected]','''140.94.133.12''','1994-06-11 02:30:05'),
('blue',35,'Randy','[email protected]','''73.255.34.196''','2005-05-26 20:28:39'),
('blue',36,'Janice','[email protected]','''22.214.227.32''','1990-02-09 04:16:52'),
('blue',37,'Laura','[email protected]','''159.148.145.73''','2015-03-17 05:59:25'),
('blue',38,'Gary','[email protected]','''40.193.124.56''','1970-01-27 10:04:51'),
('blue',39,'Jesse','[email protected]','''31.7.86.103''','2009-03-14 08:14:29'),
('blue',40,'Sandra','[email protected]','''223.80.168.239''','1993-05-21 14:08:54'),
('blue',41,'Scott','[email protected]','''38.238.46.83''','1980-08-30 11:16:56'),
('blue',42,'Phillip','[email protected]','''158.234.59.34''','2011-12-15 23:26:31'),
('blue',43,'Steven','[email protected]','''90.247.57.68''','2011-10-29 19:03:28'),
('blue',44,'Deborah','[email protected]','''179.125.143.240''','1995-04-10 14:36:07'),
('blue',45,'Lori','[email protected]','''64.80.162.180''','1980-12-27 16:49:15'),
('blue',46,'Sean','[email protected]','''240.116.183.69''','1988-06-12 21:24:45'),
('blue',47,'Terry','[email protected]','''118.38.213.137''','1997-09-22 16:43:19'),
('blue',48,'Dorothy','[email protected]','''116.81.76.49''','2005-02-28 13:33:24'),
('blue',49,'Samuel','[email protected]','''38.191.253.40''','1989-01-19 21:15:48'),
('blue',50,'Ralph','[email protected]','''104.84.60.174''','2007-08-11 10:21:49'),
('green',51,'Wayne','[email protected]','''90.61.24.102''','1983-07-03 16:58:12'),
('green',52,'Rose','[email protected]','''240.83.81.10''','1995-06-08 11:46:23'),
('green',53,'Louise','[email protected]','''105.11.82.145''','2016-09-19 14:45:51'),
('green',54,'Kenneth','[email protected]','''139.5.45.94''','1976-08-17 11:26:19'),
('green',55,'Donna','[email protected]','''19.45.169.45''','2006-05-27 16:51:40'),
('green',56,'Johnny','[email protected]','''118.202.238.23''','1975-11-17 08:42:32'),
('green',57,'Patrick','[email protected]','''231.25.153.198''','1997-08-06 11:51:09'),
('green',58,'Helen','[email protected]','''8.40.21.39''','1993-08-04 19:53:40'),
('green',59,'Patricia','[email protected]','''212.198.40.15''','1977-08-03 16:37:27'),
('green',60,'Joseph','[email protected]','''13.15.63.238''','2005-07-23 20:22:06'),
('green',61,'Phillip','[email protected]','''177.98.201.190''','1976-05-19 21:47:44'),
('green',62,'Joan','[email protected]','''105.229.170.71''','1972-09-07 17:53:47'),
('green',63,'Phyllis','[email protected]','''35.145.8.244''','2000-01-01 22:33:37'),
('green',64,'Katherine','[email protected]','''248.168.205.32''','1991-01-09 06:40:24'),
('green',65,'Laura','[email protected]','''128.129.115.152''','1997-10-23 12:04:56'),
('green',66,'Juan','[email protected]','''44.228.124.51''','2004-11-10 05:07:35'),
('green',67,'Judith','[email protected]','''40.227.179.115''','1977-08-02 17:01:45'),
('green',68,'Beverly','[email protected]','''208.34.84.59''','2016-03-06 20:07:23'),
('green',69,'Lawrence','[email protected]','''59.158.212.223''','1988-09-13 06:07:21'),
('green',70,'Gloria','[email protected]','''245.231.88.33''','1995-03-18 22:32:46'),
('green',71,'Steven','[email protected]','''104.50.58.255''','2001-08-05 21:26:20'),
('green',72,'Betty','[email protected]','''103.177.214.220''','1981-12-14 21:26:54'),
('green',73,'Mildred','[email protected]','''151.158.8.130''','2000-04-19 10:13:55'),
('green',74,'Donald','[email protected]','''9.178.102.255''','1972-12-03 00:58:24'),
('green',75,'Eric','[email protected]','''85.2.241.227''','1992-11-01 05:59:30'),
('green',76,'Joyce','[email protected]','''169.224.20.36''','1985-10-24 06:50:01'),
('green',77,'Maria','[email protected]','''143.189.167.135''','2005-10-05 05:17:42'),
('green',78,'Harry','[email protected]','''156.47.176.237''','1978-03-26 05:53:33'),
('green',79,'Kevin','[email protected]','''79.136.183.83''','1994-10-12 04:38:52'),
('green',80,'David','[email protected]','''133.149.172.153''','1976-12-15 16:24:24'),
('green',81,'Kathy','[email protected]','''229.242.72.228''','1979-03-04 02:58:56'),
('green',82,'Adam','[email protected]','''13.145.21.10''','1982-11-07 11:46:59'),
('green',83,'Brandon','[email protected]','''73.249.128.212''','2013-10-30 05:30:36'),
('green',84,'Henry','[email protected]','''211.36.214.242''','1985-01-09 06:37:27'),
('green',85,'Eric','[email protected]','''191.166.188.251''','2004-05-01 23:21:42'),
('green',86,'Jason','[email protected]','''193.92.16.182''','1973-01-08 09:05:39'),
('green',87,'Diana','[email protected]','''19.130.175.245''','1994-10-05 22:50:49'),
('green',88,'Andrea','[email protected]','''94.155.233.96''','2002-04-26 08:41:44'),
('green',89,'Louis','[email protected]','''26.217.34.111''','2003-08-25 07:56:39'),
('green',90,'Jane','[email protected]','''43.4.220.135''','1987-03-20 20:39:04'),
('green',91,'Larry','[email protected]','''97.126.79.34''','2000-09-07 20:26:19'),
('green',92,'Louis','[email protected]','''37.148.40.127''','2011-09-16 20:12:14'),
('green',93,'Jennifer','[email protected]','''38.106.254.142''','1988-07-15 05:06:49'),
('green',94,'Wayne','[email protected]','''223.28.26.187''','2009-12-15 06:16:54'),
('green',95,'Lori','[email protected]','''181.250.181.58''','1984-10-28 03:29:19'),
('green',96,'Judy','[email protected]','''180.121.239.219''','1986-02-07 15:18:10'),
('green',97,'Phillip','[email protected]','''255.247.0.175''','2002-12-26 08:44:45'),
('green',98,'Gloria','[email protected]','''156.140.7.128''','1997-10-04 07:58:58'),
('green',99,'Paul','[email protected]','''183.59.198.197''','1991-11-14 12:33:55'),
('green',100,'Frank','[email protected]','''150.143.68.121''','2010-06-12 23:55:39');
Loading

0 comments on commit 07aedc0

Please sign in to comment.