Skip to content

Commit

Permalink
Add credentials abstract block (#7908)
Browse files Browse the repository at this point in the history
  • Loading branch information
ahuang11 committed Dec 20, 2022
1 parent ec06418 commit 794b1b0
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 0 deletions.
36 changes: 36 additions & 0 deletions src/prefect/blocks/abstract.py
Expand Up @@ -12,6 +12,42 @@
T = TypeVar("T")


class CredentialsBlock(Block, ABC):
"""
Stores credentials for an external system and exposes a client for interacting
with that system. Can also hold config that is tightly coupled to credentials
(domain, endpoint, account ID, etc.) Will often be composed with other blocks.
Parent block should rely on the client provided by a credentials block for
interacting with the corresponding external system.
"""

@property
def logger(self) -> Logger:
"""
Returns a logger based on whether the CredentialsBlock
is called from within a flow or task run context.
If a run context is present, the logger property returns a run logger.
Else, it returns a default logger labeled with the class's name.
Returns:
The run logger or a default logger with the class's name.
"""
try:
return get_run_logger()
except MissingContextError:
return get_logger(self.__class__.__name__)

@abstractmethod
def get_client(self, *args, **kwargs):
"""
Returns a client for interacting with the external system.
If a service offers various clients, this method can accept
a `client_type` keyword argument to get the desired client
within the service.
"""


class NotificationBlock(Block, ABC):
"""
Block that represents a resource in an external system that is able to send notifications.
Expand Down
25 changes: 25 additions & 0 deletions tests/blocks/test_abstract.py
@@ -1,6 +1,7 @@
import pytest

from prefect.blocks.abstract import (
CredentialsBlock,
DatabaseBlock,
JobBlock,
JobRun,
Expand All @@ -11,6 +12,30 @@
from prefect.exceptions import PrefectException


class TestCredentialsBlock:
def test_credentials_block_is_abstract(self):
with pytest.raises(
TypeError, match="Can't instantiate abstract class CredentialsBlock"
):
CredentialsBlock()

def test_credentials_block_implementation(self, caplog):
class ACredentialsBlock(CredentialsBlock):
def get_client(self):
self.logger.info("Got client.")
return "client"

a_credentials_block = ACredentialsBlock()
assert a_credentials_block.get_client() == "client"

# test logging
assert hasattr(a_credentials_block, "logger")
assert len(caplog.records) == 1
record = caplog.records[0]
assert record.name == "prefect.ACredentialsBlock"
assert record.msg == "Got client."


class TestNotificationBlock:
def test_notification_block_is_abstract(self):
with pytest.raises(
Expand Down

0 comments on commit 794b1b0

Please sign in to comment.