Repository: PrefectHQ/prefect

Issue:
Migrate `prefect-gitlab` to core



Discussion:

Resolving patch:
diff --git a/src/integrations/prefect-gitlab/prefect_gitlab/__init__.py b/src/integrations/prefect-gitlab/prefect_gitlab/__init__.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-gitlab/prefect_gitlab/__init__.py
@@ -0,0 +1,6 @@
+from . import _version
+from .credentials import GitLabCredentials
+from .repositories import GitLabRepository
+
+__version__ = _version.__version__
+__all__ = ["GitLabRepository", "GitLabCredentials"]
diff --git a/src/integrations/prefect-gitlab/prefect_gitlab/credentials.py b/src/integrations/prefect-gitlab/prefect_gitlab/credentials.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-gitlab/prefect_gitlab/credentials.py
@@ -0,0 +1,56 @@
+"""Module used to enable authenticated interactions with GitLab"""
+
+from gitlab import Gitlab
+from pydantic import VERSION as PYDANTIC_VERSION
+
+from prefect.blocks.core import Block
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field, HttpUrl, SecretStr
+else:
+    from pydantic import Field, HttpUrl, SecretStr
+
+
+class GitLabCredentials(Block):
+    """
+    Store a GitLab personal access token to interact with private GitLab
+    repositories.
+
+    Attributes:
+        token: The personal access token to authenticate with GitLab.
+        url: URL to self-hosted GitLab instances.
+
+    Examples:
+        Load stored GitLab credentials:
+        ```python
+        from prefect_gitlab import GitLabCredentials
+        gitlab_credentials_block = GitLabCredentials.load("BLOCK_NAME")
+        ```
+    """
+
+    _block_type_name = "GitLab Credentials"
+    _logo_url = HttpUrl(
+        url="https://images.ctfassets.net/gm98wzqotmnx/55edIimT4g9gbjhkh5a3Sp/dfdb9391d8f45c2e93e72e3a4d350771/gitlab-logo-500.png?h=250",  # noqa
+        scheme="https",
+    )
+
+    token: SecretStr = Field(
+        name="Personal Access Token",
+        default=None,
+        description="A GitLab Personal Access Token with read_repository scope.",
+    )
+    url: str = Field(
+        default=None, title="URL", description="URL to self-hosted GitLab instances."
+    )
+
+    def get_client(self) -> Gitlab:
+        """
+        Gets an authenticated GitLab client.
+
+        Returns:
+            An authenticated GitLab client.
+        """
+        # ref: https://python-gitlab.readthedocs.io/en/stable/
+        gitlab = Gitlab(url=self.url, oauth_token=self.token.get_secret_value())
+        gitlab.auth()
+        return gitlab
diff --git a/src/integrations/prefect-gitlab/prefect_gitlab/repositories.py b/src/integrations/prefect-gitlab/prefect_gitlab/repositories.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-gitlab/prefect_gitlab/repositories.py
@@ -0,0 +1,216 @@
+"""
+Integrations with GitLab.
+
+The `GitLab` class in this collection is a storage block that lets Prefect agents
+pull Prefect flow code from GitLab repositories.
+
+The `GitLab` block is ideally configured via the Prefect UI, but can also be used
+in Python as the following examples demonstrate.
+
+Examples:
+```python
+    from prefect_gitlab.repositories import GitLabRepository
+
+    # public GitLab repository
+    public_gitlab_block = GitLabRepository(
+        name="my-gitlab-block",
+        repository="https://gitlab.com/testing/my-repository.git"
+    )
+
+    public_gitlab_block.save()
+
+
+    # specific branch or tag of a GitLab repository
+    branch_gitlab_block = GitLabRepository(
+        name="my-gitlab-block",
+        reference="branch-or-tag-name"
+        repository="https://gitlab.com/testing/my-repository.git"
+    )
+
+    branch_gitlab_block.save()
+
+
+    # private GitLab repository
+    private_gitlab_block = GitLabRepository(
+        name="my-private-gitlab-block",
+        repository="https://gitlab.com/testing/my-repository.git",
+        access_token="YOUR_GITLAB_PERSONAL_ACCESS_TOKEN"
+    )
+
+    private_gitlab_block.save()
+```
+"""
+import io
+import urllib.parse
+from distutils.dir_util import copy_tree
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Optional, Tuple, Union
+
+from pydantic import VERSION as PYDANTIC_VERSION
+from tenacity import retry, stop_after_attempt, wait_fixed, wait_random
+
+from prefect.exceptions import InvalidRepositoryURLError
+from prefect.filesystems import ReadableDeploymentStorage
+from prefect.utilities.asyncutils import sync_compatible
+from prefect.utilities.processutils import run_process
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import Field, HttpUrl, validator
+else:
+    from pydantic import Field, HttpUrl, validator
+
+from prefect_gitlab.credentials import GitLabCredentials
+
+# Create get_directory retry settings
+
+MAX_CLONE_ATTEMPTS = 3
+CLONE_RETRY_MIN_DELAY_SECONDS = 1
+CLONE_RETRY_MIN_DELAY_JITTER_SECONDS = 0
+CLONE_RETRY_MAX_DELAY_JITTER_SECONDS = 3
+
+
+class GitLabRepository(ReadableDeploymentStorage):
+    """
+    Interact with files stored in GitLab repositories.
+
+    An accessible installation of git is required for this block to function
+    properly.
+    """
+
+    _block_type_name = "GitLab Repository"
+    _logo_url = HttpUrl(
+        url="https://images.ctfassets.net/gm98wzqotmnx/55edIimT4g9gbjhkh5a3Sp/dfdb9391d8f45c2e93e72e3a4d350771/gitlab-logo-500.png?h=250",  # noqa
+        scheme="https",
+    )
+    _description = "Interact with files stored in GitLab repositories."
+
+    repository: str = Field(
+        default=...,
+        description=(
+            "The URL of a GitLab repository to read from, in either HTTP/HTTPS or SSH format."  # noqa
+        ),
+    )
+    reference: Optional[str] = Field(
+        default=None,
+        description="An optional reference to pin to; can be a branch name or tag.",
+    )
+    git_depth: Optional[int] = Field(
+        default=1,
+        gte=1,
+        description="The number of commits that Git history is truncated to "
+        "during cloning. Set to None to fetch the entire history.",
+    )
+    credentials: Optional[GitLabCredentials] = Field(
+        default=None,
+        description="An optional GitLab Credentials block for authenticating with "
+        "private GitLab repos.",
+    )
+
+    @validator("credentials")
+    def _ensure_credentials_go_with_http(cls, v: str, values: dict) -> str:
+        """Ensure that credentials are not provided with 'SSH' formatted GitLub URLs.
+        Note: validates `access_token` specifically so that it only fires when
+        private repositories are used.
+        """
+        if v is not None:
+            if urllib.parse.urlparse(values["repository"]).scheme not in [
+                "https",
+                "http",
+            ]:
+                raise InvalidRepositoryURLError(
+                    (
+                        "Credentials can only be used with GitLab repositories "
+                        "using the 'HTTPS'/'HTTP' format. You must either remove the "
+                        "credential if you wish to use the 'SSH' format and are not "
+                        "using a private repository, or you must change the repository "
+                        "URL to the 'HTTPS'/'HTTP' format."
+                    )
+                )
+
+        return v
+
+    def _create_repo_url(self) -> str:
+        """Format the URL provided to the `git clone` command.
+        For private repos: https://<oauth-key>@gitlab.com/<username>/<repo>.git
+        All other repos should be the same as `self.repository`.
+        """
+        url_components = urllib.parse.urlparse(self.repository)
+        if url_components.scheme in ["https", "http"] and self.credentials is not None:
+            token = self.credentials.token.get_secret_value()
+            updated_components = url_components._replace(
+                netloc=f"oauth2:{token}@{url_components.netloc}"
+            )
+            full_url = urllib.parse.urlunparse(updated_components)
+        else:
+            full_url = self.repository
+
+        return full_url
+
+    @staticmethod
+    def _get_paths(
+        dst_dir: Union[str, None], src_dir: str, sub_directory: Optional[str]
+    ) -> Tuple[str, str]:
+        """Returns the fully formed paths for GitLabRepository contents in the form
+        (content_source, content_destination).
+        """
+        if dst_dir is None:
+            content_destination = Path(".").absolute()
+        else:
+            content_destination = Path(dst_dir)
+
+        content_source = Path(src_dir)
+
+        if sub_directory:
+            content_destination = content_destination.joinpath(sub_directory)
+            content_source = content_source.joinpath(sub_directory)
+
+        return str(content_source), str(content_destination)
+
+    @sync_compatible
+    @retry(
+        stop=stop_after_attempt(MAX_CLONE_ATTEMPTS),
+        wait=wait_fixed(CLONE_RETRY_MIN_DELAY_SECONDS)
+        + wait_random(
+            CLONE_RETRY_MIN_DELAY_JITTER_SECONDS,
+            CLONE_RETRY_MAX_DELAY_JITTER_SECONDS,
+        ),
+        reraise=True,
+    )
+    async def get_directory(
+        self, from_path: Optional[str] = None, local_path: Optional[str] = None
+    ) -> None:
+        """
+        Clones a GitLab project specified in `from_path` to the provided `local_path`;
+        defaults to cloning the repository reference configured on the Block to the
+        present working directory.
+        Args:
+            from_path: If provided, interpreted as a subdirectory of the underlying
+                repository that will be copied to the provided local path.
+            local_path: A local path to clone to; defaults to present working directory.
+        """
+        # CONSTRUCT COMMAND
+        cmd = ["git", "clone", self._create_repo_url()]
+        if self.reference:
+            cmd += ["-b", self.reference]
+
+        # Limit git history
+        if self.git_depth is not None:
+            cmd += ["--depth", f"{self.git_depth}"]
+
+        # Clone to a temporary directory and move the subdirectory over
+        with TemporaryDirectory(suffix="prefect") as tmp_dir:
+            cmd.append(tmp_dir)
+
+            err_stream = io.StringIO()
+            out_stream = io.StringIO()
+            process = await run_process(cmd, stream_output=(out_stream, err_stream))
+            if process.returncode != 0:
+                err_stream.seek(0)
+                raise OSError(f"Failed to pull from remote:\n {err_stream.read()}")
+
+            content_source, content_destination = self._get_paths(
+                dst_dir=local_path, src_dir=tmp_dir, sub_directory=from_path
+            )
+
+            copy_tree(src=content_source, dst=content_destination)


Test patch:
diff --git a/src/integrations/prefect-gitlab/tests/conftest.py b/src/integrations/prefect-gitlab/tests/conftest.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-gitlab/tests/conftest.py
@@ -0,0 +1,23 @@
+import pytest
+
+from prefect.testing.utilities import prefect_test_harness
+
+
+@pytest.fixture(scope="session", autouse=True)
+def prefect_db():
+    """
+    Sets up test harness for temporary DB during test runs.
+    """
+    with prefect_test_harness():
+        yield
+
+
+@pytest.fixture(autouse=True)
+def reset_object_registry():
+    """
+    Ensures each test has a clean object registry.
+    """
+    from prefect.context import PrefectObjectRegistry
+
+    with PrefectObjectRegistry():
+        yield
diff --git a/src/integrations/prefect-gitlab/tests/test_credentials.py b/src/integrations/prefect-gitlab/tests/test_credentials.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-gitlab/tests/test_credentials.py
@@ -0,0 +1,17 @@
+from unittest.mock import MagicMock
+
+from prefect_gitlab.credentials import GitLabCredentials
+
+
+def test_gitlab_credentials_get_client(monkeypatch):
+    mock_gitlab = MagicMock()
+    monkeypatch.setattr("prefect_gitlab.credentials.Gitlab", mock_gitlab)
+    gitlab_credentials = GitLabCredentials(
+        url="https://gitlab.example.com", token="my-token"
+    )
+    gitlab_credentials.get_client()
+    mock_gitlab.assert_called_once_with(
+        url=gitlab_credentials.url,
+        oauth_token=gitlab_credentials.token.get_secret_value(),
+    )
+    mock_gitlab.assert_called_once()
diff --git a/src/integrations/prefect-gitlab/tests/test_repositories.py b/src/integrations/prefect-gitlab/tests/test_repositories.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-gitlab/tests/test_repositories.py
@@ -0,0 +1,286 @@
+import os
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import Set, Tuple
+
+import pytest
+from pydantic import VERSION as PYDANTIC_VERSION
+
+from prefect.exceptions import InvalidRepositoryURLError
+from prefect.testing.utilities import AsyncMock
+
+if PYDANTIC_VERSION.startswith("2."):
+    from pydantic.v1 import SecretStr
+else:
+    from pydantic import SecretStr
+
+import prefect_gitlab
+from prefect_gitlab.credentials import GitLabCredentials
+from prefect_gitlab.repositories import GitLabRepository  # noqa: E402
+
+
+class TestGitLab:
+    def setup_test_directory(
+        self, tmp_src: str, sub_dir: str = "puppy"
+    ) -> Tuple[Set[str], Set[str]]:
+        """Add files and directories to a temporary directory. Returns a tuple with the
+        expected parent-level contents and the expected child-level contents.
+        """
+        # add file to tmp_src
+        f1_name = "dog.text"
+        f1_path = Path(tmp_src) / f1_name
+        f1 = open(f1_path, "w")
+        f1.close()
+
+        # add sub-directory to tmp_src
+        sub_dir_path = Path(tmp_src) / sub_dir
+        os.mkdir(sub_dir_path)
+
+        # add file to sub-directory
+        f2_name = "cat.txt"
+        f2_path = sub_dir_path / f2_name
+        f2 = open(f2_path, "w")
+        f2.close()
+
+        parent_contents = {f1_name, sub_dir}
+        child_contents = {f2_name}
+
+        assert set(os.listdir(tmp_src)) == parent_contents
+        assert set(os.listdir(sub_dir_path)) == child_contents
+
+        return parent_contents, child_contents
+
+    class MockTmpDir:
+        """Utility for having `TemporaryDirectory` return a known location."""
+
+        dir = None
+
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def __enter__(self):
+            return self.dir
+
+        def __exit__(self, *args, **kwargs):
+            pass
+
+    async def test_subprocess_errors_are_surfaced(self):
+        g = GitLabRepository(repository="incorrect-url-scheme")
+        with pytest.raises(
+            OSError, match="fatal: repository 'incorrect-url-scheme' does not exist"
+        ):
+            await g.get_directory()
+
+    async def test_repository_default(self, monkeypatch):
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+        g = GitLabRepository(repository="prefect")
+        await g.get_directory()
+
+        assert mock.await_count == 1
+        expected_cmd = ["git", "clone", "prefect"]
+        assert mock.await_args[0][0][: len(expected_cmd)] == expected_cmd
+
+    async def test_reference_default(self, monkeypatch):
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+        g = GitLabRepository(repository="prefect", reference="2.0.0")
+        await g.get_directory()
+
+        assert mock.await_count == 1
+        expected_cmd = ["git", "clone", "prefect", "-b", "2.0.0", "--depth", "1"]
+        assert mock.await_args[0][0][: len(expected_cmd)] == expected_cmd
+
+    async def test_https_connection_with_token_added_correctly(self, monkeypatch):
+        """Ensure that the repo url is in the format `https://<oauth-key>@gitlab.com/<username>/<repo>.git`."""  # noqa: E501
+
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+        credential = "XYZ"
+        repo = "https://gitlab.com/PrefectHQ/prefect.git"
+        g = GitLabRepository(
+            repository=repo,
+            credentials=GitLabCredentials(token=SecretStr(credential)),
+        )
+        await g.get_directory()
+        assert mock.await_count == 1
+        expected_cmd = [
+            "git",
+            "clone",
+            f"https://oauth2:{credential}@gitlab.com/PrefectHQ/prefect.git",
+            "--depth",
+            "1",
+        ]
+        assert mock.await_args[0][0][: len(expected_cmd)] == expected_cmd
+
+    async def test_http_connection_with_token_added_correctly(self, monkeypatch):
+        """Ensure that the repo url is in the format `http://<oauth-key>@gitlab.com/<username>/<repo>.git`."""  # noqa: E501
+
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+        credential = "XYZ"
+        repo = "http://gitlab.xxx.com/PrefectHQ/prefect.git"
+        g = GitLabRepository(
+            repository=repo,
+            credentials=GitLabCredentials(token=SecretStr(credential)),
+        )
+        await g.get_directory()
+        assert mock.await_count == 1
+        expected_cmd = [
+            "git",
+            "clone",
+            f"http://oauth2:{credential}@gitlab.xxx.com/PrefectHQ/prefect.git",
+            "--depth",
+            "1",
+        ]
+        assert mock.await_args[0][0][: len(expected_cmd)] == expected_cmd
+
+    async def test_cloning_with_custom_depth(self, monkeypatch):
+        """Ensure that we can retrieve the whole history, i.e. support true git clone"""  # noqa: E501
+
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+        repo = "git@gitlab.com:PrefectHQ/prefect.git"
+        depth = None
+        g = GitLabRepository(
+            repository=repo,
+            git_depth=depth,
+        )
+        await g.get_directory()
+        assert mock.await_count == 1
+        expected_cmd = [
+            "git",
+            "clone",
+            repo,
+        ]
+        assert mock.await_args[0][0][: len(expected_cmd)] == expected_cmd
+
+    async def test_ssh_fails_with_credential(self, monkeypatch):
+        """Ensure that credentials cannot be passed in if the URL is not in the HTTPS/HTTP
+        format.
+        """
+
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+        credential = "XYZ"
+        error_msg = (
+            "Credentials can only be used with GitLab repositories "
+            "using the 'HTTPS'/'HTTP' format. You must either remove the "
+            "credential if you wish to use the 'SSH' format and are not "
+            "using a private repository, or you must change the repository "
+            "URL to the 'HTTPS'/'HTTP' format."
+        )
+        with pytest.raises(InvalidRepositoryURLError, match=error_msg):
+            GitLabRepository(
+                repository="git@gitlab.com:PrefectHQ/prefect.git",
+                credentials=GitLabCredentials(token=SecretStr(credential)),
+            )
+
+    async def test_dir_contents_copied_correctly_with_get_directory(self, monkeypatch):  # noqa
+        """Check that `get_directory` is able to correctly copy contents from src->dst"""  # noqa
+
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+
+        sub_dir_name = "puppy"
+
+        with TemporaryDirectory() as tmp_src:
+            parent_contents, child_contents = self.setup_test_directory(
+                tmp_src, sub_dir_name
+            )
+            self.MockTmpDir.dir = tmp_src
+
+            # move file contents to tmp_dst
+            with TemporaryDirectory() as tmp_dst:
+                monkeypatch.setattr(
+                    prefect_gitlab.repositories,
+                    "TemporaryDirectory",
+                    self.MockTmpDir,
+                )
+
+                g = GitLabRepository(
+                    repository="https://gitlab.com/PrefectHQ/prefect.git",
+                )
+                await g.get_directory(local_path=tmp_dst)
+
+                assert set(os.listdir(tmp_dst)) == parent_contents
+                assert set(os.listdir(Path(tmp_dst) / sub_dir_name)) == child_contents
+
+    async def test_dir_contents_copied_correctly_with_get_directory_and_from_path(
+        self, monkeypatch
+    ):  # noqa
+        """Check that `get_directory` is able to correctly copy contents from src->dst
+        when `from_path` is included.
+        It is expected that the directory specified by `from_path` will be moved to the
+        specified destination, along with all of its contents.
+        """
+
+        class p:
+            returncode = 0
+
+        mock = AsyncMock(return_value=p())
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+
+        sub_dir_name = "puppy"
+
+        with TemporaryDirectory() as tmp_src:
+            parent_contents, child_contents = self.setup_test_directory(
+                tmp_src, sub_dir_name
+            )
+            self.MockTmpDir.dir = tmp_src
+
+            # move file contents to tmp_dst
+            with TemporaryDirectory() as tmp_dst:
+                monkeypatch.setattr(
+                    prefect_gitlab.repositories,
+                    "TemporaryDirectory",
+                    self.MockTmpDir,
+                )
+
+                g = GitLabRepository(
+                    repository="https://gitlab.com/PrefectHQ/prefect.git",
+                )
+                await g.get_directory(local_path=tmp_dst, from_path=sub_dir_name)
+
+                assert set(os.listdir(tmp_dst)) == set([sub_dir_name])
+                assert set(os.listdir(Path(tmp_dst) / sub_dir_name)) == child_contents
+
+    async def test_get_directory_retries(self, monkeypatch):
+        # Constants for the retry decorator
+        MAX_CLONE_ATTEMPTS = 3
+
+        # Create an instance of GitLabRepository
+        g = GitLabRepository(repository="https://gitlab.com/prefectHQ/prefect.git")
+
+        # Prepare a MagicMock to simulate the process call within get_directory
+        mock = AsyncMock()
+        mock.return_value = AsyncMock(returncode=1)  # Simulate failure
+        monkeypatch.setattr(prefect_gitlab.repositories, "run_process", mock)
+
+        # Call get_directory and expect it to raise a RetryError after maximum attempts
+        with pytest.raises(OSError):
+            await g.get_directory()
+        print(mock.call_count)
+        # Verify that the function retried the expected number of times
+        assert mock.call_count == MAX_CLONE_ATTEMPTS
diff --git a/src/integrations/prefect-gitlab/tests/test_version.py b/src/integrations/prefect-gitlab/tests/test_version.py
new file mode 100644
--- /dev/null
+++ b/src/integrations/prefect-gitlab/tests/test_version.py
@@ -0,0 +1,9 @@
+from packaging.version import Version
+
+
+def test_version():
+    from prefect_gitlab import __version__
+
+    assert isinstance(__version__, str)
+    assert Version(__version__)
+    assert __version__.startswith("0.")