Repository: incubator-airflow Updated Branches: refs/heads/master a52123dca -> ffe6162f2
[AIRFLOW-1242] Allowing project_id to have a colon in it. Closes #2335 from zoyahav/master Project: http://git-wip-us.apache.org/repos/asf/incubator-airflow/repo Commit: http://git-wip-us.apache.org/repos/asf/incubator-airflow/commit/ffe6162f Tree: http://git-wip-us.apache.org/repos/asf/incubator-airflow/tree/ffe6162f Diff: http://git-wip-us.apache.org/repos/asf/incubator-airflow/diff/ffe6162f Branch: refs/heads/master Commit: ffe6162f2b9e61f686b1fe993eb95e43e60c7dd5 Parents: a52123dc Author: Zohar Yahav <[email protected]> Authored: Wed Jun 7 11:30:40 2017 -0700 Committer: Chris Riccomini <[email protected]> Committed: Wed Jun 7 11:30:40 2017 -0700 ---------------------------------------------------------------------- airflow/contrib/hooks/bigquery_hook.py | 17 ++++++++--- tests/contrib/hooks/test_bigquery_hook.py | 41 ++++++++++++++++---------- 2 files changed, 39 insertions(+), 19 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/ffe6162f/airflow/contrib/hooks/bigquery_hook.py ---------------------------------------------------------------------- diff --git a/airflow/contrib/hooks/bigquery_hook.py b/airflow/contrib/hooks/bigquery_hook.py index bf18209..cc76953 100644 --- a/airflow/contrib/hooks/bigquery_hook.py +++ b/airflow/contrib/hooks/bigquery_hook.py @@ -934,13 +934,22 @@ def _split_tablename(table_input, default_project_id, var_name=None): else: return "Format exception for {var}: ".format(var=var_name) - cmpt = table_input.split(':') + if table_input.count('.') + table_input.count(':') > 3: + raise Exception(( + '{var}Use either : or . to specify project ' + 'got {input}' + ).format(var=var_print(var_name), input=table_input)) + + cmpt = table_input.rsplit(':', 1) + project_id = None + rest = table_input if len(cmpt) == 1: project_id = None rest = cmpt[0] - elif len(cmpt) == 2: - project_id = cmpt[0] - rest = cmpt[1] + elif len(cmpt) == 2 and cmpt[0].count(':') <= 1: + if cmpt[-1].count('.') != 2: + project_id = cmpt[0] + rest = cmpt[1] else: raise Exception(( '{var}Expect format of (<project:)<dataset>.<table>, ' http://git-wip-us.apache.org/repos/asf/incubator-airflow/blob/ffe6162f/tests/contrib/hooks/test_bigquery_hook.py ---------------------------------------------------------------------- diff --git a/tests/contrib/hooks/test_bigquery_hook.py b/tests/contrib/hooks/test_bigquery_hook.py index 9970d8d..2fe9ddb 100644 --- a/tests/contrib/hooks/test_bigquery_hook.py +++ b/tests/contrib/hooks/test_bigquery_hook.py @@ -88,25 +88,35 @@ class TestBigQueryTableSplitter(unittest.TestCase): self.assertEqual("dataset", dataset) self.assertEqual("table", table) - def test_invalid_syntax_column_double_project(self): - with self.assertRaises(Exception) as context: - hook._split_tablename('alt1:alt.dataset.table', + def test_colon_in_project(self): + project, dataset, table = hook._split_tablename('alt1:alt.dataset.table', + 'project') + + self.assertEqual('alt1:alt', project) + self.assertEqual("dataset", dataset) + self.assertEqual("table", table) + + + def test_valid_double_column(self): + project, dataset, table = hook._split_tablename('alt1:alt:dataset.table', 'project') - self.assertIn('Use either : or . to specify project', - str(context.exception), "") - self.assertFalse('Format exception for' in str(context.exception)) + self.assertEqual('alt1:alt', project) + self.assertEqual("dataset", dataset) + self.assertEqual("table", table) + - def test_invalid_syntax_double_column(self): + def test_invalid_syntax_triple_colon(self): with self.assertRaises(Exception) as context: - hook._split_tablename('alt1:alt:dataset.table', + hook._split_tablename('alt1:alt2:alt3:dataset.table', 'project') - self.assertIn('Expect format of (<project:)<dataset>.<table>', + self.assertIn('Use either : or . to specify project', str(context.exception), "") self.assertFalse('Format exception for' in str(context.exception)) - def test_invalid_syntax_tiple_dot(self): + + def test_invalid_syntax_triple_dot(self): with self.assertRaises(Exception) as context: hook._split_tablename('alt1.alt.dataset.table', 'project') @@ -117,7 +127,7 @@ class TestBigQueryTableSplitter(unittest.TestCase): def test_invalid_syntax_column_double_project_var(self): with self.assertRaises(Exception) as context: - hook._split_tablename('alt1:alt.dataset.table', + hook._split_tablename('alt1:alt2:alt.dataset.table', 'project', 'var_x') self.assertIn('Use either : or . to specify project', @@ -125,17 +135,17 @@ class TestBigQueryTableSplitter(unittest.TestCase): self.assertIn('Format exception for var_x:', str(context.exception), "") - def test_invalid_syntax_double_column_var(self): + def test_invalid_syntax_triple_colon_project_var(self): with self.assertRaises(Exception) as context: - hook._split_tablename('alt1:alt:dataset.table', + hook._split_tablename('alt1:alt2:alt:dataset.table', 'project', 'var_x') - self.assertIn('Expect format of (<project:)<dataset>.<table>', + self.assertIn('Use either : or . to specify project', str(context.exception), "") self.assertIn('Format exception for var_x:', str(context.exception), "") - def test_invalid_syntax_tiple_dot_var(self): + def test_invalid_syntax_triple_dot_var(self): with self.assertRaises(Exception) as context: hook._split_tablename('alt1.alt.dataset.table', 'project', 'var_x') @@ -178,3 +188,4 @@ class TestBigQueryBaseCursor(unittest.TestCase): if __name__ == '__main__': unittest.main() +
