Merge pull request #1509 from heisencoder/feature/upgrade-to-networkx-2

Feature/upgrade to networkx 2

core/dbt/linker.py

@@ -136,7 +136,7 @@ def _find_new_additions(self):
 
         Callers must hold the lock.
         """
-        for node, in_degree in self.graph.in_degree_iter():
+        for node, in_degree in dict(self.graph.in_degree()).items():
             if not self._already_known(node) and in_degree == 0:
                 self.inner.put((self._scores[node], node))
                 self.queued.add(node)
@@ -268,5 +268,5 @@ def _updated_graph(graph, manifest):
         for key in GRAPH_SERIALIZE_BLACKLIST:
             if key in data:
                 del data[key]
-        graph.add_node(node_id, data)
+        graph.add_node(node_id, **data)
     return graph

core/setup.py

@@ -46,7 +46,7 @@ def read(fname):
         'Jinja2>=2.10',
         'PyYAML>=3.11',
         'sqlparse==0.2.3',
-        'networkx==1.11',
+        'networkx>=1.11,<3',
         'minimal-snowplow-tracker==0.0.2',
         'requests>=2.18.0,<3',
         'colorama==0.3.9',

test/unit/test_graph.py

@@ -14,7 +14,7 @@
 
 try:
     from queue import Empty
-except KeyError:
+except ImportError:
     from Queue import Empty
 
 
@@ -128,12 +128,12 @@ def test__single_model(self):
         compiler = self.get_compiler(config)
         linker = compiler.compile(manifest)
 
-        self.assertEquals(
-            linker.nodes(),
+        self.assertEqual(
+            list(linker.nodes()),
             ['model.test_models_compile.model_one'])
 
-        self.assertEquals(
-            linker.edges(),
+        self.assertEqual(
+            list(linker.edges()),
             [])
 
     def test__two_models_simple_ref(self):
@@ -196,7 +196,7 @@ def test__model_materializations(self):
             key = 'model.test_models_compile.{}'.format(model)
             actual = manifest.nodes[key].get('config', {}) \
                                              .get('materialized')
-            self.assertEquals(actual, expected)
+            self.assertEqual(actual, expected)
 
     def test__model_incremental(self):
         self.use_models({
@@ -221,8 +221,8 @@ def test__model_incremental(self):
 
         node = 'model.test_models_compile.model_one'
 
-        self.assertEqual(linker.nodes(), [node])
-        self.assertEqual(linker.edges(), [])
+        self.assertEqual(list(linker.nodes()), [node])
+        self.assertEqual(list(linker.edges()), [])
 
         self.assertEqual(
             manifest.nodes[node].get('config', {}).get('materialized'),

test/unit/test_linker.py

@@ -1,12 +1,14 @@
 import mock
+import os
+import tempfile
 import unittest
 
 import dbt.utils
 
 from dbt import linker
 try:
     from queue import Empty
-except KeyError:
+except ImportError:
     from Queue import Empty
 
 
@@ -37,6 +39,26 @@ def test_linker_add_node(self):
 
         self.assertEqual(len(actual_nodes), len(expected_nodes))
 
+    def test_linker_write_and_read_graph(self):
+        expected_nodes = ['A', 'B', 'C']
+        for node in expected_nodes:
+            self.linker.add_node(node)
+
+        manifest = _mock_manifest('ABC')
+        (fd, fname) = tempfile.mkstemp()
+        os.close(fd)
+        try:
+            self.linker.write_graph(fname, manifest)
+            new_linker = linker.from_file(fname)
+        finally:
+            os.unlink(fname)
+
+        actual_nodes = new_linker.nodes()
+        for node in expected_nodes:
+            self.assertIn(node, actual_nodes)
+
+        self.assertEqual(len(actual_nodes), len(expected_nodes))
+
     def assert_would_join(self, queue):
         """test join() without timeout risk"""
         self.assertEqual(queue.inner.unfinished_tasks, 0)