fal 0.9.2__py3-none-any.whl → 0.9.4__py3-none-any.whl

dbt/fal/adapters/python/connections.py

@@ -0,0 +1,319 @@
+import abc
+import os
+from time import sleep
+import sys
+from multiprocessing.synchronize import RLock
+from threading import get_ident
+from typing import (
+    Any,
+    Dict,
+    Tuple,
+    Hashable,
+    Optional,
+    List,
+    Type,
+    Union,
+    Iterable,
+    Callable,
+)
+
+from dbt.exceptions import (
+    NotImplementedError,
+    InvalidConnectionError,
+    DbtInternalError,
+    CompilationError,
+    FailedToConnectError,
+)
+from dbt.contracts.connection import (
+    Connection,
+    Identifier,
+    ConnectionState,
+    AdapterRequiredConfig,
+    LazyHandle,
+    AdapterResponse,
+)
+from dbt.events import AdapterLogger
+from dbt.events.functions import fire_event
+from dbt.events.types import (
+    NewConnection,
+    ConnectionReused,
+    ConnectionLeftOpen,
+    ConnectionClosed,
+)
+from dbt import flags
+
+SleepTime = Union[int, float]  # As taken by time.sleep.
+AdapterHandle = Any  # Adapter connection handle objects can be any class.
+
+
+class PythonConnectionManager(metaclass=abc.ABCMeta):
+    """Methods to implement:
+        - open
+        - execute
+
+    You must also set the 'TYPE' class attribute with a class-unique constant
+    string.
+    """
+
+    TYPE: str = NotImplemented
+
+    def __init__(self, profile: AdapterRequiredConfig):
+        self.profile = profile
+        if profile.credentials.type == self.TYPE:
+            self.credentials = profile.credentials
+        else:
+            self.credentials = profile.python_adapter_credentials
+
+        self.thread_connections: Dict[Hashable, Connection] = {}
+        self.lock: RLock = flags.MP_CONTEXT.RLock()
+
+    @staticmethod
+    def get_thread_identifier() -> Hashable:
+        # note that get_ident() may be re-used, but we should never experience
+        # that within a single process
+        return (os.getpid(), get_ident())
+
+    def get_thread_connection(self) -> Connection:
+        key = self.get_thread_identifier()
+        with self.lock:
+            if key not in self.thread_connections:
+                raise InvalidConnectionError(key, list(self.thread_connections))
+            return self.thread_connections[key]
+
+    def set_thread_connection(self, conn: Connection) -> None:
+        key = self.get_thread_identifier()
+        if key in self.thread_connections:
+            raise DbtInternalError(
+                "In set_thread_connection, existing connection exists for {}"
+            )
+        self.thread_connections[key] = conn
+
+    def get_if_exists(self) -> Optional[Connection]:
+        key = self.get_thread_identifier()
+        with self.lock:
+            return self.thread_connections.get(key)
+
+    def clear_thread_connection(self) -> None:
+        key = self.get_thread_identifier()
+        with self.lock:
+            if key in self.thread_connections:
+                del self.thread_connections[key]
+
+    def set_connection_name(self, name: Optional[str] = None) -> Connection:
+        conn_name: str
+        if name is None:
+            # if a name isn't specified, we'll re-use a single handle
+            # named 'master'
+            conn_name = "master"
+        else:
+            if not isinstance(name, str):
+                raise CompilationError(
+                    f"For connection name, got {name} - not a string!"
+                )
+            assert isinstance(name, str)
+            conn_name = name
+
+        conn = self.get_if_exists()
+        if conn is None:
+            conn = Connection(
+                type=Identifier(self.TYPE),
+                name=None,
+                state=ConnectionState.INIT,
+                transaction_open=False,
+                handle=None,
+                credentials=self.credentials,
+            )
+            self.set_thread_connection(conn)
+
+        if conn.name == conn_name and conn.state == "open":
+            return conn
+
+        fire_event(NewConnection(conn_name=conn_name, conn_type=self.TYPE))
+
+        if conn.state == "open":
+            fire_event(ConnectionReused(conn_name=conn_name))
+        else:
+            conn.handle = LazyHandle(self.open)
+
+        conn.name = conn_name
+        return conn
+
+    @classmethod
+    def retry_connection(
+        cls,
+        connection: Connection,
+        connect: Callable[[], AdapterHandle],
+        logger: AdapterLogger,
+        retryable_exceptions: Iterable[Type[Exception]],
+        retry_limit: int = 1,
+        retry_timeout: Union[Callable[[int], SleepTime], SleepTime] = 1,
+        _attempts: int = 0,
+    ) -> Connection:
+        """Given a Connection, set its handle by calling connect.
+
+        The calls to connect will be retried up to retry_limit times to deal with transient
+        connection errors. By default, one retry will be attempted if retryable_exceptions is set.
+
+        :param Connection connection: An instance of a Connection that needs a handle to be set,
+            usually when attempting to open it.
+        :param connect: A callable that returns the appropiate connection handle for a
+            given adapter. This callable will be retried retry_limit times if a subclass of any
+            Exception in retryable_exceptions is raised by connect.
+        :type connect: Callable[[], AdapterHandle]
+        :param AdapterLogger logger: A logger to emit messages on retry attempts or errors. When
+            handling expected errors, we call debug, and call warning on unexpected errors or when
+            all retry attempts have been exhausted.
+        :param retryable_exceptions: An iterable of exception classes that if raised by
+            connect should trigger a retry.
+        :type retryable_exceptions: Iterable[Type[Exception]]
+        :param int retry_limit: How many times to retry the call to connect. If this limit
+            is exceeded before a successful call, a FailedToConnectError will be raised.
+            Must be non-negative.
+        :param retry_timeout: Time to wait between attempts to connect. Can also take a
+            Callable that takes the number of attempts so far, beginning at 0, and returns an int
+            or float to be passed to time.sleep.
+        :type retry_timeout: Union[Callable[[int], SleepTime], SleepTime] = 1
+        :param int _attempts: Parameter used to keep track of the number of attempts in calling the
+            connect function across recursive calls. Passed as an argument to retry_timeout if it
+            is a Callable. This parameter should not be set by the initial caller.
+        :raises dbt.exceptions.FailedToConnectError: Upon exhausting all retry attempts without
+            successfully acquiring a handle.
+        :return: The given connection with its appropriate state and handle attributes set
+            depending on whether we successfully acquired a handle or not.
+        """
+        timeout = retry_timeout(_attempts) if callable(retry_timeout) else retry_timeout
+        if timeout < 0:
+            raise FailedToConnectError(
+                "retry_timeout cannot be negative or return a negative time."
+            )
+
+        if retry_limit < 0 or retry_limit > sys.getrecursionlimit():
+            # This guard is not perfect others may add to the recursion limit (e.g. built-ins).
+            connection.handle = None
+            connection.state = ConnectionState.FAIL
+            raise FailedToConnectError("retry_limit cannot be negative")
+
+        try:
+            connection.handle = connect()
+            connection.state = ConnectionState.OPEN
+            return connection
+
+        except tuple(retryable_exceptions) as e:
+            if retry_limit <= 0:
+                connection.handle = None
+                connection.state = ConnectionState.FAIL
+                raise FailedToConnectError(str(e))
+
+            logger.debug(
+                f"Got a retryable error when attempting to open a {cls.TYPE} connection.\n"
+                f"{retry_limit} attempts remaining. Retrying in {timeout} seconds.\n"
+                f"Error:\n{e}"
+            )
+
+            sleep(timeout)
+            return cls.retry_connection(
+                connection=connection,
+                connect=connect,
+                logger=logger,
+                retry_limit=retry_limit - 1,
+                retry_timeout=retry_timeout,
+                retryable_exceptions=retryable_exceptions,
+                _attempts=_attempts + 1,
+            )
+
+        except Exception as e:
+            connection.handle = None
+            connection.state = ConnectionState.FAIL
+            raise FailedToConnectError(str(e))
+
+    @abc.abstractmethod
+    def cancel(self, connection: Connection):
+        """Cancel the given connection."""
+        raise NotImplementedError("`cancel` is not implemented for this adapter!")
+
+    def cancel_open(self) -> List[str]:
+        names = []
+        this_connection = self.get_if_exists()
+        with self.lock:
+            for connection in self.thread_connections.values():
+                if connection is this_connection:
+                    continue
+
+                # if the connection failed, the handle will be None so we have
+                # nothing to cancel.
+                if (
+                    connection.handle is not None
+                    and connection.state == ConnectionState.OPEN
+                ):
+                    self.cancel(connection)
+                if connection.name is not None:
+                    names.append(connection.name)
+        return names
+
+    @classmethod
+    @abc.abstractmethod
+    def open(cls, connection: Connection) -> Connection:
+        """Open the given connection on the adapter and return it.
+
+        This may mutate the given connection (in particular, its state and its
+        handle).
+
+        This should be thread-safe, or hold the lock if necessary. The given
+        connection should not be in either in_use or available.
+        """
+        raise NotImplementedError("`open` is not implemented for this adapter!")
+
+    def release(self) -> None:
+        with self.lock:
+            conn = self.get_if_exists()
+            if conn is None:
+                return
+
+        try:
+            # always close the connection
+            self.close(conn)
+        except Exception:
+            # if rollback or close failed, remove our busted connection
+            self.clear_thread_connection()
+            raise
+
+    def cleanup_all(self) -> None:
+        with self.lock:
+            for connection in self.thread_connections.values():
+                if connection.state not in {"closed", "init"}:
+                    fire_event(ConnectionLeftOpen(conn_name=connection.name))
+                else:
+                    fire_event(ConnectionClosed(conn_name=connection.name))
+                self.close(connection)
+
+            # garbage collect these connections
+            self.thread_connections.clear()
+
+    @classmethod
+    def _close_handle(cls, connection: Connection) -> None:
+        """Perform the actual close operation."""
+        # On windows, sometimes connection handles don't have a close() attr.
+        if hasattr(connection.handle, "close"):
+            connection.handle.close()
+
+    @classmethod
+    def close(cls, connection: Connection) -> Connection:
+        # if the connection is in closed or init, there's nothing to do
+        if connection.state in {ConnectionState.CLOSED, ConnectionState.INIT}:
+            return connection
+
+        cls._close_handle(connection)
+        connection.state = ConnectionState.CLOSED
+
+        return connection
+
+    @abc.abstractmethod
+    def execute(self, compiled_code: str) -> Tuple[AdapterResponse, Any]:
+        """Execute the given Python code.
+
+        :param str compiled_code: The Python code to execute.
+        :return: A tuple of the run status and results
+                (type is left to be decided by the Adapter implementation for now).
+        :rtype: Tuple[AdapterResponse, Any]
+        """
+        raise NotImplementedError("`execute` is not implemented for this adapter!")

dbt/fal/adapters/python/impl.py

@@ -0,0 +1,291 @@
+import abc
+from contextlib import contextmanager
+import time
+from typing import (
+    Optional,
+    Type,
+    Dict,
+    Any,
+    Mapping,
+    Iterator,
+)
+
+from dbt.adapters.factory import get_adapter
+from dbt.exceptions import (
+    NotImplementedError,
+    DbtRuntimeError,
+)
+from dbt.contracts.graph.nodes import ResultNode
+from dbt.adapters.protocol import AdapterConfig, ConnectionManagerProtocol
+from dbt.events.functions import fire_event
+from dbt.events.types import (
+    CodeExecution,
+    CodeExecutionStatus,
+)
+from dbt.adapters.base.meta import AdapterMeta, available
+from dbt.contracts.connection import Credentials, Connection, AdapterResponse
+
+
+def log_code_execution(code_execution_function):
+    # decorator to log code and execution time
+    if code_execution_function.__name__ != "submit_python_job":
+        raise ValueError("this should be only used to log submit_python_job now")
+
+    def execution_with_log(*args):
+        self = args[0]
+        connection_name = self.connections.get_thread_connection().name
+        fire_event(CodeExecution(conn_name=connection_name, code_content=args[2]))
+        start_time = time.time()
+        response = code_execution_function(*args)
+        fire_event(
+            CodeExecutionStatus(
+                status=response._message, elapsed=round((time.time() - start_time), 2)
+            )
+        )
+        return response
+
+    return execution_with_log
+
+
+class PythonJobHelper:
+    def __init__(self, parsed_model: Dict, credential: Credentials) -> None:
+        raise NotImplementedError("PythonJobHelper is not implemented yet")
+
+    def submit(self, compiled_code: str) -> Any:
+        raise NotImplementedError(
+            "PythonJobHelper submit function is not implemented yet"
+        )
+
+
+class PythonAdapter(metaclass=AdapterMeta):
+    """The PythonAdapter provides an abstract base class for Python adapters.
+
+    Adapters must implement the following methods and macros. Some of the
+    methods can be safely overridden as a noop, where it makes sense. Those
+    methods are marked with a (passable) in their docstrings. Check docstrings
+    for type information, etc.
+
+    To implement a macro, implement "${adapter_type}__${macro_name}" in the
+    adapter's internal project.
+
+    To invoke a method in an adapter macro, call it on the 'adapter' Jinja
+    object using dot syntax.
+
+    To invoke a method in model code, add the @available decorator atop a method
+    declaration. Methods are invoked as macros.
+
+    Methods:
+        - is_cancelable
+
+        - python_submission_helpers
+        - default_python_submission_method
+        - generate_python_submission_response
+        or
+        - submit_python_job
+
+    Macros:
+
+    """
+
+    ConnectionManager: Type[ConnectionManagerProtocol]
+
+    # A set of clobber config fields accepted by this adapter
+    # for use in materializations
+    AdapterSpecificConfigs: Type[AdapterConfig] = AdapterConfig
+
+    def __init__(self, config):
+        self.config = config
+        self.connections = self.ConnectionManager(config)
+
+        # HACK: A Python adapter does not have _available_ all the attributes a DB adapter does.
+        # Since we use the DB adapter as the storage for the Python adapter, we must proxy to it
+        # all the unhandled calls.
+        #
+        # Another option is to write a PythonAdapter-specific DBWrapper (PythonWrapper?) that is
+        # aware of this case. This may be appealing because a _complete_ adapter (DB+Python) would
+        # then be more easily used to replace the Python part of any other adapter.
+
+        self._db_adapter = get_adapter(config)
+        self.Relation = self._db_adapter.Relation
+        self.Column = self._db_adapter.Column
+
+    def cache_added(self, *args, **kwargs):
+        return self._db_adapter.cache_added(*args, **kwargs)
+
+    @available.parse_none
+    def get_relation(self, *args, **kwargs):
+        return self._db_adapter.get_relation(*args, **kwargs)
+
+    @classmethod
+    def date_function(cls):
+        # HACK: to appease the ProviderContext
+        return """
+        import datetime
+        return datetime.datetime.now()
+        """
+
+    ###
+    # Methods that pass through to the connection manager
+    ###
+    def acquire_connection(self, name=None) -> Connection:
+        return self.connections.set_connection_name(name)
+
+    def release_connection(self) -> None:
+        self.connections.release()
+
+    def cleanup_connections(self) -> None:
+        self.connections.cleanup_all()
+
+    def nice_connection_name(self) -> str:
+        conn = self.connections.get_if_exists()
+        if conn is None or conn.name is None:
+            return "<None>"
+        return conn.name
+
+    @contextmanager
+    def connection_named(
+        self, name: str, node: Optional[ResultNode] = None
+    ) -> Iterator[None]:
+        try:
+            self.acquire_connection(name)
+            yield
+        finally:
+            self.release_connection()
+
+    @contextmanager
+    def connection_for(self, node: ResultNode) -> Iterator[None]:
+        with self.connection_named(node.unique_id, node):
+            yield
+
+    @classmethod
+    @abc.abstractmethod
+    def is_cancelable(cls) -> bool:
+        raise NotImplementedError(
+            "`is_cancelable` is not implemented for this adapter!"
+        )
+
+    ###
+    # Methods that should never be overridden
+    ###
+    @classmethod
+    def type(cls) -> str:
+        """Get the type of this adapter. Types must be class-unique and
+        consistent.
+
+        :return: The type name
+        :rtype: str
+        """
+        return cls.ConnectionManager.TYPE
+
+    ###
+    # ODBC FUNCTIONS -- these should not need to change for every adapter,
+    #                   although some adapters may override them
+    ###
+    def cancel_open_connections(self):
+        """Cancel all open connections."""
+        return self.connections.cancel_open()
+
+    def pre_model_hook(self, config: Mapping[str, Any]) -> Any:
+        """A hook for running some operation before the model materialization
+        runs. The hook can assume it has a connection available.
+
+        The only parameter is a configuration dictionary (the same one
+        available in the materialization context). It should be considered
+        read-only.
+
+        The pre-model hook may return anything as a context, which will be
+        passed to the post-model hook.
+        """
+        pass
+
+    def post_model_hook(self, config: Mapping[str, Any], context: Any) -> None:
+        """A hook for running some operation after the model materialization
+        runs. The hook can assume it has a connection available.
+
+        The first parameter is a configuration dictionary (the same one
+        available in the materialization context). It should be considered
+        read-only.
+
+        The second parameter is the value returned by pre_mdoel_hook.
+        """
+        pass
+
+    def get_compiler(self):
+        from dbt.compilation import Compiler
+
+        return Compiler(self.config)
+
+    ###
+    # Python
+    ###
+    @property
+    def python_submission_helpers(self) -> Dict[str, Type[PythonJobHelper]]:
+        raise NotImplementedError("python_submission_helpers is not specified")
+
+    @property
+    def default_python_submission_method(self) -> str:
+        raise NotImplementedError("default_python_submission_method is not specified")
+
+    @log_code_execution
+    def submit_python_job(
+        self, parsed_model: dict, compiled_code: str
+    ) -> AdapterResponse:
+        submission_method = parsed_model["config"].get(
+            "submission_method", self.default_python_submission_method
+        )
+        if submission_method not in self.python_submission_helpers:
+            raise NotImplementedError(
+                "Submission method {} is not supported for current adapter".format(
+                    submission_method
+                )
+            )
+        job_helper = self.python_submission_helpers[submission_method](
+            parsed_model, self.connections.credentials
+        )
+        submission_result = job_helper.submit(compiled_code)
+        # process submission result to generate adapter response
+        return self.generate_python_submission_response(submission_result)
+
+    def generate_python_submission_response(
+        self, submission_result: Any
+    ) -> AdapterResponse:
+        raise NotImplementedError(
+            "Your adapter need to implement generate_python_submission_response"
+        )
+
+    def valid_incremental_strategies(self):
+        """The set of standard builtin strategies which this adapter supports out-of-the-box.
+        Not used to validate custom strategies defined by end users.
+        """
+        return ["append"]
+
+    def builtin_incremental_strategies(self):
+        return ["append", "delete+insert", "merge", "insert_overwrite"]
+
+    @available.parse_none
+    def get_incremental_strategy_macro(self, model_context, strategy: str):
+        # Construct macro_name from strategy name
+        if strategy is None:
+            strategy = "default"
+
+        # validate strategies for this adapter
+        valid_strategies = self.valid_incremental_strategies()
+        valid_strategies.append("default")
+        builtin_strategies = self.builtin_incremental_strategies()
+        if strategy in builtin_strategies and strategy not in valid_strategies:
+            raise DbtRuntimeError(
+                f"The incremental strategy '{strategy}' is not valid for this adapter"
+            )
+
+        strategy = strategy.replace("+", "_")
+        macro_name = f"get_incremental_{strategy}_sql"
+        # The model_context should have MacroGenerator callable objects for all macros
+        if macro_name not in model_context:
+            raise DbtRuntimeError(
+                'dbt could not find an incremental strategy macro with the name "{}" in {}'.format(
+                    macro_name, self.config.project_name
+                )
+            )
+
+        # This returns a callable macro
+        return model_context[macro_name]

dbt/fal/adapters/teleport/__init__.py

@@ -0,0 +1,3 @@
+# these are all just exports, #noqa them so flake8 will be happy
+from .impl import TeleportAdapter  # noqa
+from .info import TeleportInfo, S3TeleportInfo, LocalTeleportInfo  # noqa