Misc code repository improvements

docs/book/how-to/project-setup-and-management/setting-up-a-project-repository/connect-your-git-repository.md

@@ -48,11 +48,13 @@ Afterward, you can register a GitHub code repository by running the following CL
 
 ```shell
 zenml code-repository register <NAME> --type=github \
---url=<GITHUB_URL> --owner=<OWNER> --repository=<REPOSITORY> \
+--owner=<OWNER> --repository=<REPOSITORY> \
 --token=<GITHUB_TOKEN>
 ```
 
-where \<REPOSITORY> is the name of the code repository you are registering, \<OWNER> is the owner of the repository, \<NAME> is the name of the repository, \<GITHUB\_TOKEN> is your GitHub Personal Access Token and \<GITHUB\_URL> is the URL of the GitHub instance which defaults to `https://github.com.` You will need to set a URL if you are using GitHub Enterprise.
+where \<REPOSITORY> is the name of the code repository you are registering, \<OWNER> is the owner of the repository, \<NAME> is the name of the repository, \<GITHUB\_TOKEN> is your GitHub Personal Access Token.
+
+If you're using a self-hosted GitHub Enterprise instance, you'll need to also pass the `--api_url=<API_URL>` and `--host=<HOST>` options. \<API_URL> should point to where the GitHub API is reachable (defaults to `https://api.github.com/`) and \<HOST> should be the [hostname of your GitHub instance](https://docs.github.com/en/[email protected]/admin/configuring-settings/configuring-network-settings/configuring-the-hostname-for-your-instance?learn=deploy_an_instance&learnProduct=admin).
 
 {% hint style="warning" %}
 Please refer to the section on using secrets for stack configuration in order to securely store your GitHub 
@@ -105,11 +107,14 @@ Afterward, you can register a GitLab code repository by running the following CL
 
 ```shell
 zenml code-repository register <NAME> --type=gitlab \
---url=<GITLAB_URL> --group=<GROUP> --project=<PROJECT> \
+--group=<GROUP> --project=<PROJECT> \
 --token=<GITLAB_TOKEN>
 ```
 
-where `<NAME>` is the name of the code repository you are registering, `<GROUP>` is the group of the project, `<PROJECT>` is the name of the project, \<GITLAB\_TOKEN> is your GitLab Personal Access Token, and \<GITLAB\_URL> is the URL of the GitLab instance which defaults to `https://gitlab.com.` You will need to set a URL if you have a self-hosted GitLab instance.
+where `<NAME>` is the name of the code repository you are registering, `<GROUP>` is the group of the project, `<PROJECT>` is the name of the project, \<GITLAB\_TOKEN> is your GitLab Personal Access Token.
+
+If you're using a self-hosted GitLab instance, you'll need to also pass the `--instance_url=<INSTANCE_URL>` and `--host=<HOST>` options. \<INSTANCE_URL> should point to your GitLab instance (defaults to `https://gitlab.com/`) and \<HOST> should be the hostname of your GitLab instance (defaults to `gitlab.com`).
+
 
 {% hint style="warning" %}
 Please refer to the section on using secrets for stack configuration in order to securely store your GitLab

docs/book/user-guide/production-guide/connect-code-repository.md

@@ -82,12 +82,11 @@ Now, we can install the GitHub integration and register your repository:
 ```sh
 zenml integration install github
 zenml code-repository register <REPO_NAME> --type=github \
---url=https://github.com/YOUR_USERNAME/YOUR_REPOSITORY_NAME.git \
---owner=YOUR_USERNAME --repository=YOUR_REPOSITORY_NAME \
---token=YOUR_GITHUB_PERSONAL_ACCESS_TOKEN
+--owner=<YOUR_USERNAME> --repository=<YOUR_REPOSITORY_NAME> \
+--token=<YOUR_GITHUB_PERSONAL_ACCESS_TOKEN>
 ```
 
-Fill in `<REPO_NAME>`, `YOUR_USERNAME`, `YOUR_REPOSITORY_NAME`, and `YOUR_GITHUB_PERSONAL_ACCESS_TOKEN` with your details.
+Fill in `<REPO_NAME>`, `<YOUR_USERNAME>`, `<YOUR_REPOSITORY_NAME>`, and `<YOUR_GITHUB_PERSONAL_ACCESS_TOKEN>` with your details.
 
 Your code is now connected to your ZenML server. ZenML will automatically detect if your source files are being tracked by GitHub and store the commit hash for each subsequent pipeline run.
 

src/zenml/cli/code_repository.py

@@ -42,7 +42,7 @@ def code_repository() -> None:
     context_settings={"ignore_unknown_options": True},
     help="Register a code repository.",
 )
-@click.argument("name", type=click.STRING)
+@click.argument("name", type=str)
 @click.option(
     "--type",
     "-t",
@@ -183,6 +183,74 @@ def list_code_repositories(**kwargs: Any) -> None:
         )
 
 
+@code_repository.command(
+    "update",
+    help="Update a code repository.",
+    context_settings={"ignore_unknown_options": True},
+)
+@click.argument("name_or_id", type=str, required=True)
+@click.option(
+    "--name",
+    "-n",
+    type=str,
+    required=False,
+    help="The new code repository name.",
+)
+@click.option(
+    "--description",
+    "-d",
+    type=str,
+    required=False,
+    help="The new code repository description.",
+)
+@click.option(
+    "--logo-url",
+    "-l",
+    type=str,
+    required=False,
+    help="New URL of a logo (png, jpg or svg) for the code repository.",
+)
+@click.argument(
+    "args",
+    nargs=-1,
+    type=click.UNPROCESSED,
+)
+def update_code_repository(
+    name_or_id: str,
+    name: Optional[str],
+    description: Optional[str],
+    logo_url: Optional[str],
+    args: List[str],
+) -> None:
+    """Update a code repository.
+
+    Args:
+        name_or_id: Name or ID of the code repository to update.
+        name: New name of the code repository.
+        description: New description of the code repository.
+        logo_url: New logo URL of the code repository.
+        args: Code repository configurations.
+    """
+    parsed_name_or_id, parsed_args = cli_utils.parse_name_and_extra_arguments(
+        list(args) + [name_or_id], expand_args=True, name_mandatory=True
+    )
+    assert parsed_name_or_id
+
+    with console.status(
+        f"Updating code repository '{parsed_name_or_id}'...\n"
+    ):
+        Client().update_code_repository(
+            name_id_or_prefix=parsed_name_or_id,
+            name=name,
+            description=description,
+            logo_url=logo_url,
+            config=parsed_args,
+        )
+        cli_utils.declare(
+            f"Successfully updated code repository `{parsed_name_or_id}`."
+        )
+
+
 @code_repository.command("delete")
 @click.argument("name_or_id", type=str, required=True)
 @click.option(

src/zenml/client.py

@@ -4952,25 +4952,15 @@ def restore_secrets(
 
     # --------------------------- Code repositories ---------------------------
 
-    def create_code_repository(
-        self,
-        name: str,
-        config: Dict[str, Any],
-        source: Source,
-        description: Optional[str] = None,
-        logo_url: Optional[str] = None,
-    ) -> CodeRepositoryResponse:
-        """Create a new code repository.
+    @staticmethod
+    def _validate_code_repository_config(
+        source: Source, config: Dict[str, Any]
+    ) -> None:
+        """Validate a code repository config.
 
         Args:
-            name: Name of the code repository.
-            config: The configuration for the code repository.
-            source: The code repository implementation source.
-            description: The code repository description.
-            logo_url: URL of a logo (png, jpg or svg) for the code repository.
-
-        Returns:
-            The created code repository.
+            source: The code repository source.
+            config: The code repository config.
 
         Raises:
             RuntimeError: If the provided config is invalid.
@@ -4983,13 +4973,38 @@ def create_code_repository(
             )
         )
         try:
-            # Validate the repo config
+            # This does a login to verify the credentials
             code_repo_class(id=uuid4(), config=config)
+
+            # Explicitly access the config for pydantic validation, in case
+            # the login for some reason did not do that.
+            _ = code_repo_class.config
         except Exception as e:
             raise RuntimeError(
                 "Failed to validate code repository config."
             ) from e
 
+    def create_code_repository(
+        self,
+        name: str,
+        config: Dict[str, Any],
+        source: Source,
+        description: Optional[str] = None,
+        logo_url: Optional[str] = None,
+    ) -> CodeRepositoryResponse:
+        """Create a new code repository.
+
+        Args:
+            name: Name of the code repository.
+            config: The configuration for the code repository.
+            source: The code repository implementation source.
+            description: The code repository description.
+            logo_url: URL of a logo (png, jpg or svg) for the code repository.
+
+        Returns:
+            The created code repository.
+        """
+        self._validate_code_repository_config(source=source, config=config)
         repo_request = CodeRepositoryRequest(
             user=self.active_user.id,
             workspace=self.active_workspace.id,
@@ -5088,6 +5103,7 @@ def update_code_repository(
         name: Optional[str] = None,
         description: Optional[str] = None,
         logo_url: Optional[str] = None,
+        config: Optional[Dict[str, Any]] = None,
     ) -> CodeRepositoryResponse:
         """Update a code repository.
 
@@ -5097,6 +5113,10 @@ def update_code_repository(
             name: New name of the code repository.
             description: New description of the code repository.
             logo_url: New logo URL of the code repository.
+            config: New configuration options for the code repository. Will
+                be used to update the existing configuration values. To remove
+                values from the existing configuration, set the value for that
+                key to `None`.
 
         Returns:
             The updated code repository.
@@ -5107,6 +5127,18 @@ def update_code_repository(
         update = CodeRepositoryUpdate(
             name=name, description=description, logo_url=logo_url
         )
+        if config is not None:
+            combined_config = repo.config
+            combined_config.update(config)
+            combined_config = {
+                k: v for k, v in combined_config.items() if v is not None
+            }
+
+            self._validate_code_repository_config(
+                source=repo.source, config=combined_config
+            )
+            update.config = combined_config
+
         return self.zen_store.update_code_repository(
             code_repository_id=repo.id, update=update
         )

src/zenml/entrypoints/base_entrypoint_configuration.py

@@ -214,12 +214,14 @@ def download_code_if_necessary(
         if not should_download_code:
             return
 
-        if code_reference := deployment.code_reference:
+        if code_path := deployment.code_path:
+            code_utils.download_code_from_artifact_store(code_path=code_path)
+        elif code_reference := deployment.code_reference:
+            # TODO: This might fail if the code repository had unpushed changes
+            # at the time the pipeline run was started.
             self.download_code_from_code_repository(
                 code_reference=code_reference
             )
-        elif code_path := deployment.code_path:
-            code_utils.download_code_from_artifact_store(code_path=code_path)
         else:
             raise RuntimeError(
                 "Code download required but no code reference or path provided."

src/zenml/integrations/github/code_repositories/github_code_repository.py

@@ -30,6 +30,7 @@
 )
 from zenml.code_repositories.git import LocalGitRepositoryContext
 from zenml.logger import get_logger
+from zenml.utils import deprecation_utils
 from zenml.utils.secret_utils import SecretField
 
 logger = get_logger(__name__)
@@ -39,19 +40,23 @@ class GitHubCodeRepositoryConfig(BaseCodeRepositoryConfig):
     """Config for GitHub code repositories.
 
     Args:
-        url: The URL of the GitHub instance.
+        api_url: The GitHub API URL.
         owner: The owner of the repository.
         repository: The name of the repository.
         host: The host of the repository.
         token: The token to access the repository.
     """
 
-    url: Optional[str]
+    api_url: Optional[str] = None
     owner: str
     repository: str
     host: Optional[str] = "github.com"
     token: Optional[str] = SecretField(default=None)
 
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        ("url", "api_url")
+    )
+
 
 class GitHubCodeRepository(BaseCodeRepository):
     """GitHub code repository."""
@@ -120,7 +125,9 @@ def login(
             RuntimeError: If the login fails.
         """
         try:
-            self._github_session = Github(self.config.token)
+            self._github_session = Github(
+                login_or_token=self.config.token, base_url=self.config.api_url
+            )
             if self.config.token:
                 user = self._github_session.get_user().login
                 logger.debug(f"Logged in as {user}")

src/zenml/integrations/gitlab/code_repositories/gitlab_code_repository.py

@@ -31,6 +31,7 @@
     LocalGitRepositoryContext,
 )
 from zenml.logger import get_logger
+from zenml.utils import deprecation_utils
 from zenml.utils.secret_utils import SecretField
 
 logger = get_logger(__name__)
@@ -40,19 +41,23 @@ class GitLabCodeRepositoryConfig(BaseCodeRepositoryConfig):
     """Config for GitLab code repositories.
 
     Args:
-        url: The full URL of the GitLab project.
+        instance_url: The URL of the GitLab instance.
         group: The group of the project.
         project: The name of the GitLab project.
         host: The host of GitLab in case it is self-hosted instance.
         token: The token to access the repository.
     """
 
-    url: Optional[str]
+    instance_url: Optional[str] = None
     group: str
     project: str
     host: Optional[str] = "gitlab.com"
     token: str = SecretField()
 
+    _deprecation_validator = deprecation_utils.deprecate_pydantic_attributes(
+        ("url", "instance_url")
+    )
+
 
 class GitLabCodeRepository(BaseCodeRepository):
     """GitLab code repository."""
@@ -85,7 +90,7 @@ def login(self) -> None:
         """
         try:
             self._gitlab_session = Gitlab(
-                self.config.url, private_token=self.config.token
+                url=self.config.instance_url, private_token=self.config.token
             )
             self._gitlab_session.auth()
             user = self._gitlab_session.user or None