cloud-availability-updater: implement git interaction

Author: alafanechere

tools/ci_connector_ops/ci_connector_ops/qa_engine/cloud_availability_updater.py

@@ -7,88 +7,93 @@
 import logging
 from pathlib import Path
 import subprocess
-from typing import Iterable, Optional
+from typing import Optional
 
 import git
-from github import Github
 
 from .models import ConnectorQAReport
+from .constants import (
+    AIRBYTE_CLOUD_GITHUB_REPO_URL, 
+    AIRBYTE_CLOUD_MAIN_BRANCH_NAME
+)
 
-logging.basicConfig(level="INFO")
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
 
-AIRBYTE_CLOUD_REPO_PATH = Path(os.environ["AIRBYTE_CLOUD_REPO_PATH"])
-GITHUB_ACCESS_TOKEN = os.environ["GITHUB_ACCESS_TOKEN"]
-AIRBYTE_CLOUD_REPO = git.Repo(AIRBYTE_CLOUD_REPO_PATH)
-PATH_TO_DEFINITIONS_MASKS = {
-    "source": AIRBYTE_CLOUD_REPO_PATH / "cloud-config/cloud-config-seed/src/main/resources/seed/source_definitions_mask.yaml",
-    "destination": AIRBYTE_CLOUD_REPO_PATH / "cloud-config/cloud-config-seed/src/main/resources/seed/destination_definitions_mask.yaml"
-}
 
-def checkout_new_branch(connector: ConnectorQAReport) -> git.Head:
-    new_branch_name = f"cloud-availability-updater/{connector.connector_technical_name}-to-cloud"
-    new_branch = AIRBYTE_CLOUD_REPO.create_head(new_branch_name)
+def clone_airbyte_cloud_repo(local_repo_path: Path) -> git.Repo:
+    logging.info(f"Cloning {AIRBYTE_CLOUD_GITHUB_REPO_URL} to {local_repo_path}")
+    return git.Repo.clone_from(AIRBYTE_CLOUD_GITHUB_REPO_URL, local_repo_path, branch=AIRBYTE_CLOUD_MAIN_BRANCH_NAME)
+
+def get_definitions_mask_path(local_repo_path, definition_type: str) -> Path:
+    definitions_mask_path = local_repo_path / f"cloud-config/cloud-config-seed/src/main/resources/seed/{definition_type}_definitions_mask.yaml"
+    if not definitions_mask_path.exists():
+        raise FileNotFoundError(f"Can't find the {definition_type} definitions mask")
+    return definitions_mask_path
+
+def checkout_new_branch(airbyte_cloud_repo: git.Repo, new_branch_name: str) -> git.Head:
+    new_branch = airbyte_cloud_repo.create_head(new_branch_name)
     new_branch.checkout()
     logging.info(f"Checked out branch {new_branch_name}.")
     return new_branch
 
-def update_definition_mask(connector: ConnectorQAReport) -> Optional[Path]:
-    definition_mask_path = PATH_TO_DEFINITIONS_MASKS[connector.connector_type]
-    with open(definition_mask_path, "r") as definition_mask:
+def update_definitions_mask(connector: ConnectorQAReport, definitions_mask_path: Path) -> Optional[Path]:
+    with open(definitions_mask_path, "r") as definition_mask:
         connector_already_in_mask = connector.connector_definition_id in definition_mask.read()
     if connector_already_in_mask:
-        logging.warning(f"{connector.connector_name}'s definition id is already in {definition_mask_path}.")
+        logging.warning(f"{connector.connector_name}'s definition id is already in {definitions_mask_path}.")
         return None
 
     to_append = f"""# {connector.connector_name} (from cloud availability updater)
 - {connector.connector_type}DefinitionId: {connector.connector_definition_id}
 """
 
-    with open(definition_mask_path, "a") as f:
+    with open(definitions_mask_path, "a") as f:
         f.write(to_append)
-    logging.info(f"Updated {definition_mask_path} with {connector.connector_name}'s definition id.")
-    return definition_mask_path
+    logging.info(f"Updated {definitions_mask_path} with {connector.connector_name}'s definition id.")
+    return definitions_mask_path
 
-def run_generate_cloud_connector_catalog() -> str:
+def run_generate_cloud_connector_catalog(airbyte_cloud_repo_path: Path) -> str:
     result = subprocess.check_output(
-        f"cd {AIRBYTE_CLOUD_REPO_PATH} && ./gradlew :cloud-config:cloud-config-seed:generateCloudConnectorCatalog", 
+        f"cd {airbyte_cloud_repo_path} && ./gradlew :cloud-config:cloud-config-seed:generateCloudConnectorCatalog", 
         shell=True
         )
     logging.info("Ran generateCloudConnectorCatalog Gradle Task")
     return result.decode()
 
-def commit_files(connector: ConnectorQAReport):
-    AIRBYTE_CLOUD_REPO.git.add('--all')
-    AIRBYTE_CLOUD_REPO.git.commit(m=f"🤖 Add {connector.connector_technical_name} to cloud")
+def commit_all_files(airbyte_cloud_repo: git.Repo, commit_message: str):
+    airbyte_cloud_repo.git.add('--all')
+    airbyte_cloud_repo.git.commit(m=commit_message)
     logging.info(f"Committed file changes.")
 
-def push_branch(branch, dry_run=True):
-    if not dry_run:
-        AIRBYTE_CLOUD_REPO.git.push("--set-upstream", "origin", branch)
-        logging.info(f"Pushed branch {branch} to origin")
-
-def create_pr(connector: ConnectorQAReport, branch: str, dry_run=True):
-    g = Github(GITHUB_ACCESS_TOKEN)
-
-    repo = g.get_repo("airbytehq/airbyte-cloud")
-    body = f"""
-    The Cloud Availability Updater decided that it's the right time to make {connector.connector_name} available on Cloud!
-    ```
-    {connector}
-    ```
+def push_branch(airbyte_cloud_repo: git.Repo, branch:str):
+    airbyte_cloud_repo.git.push("--set-upstream", "origin", branch)
+    logging.info(f"Pushed branch {branch} to origin")
+
+def deploy_new_connector_to_cloud_repo(
+    airbyte_cloud_repo_path: Path,
+    airbyte_cloud_repo: git.Repo,
+    connector: ConnectorQAReport
+ ):
+    """Updates the local definitions mask on Airbyte cloud repo.
+    Calls the generateCloudConnectorCatalog gradle task.
+    Commits these changes on a new branch.
+    Pushes these new branch to the origin.
+
+    Args:
+        airbyte_cloud_repo_path (Path): The local path to Airbyte Cloud repository.
+        airbyte_cloud_repo (git.Repo): The Airbyte Cloud repo instance.
+        connector (ConnectorQAReport): The connector to add to a definitions mask.
     """
-    if not dry_run:
-        return repo.create_pull(title=f"🤖 Add {connector.connector_technical_name} to cloud", body=body, head=branch, base="master")
-    
-def deploy_eligible_connector(connector: ConnectorQAReport, dry_run=True):
-    AIRBYTE_CLOUD_REPO.heads.master.checkout()
-    new_branch = checkout_new_branch(connector)
-    update_definition_mask(connector)
-    run_generate_cloud_connector_catalog()
-    commit_files(connector)
-    push_branch(new_branch, dry_run=dry_run)
-    create_pr(connector, new_branch, dry_run=dry_run)
-
-def deploy_eligible_connectors(connectors_eligible_for_cloud: Iterable[ConnectorQAReport]):
-    for connector in connectors_eligible_for_cloud:
-        deploy_eligible_connector(connector)
-    AIRBYTE_CLOUD_REPO.heads.master.checkout()
+    airbyte_cloud_repo.git.checkout(AIRBYTE_CLOUD_MAIN_BRANCH_NAME)
+    new_branch_name = f"cloud-availability-updater/deploy-{connector.connector_technical_name}"
+    checkout_new_branch(airbyte_cloud_repo, new_branch_name)
+    definitions_mask_path = get_definitions_mask_path(airbyte_cloud_repo_path, connector.connector_type)
+    update_definitions_mask(connector, definitions_mask_path)
+    run_generate_cloud_connector_catalog(airbyte_cloud_repo_path)
+    commit_all_files(
+        airbyte_cloud_repo, 
+        f"🤖 Add {connector.connector_name} connector to cloud"
+    )
+    push_branch(airbyte_cloud_repo, new_branch_name)
+    airbyte_cloud_repo.git.checkout(AIRBYTE_CLOUD_MAIN_BRANCH_NAME)

tools/ci_connector_ops/ci_connector_ops/qa_engine/constants.py

@@ -13,3 +13,5 @@
 ]
 
 GCS_QA_REPORT_PATH = "gs://prod-airbyte-cloud-connector-metadata-service/qa_report.json"
+AIRBYTE_CLOUD_GITHUB_REPO_URL = "https://github.com/airbytehq/airbyte-cloud.git"
+AIRBYTE_CLOUD_MAIN_BRANCH_NAME = "master"

tools/ci_connector_ops/ci_connector_ops/qa_engine/models.py

@@ -3,9 +3,8 @@
 #
 
 
-from datetime import datetime
 from enum import Enum
-from typing import List, Optional
+from typing import List
 
 from pydantic import BaseModel
 
@@ -30,8 +29,7 @@ class ConnectorQAReport(BaseModel):
     is_appropriate_for_cloud_use: bool
     latest_build_is_successful: bool
     documentation_is_available: bool
-    is_eligible_for_cloud: bool
-    report_generation_datetime: datetime
+
 
 class QAReport(BaseModel):
     connectors_qa_report: List[ConnectorQAReport]

tools/ci_connector_ops/pytest.ini

@@ -0,0 +1,4 @@
+[pytest]
+markers =
+    slow: marks tests as slow (deselect with '-m "not slow"')
+    serial

tools/ci_connector_ops/setup.py

@@ -12,7 +12,6 @@
     "pandas~=1.5.3",
     "pandas-gbq~=0.19.0",
     "pydantic~=1.10.4",
-    "PyGithub"
     "fsspec~=2023.1.0",
     "gcsfs~=2023.1.0"
 ]
@@ -24,7 +23,7 @@
 
 
 setup(
-    version="0.1.6",
+    version="0.1.9",
     name="ci_connector_ops",
     description="Packaged maintained by the connector operations team to perform CI for connectors",
     author="Airbyte",

tools/ci_connector_ops/tests/test_qa_engine/test_cloud_availability_updater.py

@@ -0,0 +1,145 @@
+#
+# Copyright (c) 2023 Airbyte, Inc., all rights reserved.
+#
+
+
+from datetime import datetime
+from pathlib import Path
+
+import pytest
+import git
+import yaml
+
+from ci_connector_ops.qa_engine import cloud_availability_updater, models
+
+@pytest.fixture(scope="module")
+def dummy_repo_path(tmp_path_factory) -> Path:
+    repo_path = tmp_path_factory.mktemp("cloud_availability_updater_tests") / "airbyte-cloud"
+    repo_path.mkdir()
+    return repo_path
+
+@pytest.fixture(scope="module")
+def dummy_repo(dummy_repo_path) -> git.Repo:
+    seed_dir = dummy_repo_path / "cloud-config/cloud-config-seed/src/main/resources/seed"
+    seed_dir.mkdir(parents=True)
+    repo = git.Repo.init(dummy_repo_path)
+    source_definitions_mask_path = seed_dir / "source_definitions_mask.yaml"
+    destination_definitions_mask_path = seed_dir / "destination_definitions_mask.yaml"
+    source_definitions_mask_path.touch()
+    destination_definitions_mask_path.touch()
+    repo.git.add("--all")
+    repo.git.commit(m=f"🤖 Initialized the repo")
+    return repo
+
+
+@pytest.fixture
+def checkout_master(dummy_repo):
+    """
+    Ensure we're always on dummy repo master before and after each test using this fixture
+    """
+    yield dummy_repo.heads.master.checkout()
+    dummy_repo.heads.master.checkout()
+
+def test_get_definitions_mask_path(checkout_master, dummy_repo_path: Path):
+    path = cloud_availability_updater.get_definitions_mask_path(dummy_repo_path, "source")
+    assert path.exists() and path.name == "source_definitions_mask.yaml"
+    path = cloud_availability_updater.get_definitions_mask_path(dummy_repo_path, "destination")
+    assert path.exists() and path.name == "destination_definitions_mask.yaml"
+    with pytest.raises(FileNotFoundError):
+        cloud_availability_updater.get_definitions_mask_path(dummy_repo_path, "foobar")
+
+def test_checkout_new_branch(mocker, checkout_master, dummy_repo):
+    new_branch = cloud_availability_updater.checkout_new_branch(dummy_repo, "test-branch")
+    assert new_branch.name == dummy_repo.active_branch.name == "test-branch"
+
+
+@pytest.mark.parametrize(
+    "definitions_mask_content_before_update, definition_id, expect_update",
+    [
+        ("", "abcdefg", True),
+        ("abcdefg", "abcdefg", False),
+    ]
+
+)
+def test_update_definitions_mask(
+    mocker,
+    tmp_path, 
+    definitions_mask_content_before_update,
+    definition_id, 
+    expect_update
+):
+    connector = mocker.Mock(
+        connector_name="foobar",
+        connector_definition_id=definition_id,
+        connector_type="unknown"
+    )
+    definitions_mask_path = tmp_path / "definitions_mask.yaml"
+    with open(definitions_mask_path, "w") as definitions_mask:
+        definitions_mask.write(definitions_mask_content_before_update)
+    updated_path = cloud_availability_updater.update_definitions_mask(connector, definitions_mask_path)
+    if not expect_update:
+        assert updated_path is None
+    else:
+        with open(updated_path, "r") as definitions_mask:
+            raw_content = definitions_mask.read()
+            definitions = yaml.safe_load(raw_content)
+        assert isinstance(definitions, list)
+        assert definitions[0]["unknownDefinitionId"] == definition_id
+        assert len(
+        [
+            d for d in definitions 
+            if d["unknownDefinitionId"] == definition_id
+        ]) == 1
+        assert "# foobar (from cloud availability updater)" in raw_content
+        assert raw_content[-1] == "\n"
+
+def test_commit_files(checkout_master, dummy_repo, dummy_repo_path):
+    cloud_availability_updater.checkout_new_branch(dummy_repo, "test-commit-files")
+    commit_message = "🤖 Add new connector to cloud"
+    with open(dummy_repo_path / "test_file.txt", "w") as f:
+        f.write(".")
+    
+    cloud_availability_updater.commit_all_files(dummy_repo, commit_message)
+    
+    assert dummy_repo.head.reference.commit.message == commit_message + "\n"
+    edited_files = dummy_repo.git.diff("--name-only", checkout_master.name).split("\n")
+    assert "test_file.txt" in edited_files
+
+def test_push_branch(mocker):
+    mock_repo = mocker.Mock()
+    cloud_availability_updater.push_branch(mock_repo, "new_branch")
+    mock_repo.git.push.assert_called_once_with("--set-upstream", "origin", "new_branch")
+
+@pytest.mark.slow
+def test_deploy_new_connector_to_cloud_repo(mocker, tmp_path):
+    mocker.patch.object(cloud_availability_updater, "push_branch")
+    mocker.patch.object(cloud_availability_updater, "run_generate_cloud_connector_catalog")
+
+    repo_path = tmp_path / "airbyte-cloud"
+    repo_path.mkdir()
+    airbyte_cloud_repo = cloud_availability_updater.clone_airbyte_cloud_repo(repo_path)
+    source_definitions_mask_path = repo_path / "cloud-config/cloud-config-seed/src/main/resources/seed/source_definitions_mask.yaml"
+    destination_definitions_mask_path = repo_path / "cloud-config/cloud-config-seed/src/main/resources/seed/destination_definitions_mask.yaml"
+    assert source_definitions_mask_path.exists() and destination_definitions_mask_path.exists()
+    
+    connector = models.ConnectorQAReport(
+        connector_type="source",
+        connector_name="foobar",
+        connector_technical_name="source-foobar",
+        connector_definition_id="abcdefg",
+        connector_version="0.0.0",
+        release_stage="alpha",
+        is_on_cloud=False,
+        is_appropriate_for_cloud_use=True,
+        latest_build_is_successful=True,
+        documentation_is_available=True
+    )
+    cloud_availability_updater.deploy_new_connector_to_cloud_repo(repo_path, airbyte_cloud_repo, connector)
+    new_branch_name = f"cloud-availability-updater/deploy-{connector.connector_technical_name}"
+
+    cloud_availability_updater.push_branch.assert_called_once_with(airbyte_cloud_repo, new_branch_name)
+    cloud_availability_updater.run_generate_cloud_connector_catalog.assert_called_once_with(repo_path)
+    airbyte_cloud_repo.git.checkout(new_branch_name)
+    edited_files = airbyte_cloud_repo.git.diff("--name-only", "master").split("\n")
+    assert edited_files == ['cloud-config/cloud-config-seed/src/main/resources/seed/source_definitions_mask.yaml']
+    assert airbyte_cloud_repo.head.reference.commit.message == "🤖 Add foobar connector to cloud\n"