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

dbt/adapters/fal_experimental/teleport_support/duckdb.py

@@ -0,0 +1,122 @@
+from contextlib import contextmanager
+from typing import Optional
+
+from dbt.exceptions import DbtRuntimeError
+from dbt.adapters.base.relation import BaseRelation
+from dbt.adapters.base.impl import BaseAdapter
+
+from dbt.adapters.fal_experimental.connections import (
+    TeleportCredentials,
+    TeleportTypeEnum,
+)
+
+from dbt.fal.adapters.teleport.impl import TeleportAdapter
+from dbt.fal.adapters.teleport.info import TeleportInfo
+
+
+class DuckDBAdapterTeleport(TeleportAdapter):
+    def __init__(
+        self, db_adapter: BaseAdapter, teleport_credentials: TeleportCredentials
+    ):
+        self._db_adapter = db_adapter
+        self.credentials = teleport_credentials
+        with self._db_adapter.connection_named("teleport:init"):
+            self._db_adapter.execute("INSTALL parquet")
+            self._db_adapter.execute("INSTALL httpfs")
+
+    @classmethod
+    def storage_formats(cls):
+        return ["parquet"]
+
+    def teleport_from_external_storage(
+        self, relation: BaseRelation, relation_path: str, teleport_info: TeleportInfo
+    ):
+        assert (
+            teleport_info.format == "parquet"
+        ), "duckdb only supports parquet format for Teleport"
+
+        url = teleport_info.build_url(relation_path)
+
+        with self._db_adapter.connection_named("teleport:copy_from"):
+            if self.credentials.type == TeleportTypeEnum.LOCAL:
+                rendered_macro = self._db_adapter.execute_macro(
+                    "duckdb__copy_from_parquet",
+                    kwargs={"relation": relation, "url": url},
+                )
+                self._db_adapter.execute(rendered_macro)
+
+            elif self.credentials.type == TeleportTypeEnum.REMOTE_S3:
+                with self._s3_setup():
+                    rendered_macro = self._db_adapter.execute_macro(
+                        "duckdb__copy_from_parquet",
+                        kwargs={"relation": relation, "url": url},
+                    )
+                    self._db_adapter.execute(rendered_macro)
+            else:
+                raise RuntimeError(
+                    f"Teleport type {self.credentials.type} not supported"
+                )
+
+    def teleport_to_external_storage(
+        self, relation: BaseRelation, teleport_info: TeleportInfo
+    ):
+        assert (
+            teleport_info.format == "parquet"
+        ), "duckdb only supports parquet format for Teleport"
+
+        rel_path = teleport_info.build_relation_path(relation)
+        url = teleport_info.build_url(rel_path)
+
+        with self._db_adapter.connection_named("teleport:copy_to"):
+            if self.credentials.type == TeleportTypeEnum.LOCAL:
+                rendered_macro = self._db_adapter.execute_macro(
+                    "duckdb__copy_to", kwargs={"relation": relation, "url": url}
+                )
+                self._db_adapter.execute(rendered_macro)
+            elif self.credentials.type == TeleportTypeEnum.REMOTE_S3:
+                with self._s3_setup():
+                    rendered_macro = self._db_adapter.execute_macro(
+                        "duckdb__copy_to", kwargs={"relation": relation, "url": url}
+                    )
+                    self._db_adapter.execute(rendered_macro)
+            else:
+                raise RuntimeError(
+                    f"Teleport type {self.credentials.type} not supported"
+                )
+
+        return rel_path
+
+    def _get_setting(self, name: str):
+        try:
+            _, table = self._db_adapter.execute(
+                f"SELECT current_setting('{name}')", fetch=True
+            )
+            return table.rows[0][0]
+        except DbtRuntimeError:
+            return None
+
+    def _set_setting(self, name: str, value: Optional[str]):
+        if value:
+            self._db_adapter.execute(f"SET {name} = '{value}'")
+        else:
+            # HACK while we get a response https://github.com/duckdb/duckdb/issues/4998
+            self._db_adapter.execute(f"SET {name} = ''")
+
+    @contextmanager
+    def _s3_setup(self):
+        self._db_adapter.execute("LOAD parquet")
+        self._db_adapter.execute("LOAD httpfs")
+
+        old_region = self._get_setting("s3_region")
+        old_access_key_id = self._get_setting("s3_access_key_id")
+        old_secret_access_key = self._get_setting("s3_secret_access_key")
+
+        self._set_setting("s3_region", self.credentials.s3_region)
+        self._set_setting("s3_access_key_id", self.credentials.s3_access_key_id)
+        self._set_setting("s3_secret_access_key", self.credentials.s3_access_key)
+
+        yield
+
+        self._set_setting("s3_region", old_region)
+        self._set_setting("s3_access_key_id", old_access_key_id)
+        self._set_setting("s3_secret_access_key", old_secret_access_key)

dbt/adapters/fal_experimental/teleport_support/snowflake.py

@@ -0,0 +1,72 @@
+from dbt.adapters.fal_experimental.connections import TeleportCredentials
+from dbt.fal.adapters.teleport.impl import TeleportAdapter
+from dbt.fal.adapters.teleport.info import TeleportInfo
+from dbt.adapters.base.relation import BaseRelation
+from dbt.adapters.base.impl import BaseAdapter
+
+
+class SnowflakeAdapterTeleport(TeleportAdapter):
+
+    def __init__(self, db_adapter: BaseAdapter, teleport_credentials: TeleportCredentials):
+        from dbt.adapters.fal_experimental.adapter_support import new_connection
+        self._db_adapter = db_adapter
+        self._credentials = teleport_credentials
+
+        # TODO: put this in teleport_info
+        url = f's3://{teleport_credentials.s3_bucket}/teleport'
+
+        with new_connection(self._db_adapter, "fal-snowflake:setup-teleport") as conn:
+            cur = conn.handle.cursor()
+
+            create_stage_query = f"""CREATE OR REPLACE STAGE falstage
+            URL = '{url}' CREDENTIALS = (
+            aws_key_id='{self._credentials.s3_access_key_id}',
+            aws_secret_key='{self._credentials.s3_access_key}');"""
+
+            create_format_query = """CREATE OR REPLACE FILE FORMAT falparquet type = 'PARQUET';"""
+
+            cur.execute(create_stage_query)
+            cur.execute(create_format_query)
+
+
+    @classmethod
+    def storage_formats(cls):
+        return ['parquet']
+
+    def teleport_from_external_storage(self, relation: BaseRelation, relation_path: str, teleport_info: TeleportInfo) -> None:
+        assert teleport_info.format == 'parquet', "snowflake only supports parquet format for Teleport"
+        location = f"@falstage/{relation_path}"
+
+        with self._db_adapter.connection_named('teleport:copy_from'):
+            create_macro = self._db_adapter.execute_macro(
+                'snowflake__create_table_from_parquet',
+                kwargs={
+                    'relation': relation,
+                    'location': location,
+                })
+
+            self._db_adapter.execute(create_macro)
+
+            copy_macro = self._db_adapter.execute_macro(
+                'snowflake__copy_from_parquet',
+                kwargs={
+                    'relation': relation,
+                    'location': location,
+                }
+            )
+            self._db_adapter.execute(copy_macro)
+
+    def teleport_to_external_storage(self, relation: BaseRelation, teleport_info: TeleportInfo) -> str:
+        assert teleport_info.format == 'parquet', "snowflake only supports parquet format for Teleport"
+        relation_path = teleport_info.build_relation_path(relation)
+        location = f"@falstage/{relation_path}"
+        rendered_macro = self._db_adapter.execute_macro(
+            'snowflake__copy_to_parquet',
+            kwargs={
+                'relation': relation,
+                'location': location
+            })
+        with self._db_adapter.connection_named('teleport:copy_to'):
+            self._db_adapter.execute(rendered_macro)
+        return relation_path
+

dbt/adapters/fal_experimental/utils/__init__.py

@@ -0,0 +1,50 @@
+from contextlib import contextmanager
+from typing import Any
+
+from dbt.config.runtime import RuntimeConfig
+
+try:
+    from functools import lru_cache
+except ImportError:
+    from backports.functools_lru_cache import lru_cache
+
+
+FAL_SCRIPTS_PATH_VAR_NAME = 'fal-scripts-path'
+
+
+def cache_static(func):
+    """Cache the result of a function."""
+    return lru_cache(maxsize=None)(func)
+
+
+def retrieve_symbol(source_code: str, symbol_name: str) -> Any:
+    """Retrieve the function with the given name from the source code."""
+    namespace = {}
+    exec(source_code, namespace)
+    return namespace[symbol_name]
+
+
+def get_fal_scripts_path(config: RuntimeConfig):
+    import pathlib
+    project_path = pathlib.Path(config.project_root)
+
+    # Default value
+    fal_scripts_path = ''
+
+    if hasattr(config, 'vars'):
+        fal_scripts_path: str = config.vars.to_dict().get(FAL_SCRIPTS_PATH_VAR_NAME, fal_scripts_path)  # type: ignore
+
+    if hasattr(config, 'cli_vars'):
+        fal_scripts_path = config.cli_vars.get(FAL_SCRIPTS_PATH_VAR_NAME, fal_scripts_path)
+
+    return project_path / fal_scripts_path
+
+
+@contextmanager
+def extra_path(path: str):
+    import sys
+    sys.path.append(path)
+    try:
+        yield
+    finally:
+        sys.path.remove(path)

dbt/adapters/fal_experimental/utils/environments.py

@@ -0,0 +1,302 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from pathlib import Path
+from typing import Any, ClassVar, Dict, Iterator, List, Optional, Tuple
+import importlib_metadata
+
+from dbt.events.adapter_endpoint import AdapterLogger
+from dbt.exceptions import DbtRuntimeError
+from dbt.config.runtime import RuntimeConfig
+
+from isolate.backends import BaseEnvironment, BasicCallable, EnvironmentConnection
+from fal_serverless import FalServerlessKeyCredentials, LocalHost
+from fal_serverless.api import Host, FalServerlessHost
+
+from . import cache_static
+from .yaml_helper import load_yaml
+
+
+CONFIG_KEYS_TO_IGNORE = ["host", "remote_type", "type", "name", "machine_type"]
+REMOTE_TYPES_DICT = {
+    "venv": "virtualenv",
+    "conda": "conda",
+}
+
+logger = AdapterLogger("fal")
+
+
+class FalParseError(Exception):
+    pass
+
+
+@dataclass
+class LocalConnection(EnvironmentConnection):
+    def run(self, executable: BasicCallable, *args, **kwargs) -> Any:
+        return executable(*args, **kwargs)
+
+
+@dataclass
+class EnvironmentDefinition():
+    host: Host
+    kind: str
+    config: dict[Any, Any]
+    machine_type: str = "S"
+
+
+def fetch_environment(
+    project_root: str,
+    environment_name: str,
+    machine_type: str = "S",
+    credentials: Optional[Any] = None,
+) -> Tuple[EnvironmentDefinition, bool]:
+    """Fetch the environment with the given name from the project's
+    fal_project.yml file."""
+    # Local is a special environment where it doesn't need to be defined
+    # since it will mirror user's execution context directly.
+    if environment_name == "local":
+        if credentials.host:
+            logger.warning(
+                "`local` environments will be executed on fal cloud."
+                + "If you don't want to use fal cloud, you can change your "
+                + "profile target to one where fal doesn't have credentials"
+            )
+            host = FalServerlessHost(
+                url=credentials.host,
+                credentials=FalServerlessKeyCredentials(credentials.key_id, credentials.key_secret))
+            return EnvironmentDefinition(
+                host=host,
+                kind="virtualenv",
+                machine_type=machine_type,
+                config={"name": "", "type": "venv"}), False
+
+        return EnvironmentDefinition(host=LocalHost(), kind="local", config={}), True
+
+    try:
+        environments = load_environments(project_root, machine_type, credentials)
+    except Exception as exc:
+        raise DbtRuntimeError(
+            "Error loading environments from fal_project.yml"
+        ) from exc
+
+    if environment_name not in environments:
+        raise DbtRuntimeError(
+            f"Environment '{environment_name}' was used but not defined in fal_project.yml"
+        )
+
+    return environments[environment_name], False
+
+
+def db_adapter_config(config: RuntimeConfig) -> RuntimeConfig:
+    """Return a config object that has the database adapter as its primary. Only
+    applicable when the underlying db adapter is encapsulated."""
+    if hasattr(config, "sql_adapter_credentials"):
+        new_config = replace(config, credentials=config.sql_adapter_credentials)
+        new_config.python_adapter_credentials = config.credentials
+    else:
+        new_config = config
+
+    return new_config
+
+
+def load_environments(
+    base_dir: str, machine_type: str = "S", credentials: Optional[Any] = None
+) -> Dict[str, EnvironmentDefinition]:
+    import os
+    fal_project_path = os.path.join(base_dir, "fal_project.yml")
+    if not os.path.exists(fal_project_path):
+        raise FalParseError(f"{fal_project_path} must exist to define environments")
+
+    fal_project = load_yaml(fal_project_path)
+    environments = {}
+    for environment in fal_project.get("environments", []):
+        env_name = _get_required_key(environment, "name")
+        if _is_local_environment(env_name):
+            raise FalParseError(
+                f"Environment name conflicts with a reserved name: {env_name}."
+            )
+
+        env_kind = _get_required_key(environment, "type")
+        if environments.get(env_name) is not None:
+            raise FalParseError("Environment names must be unique.")
+
+        environments[env_name] = create_environment(
+            env_name, env_kind, environment, machine_type, credentials
+        )
+
+    return environments
+
+
+def create_environment(
+    name: str,
+    kind: str,
+    config: Dict[str, Any],
+    machine_type: str = "S",
+    credentials: Optional[Any] = None,
+) -> EnvironmentDefinition:
+    if kind not in ["venv", "conda"]:
+        raise ValueError(
+            f"Invalid environment type (of {kind}) for {name}. Please choose from: "
+            + "venv, conda."
+        )
+
+    kind = kind if kind == "conda" else "virtualenv"
+
+    parsed_config = {
+        key: val for key, val in config.items() if key not in CONFIG_KEYS_TO_IGNORE
+    }
+
+    if credentials.key_secret and credentials.key_id:
+        host = FalServerlessHost(
+            url=credentials.host,
+            credentials=FalServerlessKeyCredentials(credentials.key_id, credentials.key_secret))
+    else:
+        host = LocalHost()
+    return EnvironmentDefinition(
+        host=host,
+        kind=kind,
+        config=parsed_config,
+        machine_type=machine_type)
+
+
+def _is_local_environment(environment_name: str) -> bool:
+    return environment_name == "local"
+
+
+def _get_required_key(data: Dict[str, Any], name: str) -> Any:
+    if name not in data:
+        raise FalParseError("Missing required key: " + name)
+    return data[name]
+
+def _parse_remote_config(
+    config: Dict[str, Any], parsed_config: Dict[str, Any]
+) -> Dict[str, Any]:
+    assert config.get("remote_type"), "remote_type needs to be specified."
+
+    remote_type = REMOTE_TYPES_DICT.get(config["remote_type"])
+
+    assert (
+        remote_type
+    ), f"{config['remote_type']} not recognised. Available remote types: {list(REMOTE_TYPES_DICT.keys())}"
+
+    env_definition = {
+        "kind": remote_type,
+        "configuration": parsed_config,
+    }
+
+    return {
+        "host": config.get("host"),
+        "target_environments": [env_definition],
+    }
+
+def _get_package_from_type(adapter_type: str):
+    SPECIAL_ADAPTERS = {
+        # Documented in https://docs.getdbt.com/docs/supported-data-platforms#community-adapters
+        "athena": "dbt-athena-community",
+    }
+    return SPECIAL_ADAPTERS.get(adapter_type, f"dbt-{adapter_type}")
+
+def _get_dbt_packages(
+    adapter_type: str,
+    is_teleport: bool = False,
+    is_remote: bool = False,
+) -> Iterator[Tuple[str, Optional[str]]]:
+    dbt_adapter = _get_package_from_type(adapter_type)
+    for dbt_plugin_name in [dbt_adapter]:
+        distribution = importlib_metadata.distribution(dbt_plugin_name)
+
+        yield dbt_plugin_name, distribution.version
+
+    try:
+        dbt_fal_version = importlib_metadata.version("dbt-fal")
+    except importlib_metadata.PackageNotFoundError:
+        # It might not be installed.
+        return None
+
+    dbt_fal_dep = "dbt-fal"
+    dbt_fal_extras = _find_adapter_extras(dbt_fal_dep, dbt_adapter)
+    if is_teleport:
+        dbt_fal_extras.add("teleport")
+    dbt_fal_suffix = ""
+
+    if _version_is_prerelease(dbt_fal_version):
+        if is_remote:
+            # If it's a pre-release and it's remote, its likely us developing, so we try installing
+            # from Github and we can get the custom branch name from FAL_GITHUB_BRANCH environment variable
+            # TODO: Handle pre-release on PyPI. How should we approach that?
+            import os
+
+            branch_name = os.environ.get("FAL_GITHUB_BRANCH", "main")
+
+            dbt_fal_suffix = f" @ git+https://github.com/fal-ai/fal.git@{branch_name}#subdirectory=projects/adapter"
+            dbt_fal_version = None
+        else:
+            dbt_fal_path = _get_project_root_path("adapter")
+            if dbt_fal_path is not None:
+                # Can be a pre-release from PyPI
+                dbt_fal_dep = str(dbt_fal_path)
+                dbt_fal_version = None
+
+    dbt_fal = f"{dbt_fal_dep}[{' ,'.join(dbt_fal_extras)}]{dbt_fal_suffix}"
+    yield dbt_fal, dbt_fal_version
+
+
+def _find_adapter_extras(package: str, plugin_package: str) -> set[str]:
+    import pkgutil
+    import dbt.adapters
+
+    all_extras = _get_extras(package)
+    available_plugins = {
+        module_info.name
+        for module_info in pkgutil.iter_modules(dbt.adapters.__path__)
+        if module_info.ispkg and module_info.name in plugin_package
+    }
+    return available_plugins.intersection(all_extras)
+
+
+def _get_extras(package: str) -> list[str]:
+    import importlib_metadata
+
+    dist = importlib_metadata.distribution(package)
+    return dist.metadata.get_all("Provides-Extra", [])
+
+
+def _version_is_prerelease(raw_version: str) -> bool:
+    from packaging.version import Version
+
+    package_version = Version(raw_version)
+    return package_version.is_prerelease
+
+
+def _get_project_root_path(package: str) -> Path:
+    from dbt.adapters import fal
+
+    # If this is a development version, we'll install
+    # the current fal itself.
+    path = Path(fal.__file__)
+    while path is not None:
+        if (path.parent / ".git").exists():
+            break
+        path = path.parent
+    return path / package
+
+
+def get_default_requirements(
+    adapter_type: str,
+    is_teleport: bool = False,
+    is_remote: bool = False,
+) -> Iterator[Tuple[str, Optional[str]]]:
+    yield from _get_dbt_packages(adapter_type, is_teleport, is_remote)
+
+@cache_static
+def get_default_pip_dependencies(
+    adapter_type: str,
+    is_teleport: bool = False,
+    is_remote: bool = False,
+) -> List[str]:
+    return [
+        f"{package}=={version}" if version else package
+        for package, version in get_default_requirements(
+            adapter_type, is_teleport, is_remote
+        )
+    ]

dbt/fal/adapters/python/__init__.py

@@ -0,0 +1,3 @@
+# these are all just exports, #noqa them so flake8 will be happy
+from dbt.fal.adapters.python.connections import PythonConnectionManager  # noqa
+from dbt.fal.adapters.python.impl import PythonAdapter  # noqa