Make conn_id unique in Connections table

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

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"""

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)
+        )

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))

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)
 

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 <environment_variables_secrets_backend>`.
 Additional sources may be enabled, e.g. :ref:`AWS SSM Parameter Store <ssm_parameter_store_secrets>`, or you may
 :ref:`roll your own secrets backend <roll_your_own_secrets_backend>`.

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
 

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',

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"))

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', {