def test_append_non_root_node(self):
chain = SqlChain().append(
SqlNode(output_name='root', source='root', query='q1'))
self.assertIsNone(chain.root.next)
node = SqlNode(output_name='next_node', source='root', query='q2')
chain.append(node)
self.assertIs(node, chain.root.next)
self.assertIs(node, chain.get(node.output_name))
def test_to_pipeline_only_evaluate_once_per_pipeline_and_node(
self, mocked_sql_transform):
p = beam.Pipeline()
ie.current_env().watch({'p': p})
pcoll_1 = p | 'create pcoll_1' >> beam.Create([1, 2, 3])
pcoll_2 = p | 'create pcoll_2' >> beam.Create([4, 5, 6])
ie.current_env().watch({'pcoll_1': pcoll_1, 'pcoll_2': pcoll_2})
node = SqlNode(output_name='root',
source={'pcoll_1', 'pcoll_2'},
query='q1')
chain = SqlChain(user_pipeline=p).append(node)
_ = chain.to_pipeline()
mocked_sql_transform.assert_called_once()
_ = chain.to_pipeline()
mocked_sql_transform.assert_called_once()
def get_sql_chain(self, pipeline, set_user_pipeline=False):
if pipeline not in self.sql_chain:
self.sql_chain[pipeline] = SqlChain()
chain = self.sql_chain[pipeline]
if set_user_pipeline:
if chain.user_pipeline and chain.user_pipeline is not pipeline:
raise ValueError(
'The beam_sql magic tries to query PCollections from multiple '
'pipelines: %s and %s', chain.user_pipeline, pipeline)
chain.user_pipeline = pipeline
return chain
def test_nodes_with_same_outputs(self, mocked_sql_transform):
p = beam.Pipeline()
ie.current_env().watch({'p_nodes_with_same_output': p})
pcoll = p | 'create pcoll' >> beam.Create([1, 2, 3])
ie.current_env().watch({'pcoll': pcoll})
chain = SqlChain(user_pipeline=p)
output_name = 'output'
with patch('IPython.get_ipython',
new_callable=mock_get_ipython) as cell:
with cell:
node_cell_1 = SqlNode(output_name, source='pcoll', query='q1')
chain.append(node_cell_1)
_ = chain.to_pipeline()
mocked_sql_transform.assert_called_with(
'schema_loaded_beam_sql_output_1')
with cell:
node_cell_2 = SqlNode(output_name, source='pcoll', query='q2')
chain.append(node_cell_2)
_ = chain.to_pipeline()
mocked_sql_transform.assert_called_with(
'schema_loaded_beam_sql_output_2')
def test_append_first_node(self):
node = SqlNode(output_name='first', source='a', query='q1')
chain = SqlChain().append(node)
self.assertIs(node, chain.get(node.output_name))
self.assertIs(node, chain.root)
self.assertIs(node, chain.current)
def test_init(self):
chain = SqlChain()
self.assertEqual({}, chain.nodes)
self.assertIsNone(chain.root)
self.assertIsNone(chain.current)
self.assertIsNone(chain.user_pipeline)