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

_fal_testing/utils.py

@@ -43,8 +43,8 @@ def get_environment_type():
     # To determine whether this is a fal-created environment or not
     # we'll check whether the executable that is running this script
     # is located under any of the designated fal environment directories.
-    from fal.packages.environments.virtual_env import _BASE_VENV_DIR
-    from fal.packages.environments.conda import _BASE_CONDA_DIR
+    from fal.dbt.packages.environments.virtual_env import _BASE_VENV_DIR
+    from fal.dbt.packages.environments.conda import _BASE_CONDA_DIR
 
     executable_path = Path(sys.executable)
     for environment_type, prefix in [

dbt/adapters/fal/__init__.py

@@ -0,0 +1,21 @@
+from dbt.adapters.base import AdapterPlugin
+
+from dbt.adapters.fal.connections import FalEncCredentials
+from dbt.include import fal
+
+# Avoid loading the plugin code for any import
+def __getattr__(name):
+    if name == "FalEncAdapter":
+        from dbt.adapters.fal.impl import FalEncAdapter
+
+        return FalEncAdapter
+    if name == "Plugin":
+        from dbt.adapters.fal.impl import FalEncAdapter
+
+        return AdapterPlugin(
+            adapter=FalEncAdapter,
+            credentials=FalEncCredentials,
+            include_path=fal.PACKAGE_PATH
+        )
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

dbt/adapters/fal/__version__.py

@@ -0,0 +1 @@
+version = '0.9.4'

dbt/adapters/fal/connections.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from dbt.adapters.fal_experimental.connections import FalCredentials
+
+
+@dataclass
+class FalEncCredentials(FalCredentials):
+    db_profile: str = ""
+
+    def _connection_keys(self):
+        return () + super()._connection_keys()
+
+    @property
+    def type(self):
+        return "fal"
+
+    @property
+    def unique_field(self):
+        return self.db_profile

dbt/adapters/fal/impl.py

@@ -0,0 +1,93 @@
+from typing import Optional
+
+from collections import defaultdict
+from contextlib import contextmanager
+from dbt.adapters.base.impl import BaseAdapter
+from dbt.adapters.base.relation import BaseRelation
+from dbt.adapters.protocol import AdapterConfig
+from dbt.adapters.factory import FACTORY
+
+# TODO: offer in `from isolate import is_agent`
+from isolate.connections.common import is_agent
+
+from .connections import FalEncCredentials
+from .wrappers import FalEncAdapterWrapper, FalCredentialsWrapper
+
+from .load_db_profile import load_profiles_info_1_5
+
+
+class FalConfigs(AdapterConfig):
+    fal_environment: Optional[str]
+    fal_machine: Optional[str]
+
+
+@contextmanager
+def _release_plugin_lock():
+    FACTORY.lock.release()
+    try:
+        yield
+    finally:
+        FACTORY.lock.acquire()
+
+
+DB_PROFILE = None
+DB_RELATION = BaseRelation
+OVERRIDE_PROPERTIES = {}
+
+# NOTE: Should this file run on isolate agents? Could we skip it entirely and build a FalEncAdapterWrapper directly?
+if not is_agent():
+    DB_PROFILE, OVERRIDE_PROPERTIES = load_profiles_info_1_5()
+    DB_RELATION = FACTORY.get_relation_class_by_name(DB_PROFILE.credentials.type)
+
+
+class FalEncAdapter(BaseAdapter):
+    Relation = DB_RELATION  # type: ignore
+
+    # TODO: how do we actually use this?
+    AdapterSpecificConfigs = FalConfigs
+
+    def __new__(cls, config):
+        # There are two different credentials types which can be passed to FalEncAdapter
+        # 1. FalEncCredentials
+        # 2. FalCredentialsWrapper
+        #
+        # For `FalEncCredentials`, we have to go through parsing the profiles.yml (so that we
+        # can obtain the real 'db' credentials). But for the other one, we can just use
+        # the bound credentials directly (e.g. in isolated mode where we don't actually
+        # have access to the profiles.yml file).
+
+        fal_credentials = config.credentials
+        if isinstance(fal_credentials, FalEncCredentials):
+            db_credentials = DB_PROFILE.credentials
+        else:
+            # Since profile construction (in the case above) already registers the
+            # adapter plugin for the db type, we need to also mimic that here.
+            assert isinstance(fal_credentials, FalCredentialsWrapper)
+            db_credentials = fal_credentials._db_creds
+            with _release_plugin_lock():
+                FACTORY.load_plugin(db_credentials.type)
+
+        # TODO: maybe we can do this better?
+        with _release_plugin_lock():
+            db_adapter_class = FACTORY.get_adapter_class_by_name(db_credentials.type)
+            original_plugin = FACTORY.get_plugin_by_name(fal_credentials.type)
+            original_plugin.dependencies = [db_credentials.type]
+
+        config.python_adapter_credentials = fal_credentials
+        config.sql_adapter_credentials = db_credentials
+
+        for key in OVERRIDE_PROPERTIES:
+            if OVERRIDE_PROPERTIES[key] is not None:
+                setattr(config, key, OVERRIDE_PROPERTIES[key])
+
+        with _release_plugin_lock():
+            # Temporary credentials for register
+            config.credentials = config.sql_adapter_credentials
+            FACTORY.register_adapter(config)
+            config.credentials = FalCredentialsWrapper(config.sql_adapter_credentials)
+
+        return FalEncAdapterWrapper(db_adapter_class, config)
+
+    @classmethod
+    def type(cls):
+        return "fal"

dbt/adapters/fal/load_db_profile.py

@@ -0,0 +1,80 @@
+from typing import Dict, Any, Tuple, Optional
+
+from dbt.flags import get_flags, Namespace
+from dbt.config.project import load_raw_project
+from dbt.config.profile import read_profile, Profile
+from dbt.config.renderer import ProfileRenderer
+
+
+# NOTE: inspired in https://github.com/dbt-labs/dbt-core/blob/828d723512fced809c63e369a82c7eb570a74986/core/dbt/config/runtime.py#L58-L74
+def find_profile_name(
+    profile_override: Optional[str],
+    project_root: str,
+    profile_renderer: ProfileRenderer,
+):
+    if profile_override is not None:
+        profile_name = profile_override
+    else:
+        raw_project = load_raw_project(project_root)
+        raw_profile_name = raw_project.get("profile")
+        profile_name = profile_renderer.render_value(raw_profile_name)
+
+    return profile_name
+
+
+# NOTE: inspired in https://github.com/dbt-labs/dbt-core/blob/828d723512fced809c63e369a82c7eb570a74986/core/dbt/config/profile.py#L279-L311
+def find_target_name(
+    target_override: Optional[str], raw_profile: dict, profile_renderer: ProfileRenderer
+):
+    if target_override is not None:
+        target_name = target_override
+    elif "target" in raw_profile:
+        # render the target if it was parsed from yaml
+        target_name = profile_renderer.render_value(raw_profile["target"])
+    else:
+        target_name = "default"
+
+    return target_name
+
+
+def load_profiles_info_1_5() -> Tuple[Profile, Dict[str, Any]]:
+    flags: Namespace = get_flags()
+
+    profile_renderer = ProfileRenderer(getattr(flags, "VARS", {}))
+
+    profile_name = find_profile_name(flags.PROFILE, flags.PROJECT_DIR, profile_renderer)
+
+    raw_profiles = read_profile(flags.PROFILES_DIR)
+    raw_profile = raw_profiles[profile_name]
+
+    target_name = find_target_name(flags.TARGET, raw_profile, profile_renderer)
+
+    fal_dict = Profile._get_profile_data(
+        profile=raw_profile,
+        profile_name=profile_name,
+        target_name=target_name,
+    )
+    db_profile_target_name = fal_dict.get("db_profile")
+    assert (
+        db_profile_target_name
+    ), "fal credentials must have a `db_profile` property set"
+
+    try:
+        db_profile = Profile.from_raw_profile_info(
+            raw_profile=raw_profile,
+            profile_name=profile_name,
+            renderer=profile_renderer,
+            # TODO: should we load the user_config?
+            user_config={},
+            target_override=db_profile_target_name,
+        )
+    except RecursionError as error:
+        raise AttributeError(
+            "Did you wrap a type 'fal' profile with another type 'fal' profile?"
+        ) from error
+
+    override_properties = {
+        "threads": getattr(flags, "THREADS", None) or fal_dict.get("threads") or db_profile.threads,
+    }
+
+    return db_profile, override_properties

dbt/adapters/fal/wrappers.py

@@ -0,0 +1,113 @@
+from typing import Any, Optional, Type, Set
+
+from dbt.adapters.factory import get_adapter_by_type
+from dbt.adapters.base.meta import available
+from dbt.adapters.base.impl import BaseAdapter
+from dbt.contracts.connection import Credentials
+from dbt.parser.manifest import ManifestLoader
+from dbt.clients.jinja import MacroGenerator
+
+from ..fal_experimental.utils import cache_static
+from ..fal_experimental.impl import FalAdapterMixin
+from ..fal_experimental import telemetry
+
+
+class FalCredentialsWrapper:
+    _db_creds: Optional[Credentials] = None
+
+    def __init__(self, db_creds: Credentials):
+        self._db_creds = db_creds
+
+    @property
+    def type(self):
+        if find_funcs_in_stack({"to_target_dict", "db_materialization"}):
+            # This makes sense for both SQL and Python because the target is always the db
+            return self._db_creds.type
+
+        return "fal"
+
+    def __getattr__(self, name: str) -> Any:
+        """
+        Directly proxy to the DB adapter, just shadowing the type
+        """
+        return getattr(self._db_creds, name)
+
+
+class FalEncAdapterWrapper(FalAdapterMixin):
+    def __init__(self, db_adapter_type: Type[BaseAdapter], config):
+        # Use the db_adapter_type connection manager
+        self.ConnectionManager = db_adapter_type.ConnectionManager
+
+        db_adapter = get_adapter_by_type(db_adapter_type.type())
+        super().__init__(config, db_adapter)
+
+        # HACK: A Python adapter does not have self._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.
+
+        # self._available_ is set by metaclass=AdapterMeta
+        self._available_ = self._db_adapter._available_.union(self._available_)
+        # self._parse_replacements_ is set by metaclass=AdapterMeta
+        self._parse_replacements_.update(self._db_adapter._parse_replacements_)
+
+        telemetry.log_api(
+            "encapsulate_init",
+            config=config,
+            additional_props={
+                "is_teleport": self.is_teleport(),
+                "fal_host": self.credentials.host,
+            },
+        )
+
+    def submit_python_job(self, *args, **kwargs):
+        return super().submit_python_job(*args, **kwargs)
+
+    @available
+    @telemetry.log_call(
+        "encapsulate_db_materialization", log_args=["materialization"], config=True
+    )
+    def db_materialization(self, context: dict, materialization: str):
+        # NOTE: inspired by https://github.com/dbt-labs/dbt-core/blob/be4a91a0fe35a619587b7a0145e190690e3771c6/core/dbt/task/run.py#L254-L290
+        materialization_macro = self.manifest.find_materialization_macro_by_name(
+            self.config.project_name, materialization, self._db_adapter.type()
+        )
+
+        # HACK: run the entire SQL materialization and return the resulting dict with relations created
+        return MacroGenerator(
+            materialization_macro, context, stack=context["context_macro_stack"]
+        )()
+
+    @property
+    @cache_static
+    def manifest(self):
+        return ManifestLoader.get_full_manifest(self.config)
+
+    def type(self):
+        # NOTE: This does not let `fal__` macros to be used
+        # Maybe for 1.5 we will get a more reliable way to detect if we are in a SQL or Python context
+        if find_funcs_in_stack({"render", "db_materialization"}):
+            return self._db_adapter.type()
+
+        return "fal"
+
+    def __getattr__(self, name):
+        """
+        Directly proxy to the DB adapter, Python adapter in this case does what we explicitly define in this class.
+        """
+        if hasattr(self._db_adapter, name):
+            return getattr(self._db_adapter, name)
+        else:
+            getattr(super(), name)
+
+
+def find_funcs_in_stack(funcs: Set[str]) -> bool:
+    import inspect
+
+    # NOTE: from https://stackoverflow.com/a/42636264/1276441
+    frame = inspect.currentframe()
+    while frame:
+        if frame.f_code.co_name in funcs:
+            return True
+        frame = frame.f_back
+
+    return False

dbt/adapters/fal_experimental/__init__.py

@@ -0,0 +1,11 @@
+from dbt.adapters.base import AdapterPlugin
+
+from dbt.adapters.fal_experimental.connections import FalCredentials
+from dbt.adapters.fal_experimental.impl import FalAdapter
+from dbt.include import fal_experimental
+
+Plugin = AdapterPlugin(
+    adapter=FalAdapter,
+    credentials=FalCredentials,
+    include_path=fal_experimental.PACKAGE_PATH,
+)

dbt/adapters/fal_experimental/__version__.py

@@ -0,0 +1 @@
+version = '0.9.4'

dbt/adapters/fal_experimental/adapter.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import zipfile
+import io
+from functools import partial
+from tempfile import NamedTemporaryFile
+from typing import Any, Optional
+
+from dbt.adapters.base.impl import BaseAdapter
+from dbt.config.runtime import RuntimeConfig
+from dbt.contracts.connection import AdapterResponse
+from dbt.flags import get_flags, Namespace
+
+from fal_serverless import FalServerlessHost, isolated
+from dbt.adapters.fal_experimental.utils.environments import (
+    EnvironmentDefinition,
+    get_default_pip_dependencies
+)
+
+from dbt.parser.manifest import MacroManifest, Manifest
+
+from .adapter_support import (
+    prepare_for_adapter,
+    read_relation_as_df,
+    reconstruct_adapter,
+    write_df_to_relation,
+)
+
+from .utils import extra_path, get_fal_scripts_path, retrieve_symbol
+
+
+def run_with_adapter(code: str, adapter: BaseAdapter, config: RuntimeConfig) -> Any:
+    # main symbol is defined during dbt-fal's compilation
+    # and acts as an entrypoint for us to run the model.
+    fal_scripts_path = str(get_fal_scripts_path(config))
+    with extra_path(fal_scripts_path):
+        main = retrieve_symbol(code, "main")
+        return main(
+            read_df=prepare_for_adapter(adapter, read_relation_as_df),
+            write_df=prepare_for_adapter(adapter, write_df_to_relation),
+        )
+
+
+def _isolated_runner(
+    code: str,
+    flags: Namespace,
+    config: RuntimeConfig,
+    manifest: Manifest,
+    macro_manifest: MacroManifest,
+    local_packages: Optional[bytes] = None,
+) -> Any:
+    # This function can be run in an entirely separate
+    # process or an environment, so we need to reconstruct
+    # the DB adapter solely from the config.
+    adapter = reconstruct_adapter(flags, config, manifest, macro_manifest)
+    fal_scripts_path = get_fal_scripts_path(config)
+    if local_packages is not None:
+        if fal_scripts_path.exists():
+            import shutil
+            shutil.rmtree(fal_scripts_path)
+        fal_scripts_path.parent.mkdir(parents=True, exist_ok=True)
+        zip_file = zipfile.ZipFile(io.BytesIO(local_packages))
+        zip_file.extractall(fal_scripts_path)
+
+    return run_with_adapter(code, adapter, config)
+
+
+def run_in_environment_with_adapter(
+    environment: EnvironmentDefinition,
+    code: str,
+    config: RuntimeConfig,
+    manifest: Manifest,
+    macro_manifest: MacroManifest,
+    adapter_type: str
+) -> AdapterResponse:
+    """Run the 'main' function inside the given code on the
+    specified environment.
+
+    The environment_name must be defined inside fal_project.yml file
+    in your project's root directory."""
+    compressed_local_packages = None
+
+    is_remote = type(environment.host) is FalServerlessHost
+
+    deps = get_default_pip_dependencies(
+        is_remote=is_remote,
+        adapter_type=adapter_type)
+
+    fal_scripts_path = get_fal_scripts_path(config)
+
+    if is_remote and fal_scripts_path.exists():
+        with NamedTemporaryFile() as temp_file:
+            with zipfile.ZipFile(
+                temp_file.name, "w", zipfile.ZIP_DEFLATED
+            ) as zip_file:
+                for entry in fal_scripts_path.rglob("*"):
+                    zip_file.write(entry, entry.relative_to(fal_scripts_path))
+
+            compressed_local_packages = temp_file.read()
+
+    execute_model = partial(
+        _isolated_runner,
+        code=code,
+        flags=get_flags(),
+        config=config,
+        manifest=manifest,
+        macro_manifest=macro_manifest,
+        local_packages=compressed_local_packages
+    )
+
+    if environment.kind == "local":
+        result = execute_model()
+        return result
+
+    extra = {}
+    if is_remote:
+        # Machine type is only applicable in FalServerlesshost
+        extra = {"machine_type": environment.machine_type}
+
+    if environment.kind == "virtualenv":
+        requirements = environment.config.get("requirements", [])
+        requirements += deps
+        isolated_function = isolated(
+            kind="virtualenv",
+            host=environment.host,
+            requirements=requirements,
+            **extra
+        )(execute_model)
+    elif environment.kind == "conda":
+        dependencies = environment.config.pop("packages", [])
+        dependencies.append({"pip": deps})
+        env_dict = {
+            "name": "dbt_fal_env",
+            "channels": ["conda-forge", "defaults"],
+            "dependencies": dependencies
+        }
+        isolated_function = isolated(
+            kind="conda",
+            host=environment.host,
+            env_dict=env_dict,
+            **extra
+        )(execute_model)
+    else:
+        # We should not reach this point, because environment types are validated when the
+        # environment objects are created (in utils/environments.py).
+        raise Exception(f"Environment type not supported: {environment.kind}")
+
+    result = isolated_function()
+    return result