diff --git a/UPDATING.md b/UPDATING.md index 6e25cc7af1219..50fd686bd5f0e 100644 --- a/UPDATING.md +++ b/UPDATING.md @@ -322,6 +322,14 @@ We should default it as `true` to avoid confusion In order to migrate the database, you should use the command `airflow db upgrade`, but in some cases manual steps are required. +#### Unique conn_id in connection table + +Previously, Airflow allowed users to add more than one connection with the same `conn_id` and on access it would choose one connection randomly. This acted as a basic load balancing and fault tolerance technique, when used in conjunction with retries. + +This behavior caused some confusion for users, and there was no clear evidence if it actually worked well or not. + +Now the `conn_id` will be unique. If you already have duplicates in your metadata database, you will have to manage those duplicate connections before upgrading the database. + #### Not-nullable conn_type column in connection table The `conn_type` column in the `connection` table must contain content. Previously, this rule was enforced diff --git a/airflow/exceptions.py b/airflow/exceptions.py index 9c2bde9599005..3038bb6ec4b4a 100644 --- a/airflow/exceptions.py +++ b/airflow/exceptions.py @@ -182,3 +182,7 @@ def __str__(self): result += "\n" + prepare_code_snippet(self.file_path, parse_error.line_no) + "\n" return result + + +class ConnectionNotUnique(AirflowException): + """Raise when multiple values are found for the same conn_id""" diff --git a/airflow/migrations/versions/8d48763f6d53_add_unique_constraint_to_conn_id.py b/airflow/migrations/versions/8d48763f6d53_add_unique_constraint_to_conn_id.py new file mode 100644 index 0000000000000..34ce9e5aa5679 --- /dev/null +++ b/airflow/migrations/versions/8d48763f6d53_add_unique_constraint_to_conn_id.py @@ -0,0 +1,67 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +"""add unique constraint to conn_id + +Revision ID: 8d48763f6d53 +Revises: 8f966b9c467a +Create Date: 2020-05-03 16:55:01.834231 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = '8d48763f6d53' +down_revision = '8f966b9c467a' +branch_labels = None +depends_on = None + + +def upgrade(): + """Apply add unique constraint to conn_id and set it as non-nullable""" + try: + with op.batch_alter_table('connection') as batch_op: + batch_op.create_unique_constraint( + constraint_name="unique_conn_id", + columns=["conn_id"] + ) + + batch_op.alter_column( + "conn_id", + nullable=False, + existing_type=sa.String(250) + ) + except sa.exc.IntegrityError: + raise Exception("Make sure there are no duplicate connections with the same conn_id or null values") + + +def downgrade(): + """Unapply add unique constraint to conn_id and set it as non-nullable""" + with op.batch_alter_table('connection') as batch_op: + batch_op.drop_constraint( + constraint_name="unique_conn_id", + type_="unique" + ) + + batch_op.alter_column( + "conn_id", + nullable=True, + existing_type=sa.String(250) + ) diff --git a/airflow/models/connection.py b/airflow/models/connection.py index 75174c41784cc..37c3d1b1b779e 100644 --- a/airflow/models/connection.py +++ b/airflow/models/connection.py @@ -146,7 +146,7 @@ class Connection(Base, LoggingMixin): __tablename__ = "connection" id = Column(Integer(), primary_key=True) - conn_id = Column(String(ID_LEN)) + conn_id = Column(String(ID_LEN), unique=True, nullable=False) conn_type = Column(String(500), nullable=False) host = Column(String(500)) schema = Column(String(500)) diff --git a/airflow/secrets/local_filesystem.py b/airflow/secrets/local_filesystem.py index 0c246afc4054e..29754d52dc2fa 100644 --- a/airflow/secrets/local_filesystem.py +++ b/airflow/secrets/local_filesystem.py @@ -28,7 +28,9 @@ import yaml -from airflow.exceptions import AirflowException, AirflowFileParseException, FileSyntaxError +from airflow.exceptions import ( + AirflowException, AirflowFileParseException, ConnectionNotUnique, FileSyntaxError, +) from airflow.secrets.base_secrets import BaseSecretsBackend from airflow.utils.file import COMMENT_PATTERN from airflow.utils.log.logging_mixin import LoggingMixin @@ -252,6 +254,10 @@ def load_connections(file_path: str): connections_by_conn_id[key].append(_create_connection(key, secret_value)) else: connections_by_conn_id[key].append(_create_connection(key, secret_values)) + + if len(connections_by_conn_id[key]) > 1: + raise ConnectionNotUnique(f"Found multiple values for {key} in {file_path}") + num_conn = sum(map(len, connections_by_conn_id.values())) log.debug("Loaded %d connections", num_conn) diff --git a/docs/concepts.rst b/docs/concepts.rst index 3d2bc891b2b33..b75cb4105a6c4 100644 --- a/docs/concepts.rst +++ b/docs/concepts.rst @@ -675,11 +675,6 @@ managed in the UI (``Menu -> Admin -> Connections``). A ``conn_id`` is defined password / schema information attached to it. Airflow pipelines retrieve centrally-managed connections information by specifying the relevant ``conn_id``. -You may add more than one connection with the same ``conn_id``. When there is more than one connection -with the same ``conn_id``, the :py:meth:`~airflow.hooks.base_hook.BaseHook.get_connection` method on -:py:class:`~airflow.hooks.base_hook.BaseHook` will choose one connection randomly. This can be be used to -provide basic load balancing and fault tolerance, when used in conjunction with retries. - Airflow also provides a mechanism to store connections outside the database, e.g. in :ref:`environment variables `. Additional sources may be enabled, e.g. :ref:`AWS SSM Parameter Store `, or you may :ref:`roll your own secrets backend `. diff --git a/docs/howto/use-alternative-secrets-backend.rst b/docs/howto/use-alternative-secrets-backend.rst index 6b4c0c95198a4..edc313c609623 100644 --- a/docs/howto/use-alternative-secrets-backend.rst +++ b/docs/howto/use-alternative-secrets-backend.rst @@ -101,18 +101,16 @@ The key ``extra_dejson`` can be used to provide parameters as JSON object where The keys ``extra`` and ``extra_dejson`` are mutually exclusive. The JSON file must contain an object where the key contains the connection ID and the value contains -the definitions of one or more connections. In this format, the connection can be defined as a URI (string) or JSON object. +the definition of one connection. The connection can be defined as a URI (string) or JSON object. +For a guide about defining a connection as a URI, see:: :ref:`generating_connection_uri`. +For a description of the connection object parameters see :class:`~airflow.models.connection.Connection`. The following is a sample JSON file. .. code-block:: json { "CONN_A": "mysq://host_a", - "CONN_B": [ - "mysq://host_a", - "mysq://host_a" - ], - "CONN_C": { + "CONN_B": { "conn_type": "scheme", "host": "host", "schema": "lschema", @@ -146,8 +144,8 @@ In this format, the connection can be defined as a URI (string) or JSON object. x: y You can also define connections using a ``.env`` file. Then the key is the connection ID, and -the value should describe the connection using the URI. If the connection ID is repeated, all values will -be returned. The following is a sample file. +the value should describe the connection using the URI. Connection ID should not be repeated, it will +raise an exception. The following is a sample file. .. code-block:: text diff --git a/tests/providers/pagerduty/hooks/test_pagerduty.py b/tests/providers/pagerduty/hooks/test_pagerduty.py index 254d5c0a31cda..f614e33f73a4c 100644 --- a/tests/providers/pagerduty/hooks/test_pagerduty.py +++ b/tests/providers/pagerduty/hooks/test_pagerduty.py @@ -28,8 +28,9 @@ class TestPagerdutyHook(unittest.TestCase): + @classmethod @provide_session - def setUp(self, session=None): + def setUpClass(cls, session=None): session.add(Connection( conn_id=DEFAULT_CONN_ID, conn_type='http', diff --git a/tests/secrets/test_local_filesystem.py b/tests/secrets/test_local_filesystem.py index 6f58850f9e089..97f6d423b9a33 100644 --- a/tests/secrets/test_local_filesystem.py +++ b/tests/secrets/test_local_filesystem.py @@ -24,7 +24,7 @@ from parameterized import parameterized -from airflow.exceptions import AirflowException, AirflowFileParseException +from airflow.exceptions import AirflowException, AirflowFileParseException, ConnectionNotUnique from airflow.secrets import local_filesystem from airflow.secrets.local_filesystem import LocalFilesystemBackend @@ -124,16 +124,16 @@ class TestLoadConnection(unittest.TestCase): ( ("CONN_ID=mysql://host_1/", {"CONN_ID": ["mysql://host_1"]}), ( - "CONN_ID=mysql://host_1/\nCONN_ID=mysql://host_2/", - {"CONN_ID": ["mysql://host_1", "mysql://host_2"]}, + "CONN_ID1=mysql://host_1/\nCONN_ID2=mysql://host_2/", + {"CONN_ID1": ["mysql://host_1"], "CONN_ID2": ["mysql://host_2"]}, ), ( - "CONN_ID=mysql://host_1/\n # AAAA\nCONN_ID=mysql://host_2/", - {"CONN_ID": ["mysql://host_1", "mysql://host_2"]}, + "CONN_ID1=mysql://host_1/\n # AAAA\nCONN_ID2=mysql://host_2/", + {"CONN_ID1": ["mysql://host_1"], "CONN_ID2": ["mysql://host_2"]}, ), ( - "\n\n\n\nCONN_ID=mysql://host_1/\n\n\n\n\nCONN_ID=mysql://host_2/\n\n\n", - {"CONN_ID": ["mysql://host_1", "mysql://host_2"]}, + "\n\n\n\nCONN_ID1=mysql://host_1/\n\n\n\n\nCONN_ID2=mysql://host_2/\n\n\n", + {"CONN_ID1": ["mysql://host_1"], "CONN_ID2": ["mysql://host_2"]}, ), ) ) @@ -162,16 +162,8 @@ def test_env_file_invalid_format(self, content, expected_message): ( ({"CONN_ID": "mysql://host_1"}, {"CONN_ID": ["mysql://host_1"]}), ({"CONN_ID": ["mysql://host_1"]}, {"CONN_ID": ["mysql://host_1"]}), - ( - {"CONN_ID": ["mysql://host_1", "mysql://host_2"]}, - {"CONN_ID": ["mysql://host_1", "mysql://host_2"]}, - ), ({"CONN_ID": {"uri": "mysql://host_1"}}, {"CONN_ID": ["mysql://host_1"]}), ({"CONN_ID": [{"uri": "mysql://host_1"}]}, {"CONN_ID": ["mysql://host_1"]}), - ( - {"CONN_ID": [{"uri": "mysql://host_1"}, {"uri": "mysql://host_2"}]}, - {"CONN_ID": ["mysql://host_1", "mysql://host_2"]}, - ), ) ) def test_json_file_should_load_connection(self, file_content, expected_connection_uris): @@ -211,16 +203,8 @@ def test_missing_file(self, mock_exists): ( ("""CONN_A: 'mysql://host_a'""", {"CONN_A": ["mysql://host_a"]}), (""" - CONN_B: - - 'mysql://host_a' - - 'mysql://host_b' - """, {"CONN_B": ["mysql://host_a", "mysql://host_b"]}), - (""" conn_a: mysql://hosta conn_b: - - mysql://hostb - - mysql://hostc - conn_c: conn_type: scheme host: host schema: lschema @@ -231,8 +215,8 @@ def test_missing_file(self, mock_exists): extra__google_cloud_platform__keyfile_dict: a: b extra__google_cloud_platform__keyfile_path: asaa""", - {"conn_a": ["mysql://hosta"], "conn_b": ["mysql://hostb", "mysql://hostc"], - "conn_c": [''.join("""scheme://Login:None@host:1234/lschema? + {"conn_a": ["mysql://hosta"], + "conn_b": [''.join("""scheme://Login:None@host:1234/lschema? extra__google_cloud_platform__keyfile_dict=%7B%27a%27%3A+%27b%27%7D &extra__google_cloud_platform__keyfile_path=asaa""".split())]}), ) @@ -316,6 +300,44 @@ def test_yaml_invalid_extra(self, file_content, expected_message): with self.assertRaisesRegex(AirflowException, re.escape(expected_message)): local_filesystem.load_connections("a.yaml") + @parameterized.expand( + ( + "CONN_ID=mysql://host_1/\nCONN_ID=mysql://host_2/", + ), + ) + def test_ensure_unique_connection_env(self, file_content): + with mock_local_file(file_content): + with self.assertRaises(ConnectionNotUnique): + local_filesystem.load_connections("a.env") + + @parameterized.expand( + ( + ( + {"CONN_ID": ["mysql://host_1", "mysql://host_2"]}, + ), + ( + {"CONN_ID": [{"uri": "mysql://host_1"}, {"uri": "mysql://host_2"}]}, + ), + ) + ) + def test_ensure_unique_connection_json(self, file_content): + with mock_local_file(json.dumps(file_content)): + with self.assertRaises(ConnectionNotUnique): + local_filesystem.load_connections("a.json") + + @parameterized.expand( + ( + (""" + conn_a: + - mysql://hosta + - mysql://hostb"""), + ), + ) + def test_ensure_unique_connection_yaml(self, file_content): + with mock_local_file(file_content): + with self.assertRaises(ConnectionNotUnique): + local_filesystem.load_connections("a.yaml") + class TestLocalFileBackend(unittest.TestCase): def test_should_read_variable(self): @@ -328,11 +350,11 @@ def test_should_read_variable(self): def test_should_read_connection(self): with NamedTemporaryFile(suffix=".env") as tmp_file: - tmp_file.write("CONN_A=mysql://host_a\nCONN_A=mysql://host_b".encode()) + tmp_file.write("CONN_A=mysql://host_a".encode()) tmp_file.flush() backend = LocalFilesystemBackend(connections_file_path=tmp_file.name) self.assertEqual( - ["mysql://host_a", "mysql://host_b"], + ["mysql://host_a"], [conn.get_uri() for conn in backend.get_connections("CONN_A")], ) self.assertIsNone(backend.get_variable("CONN_B")) diff --git a/tests/secrets/test_secrets_backends.py b/tests/secrets/test_secrets_backends.py index a6917cc14e141..08f8dd4573c8e 100644 --- a/tests/secrets/test_secrets_backends.py +++ b/tests/secrets/test_secrets_backends.py @@ -71,17 +71,15 @@ def test_connection_env_secrets_backend(self): self.assertEqual(sample_conn_1.host.lower(), conn.host) def test_connection_metastore_secrets_backend(self): - sample_conn_2a = SampleConn("sample_2", "A") - sample_conn_2b = SampleConn("sample_2", "B") + sample_conn_2 = SampleConn("sample_2", "A") with create_session() as session: - session.add(sample_conn_2a.conn) - session.add(sample_conn_2b.conn) + session.add(sample_conn_2.conn) session.commit() metastore_backend = MetastoreBackend() conn_list = metastore_backend.get_connections("sample_2") host_list = {x.host for x in conn_list} self.assertEqual( - {sample_conn_2a.host.lower(), sample_conn_2b.host.lower()}, set(host_list) + {sample_conn_2.host.lower()}, set(host_list) ) @mock.patch.dict('os.environ', {