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

dbt/adapters/fal_experimental/adapter_support.py

@@ -0,0 +1,234 @@
+import functools
+from typing import Any
+
+import pandas as pd
+import sqlalchemy
+from contextlib import contextmanager
+from dbt.adapters.base import BaseAdapter, BaseRelation, RelationType
+from dbt.adapters.base.connections import AdapterResponse, Connection
+from dbt.config import RuntimeConfig
+from dbt.parser.manifest import MacroManifest, Manifest
+from dbt.flags import Namespace
+
+from dbt.adapters import factory
+
+_SQLALCHEMY_DIALECTS = {
+    "sqlserver": "mssql+pyodbc",
+}
+
+
+def _get_alchemy_engine(adapter: BaseAdapter, connection: Connection) -> Any:
+    # The following code heavily depends on the implementation
+    # details of the known adapters, hence it can't work for
+    # arbitrary ones.
+    adapter_type = adapter.type()
+
+    sqlalchemy_kwargs = {}
+    format_url = lambda url: url
+
+    if adapter_type == "trino":
+        import dbt.adapters.fal_experimental.support.trino as support_trino
+
+        return support_trino.create_engine(adapter)
+
+    elif adapter_type == "sqlserver":
+        sqlalchemy_kwargs["creator"] = lambda *args, **kwargs: connection.handle
+        url = _SQLALCHEMY_DIALECTS.get(adapter_type, adapter_type) + "://"
+        url = format_url(url)
+    else:
+        message = (
+            f"dbt-fal does not support {adapter_type} adapter. ",
+            f"If you need {adapter_type} support, you can create an issue ",
+            "in our GitHub repository: https://github.com/fal-ai/fal. ",
+            "We will look into it ASAP.",
+        )
+        raise NotImplementedError(message)
+
+    return sqlalchemy.create_engine(url, **sqlalchemy_kwargs)
+
+
+def drop_relation_if_it_exists(adapter: BaseAdapter, relation: BaseRelation) -> None:
+    if adapter.get_relation(
+        database=relation.database,
+        schema=relation.schema,
+        identifier=relation.identifier,
+    ):
+        adapter.drop_relation(relation)
+
+
+def write_df_to_relation(
+    adapter: BaseAdapter,
+    relation: BaseRelation,
+    dataframe: pd.DataFrame,
+    *,
+    if_exists: str = "replace",
+) -> AdapterResponse:
+    """Generic version of the write_df_to_relation. Materialize the given
+    dataframe to the targeted relation on the adapter."""
+
+    adapter_type = adapter.type()
+
+    if adapter_type == "snowflake":
+        import dbt.adapters.fal_experimental.support.snowflake as support_snowflake
+
+        return support_snowflake.write_df_to_relation(adapter, dataframe, relation)
+
+    elif adapter_type == "bigquery":
+        import dbt.adapters.fal_experimental.support.bigquery as support_bq
+
+        return support_bq.write_df_to_relation(adapter, dataframe, relation)
+
+    elif adapter_type == "duckdb":
+        import dbt.adapters.fal_experimental.support.duckdb as support_duckdb
+
+        return support_duckdb.write_df_to_relation(adapter, dataframe, relation)
+
+    elif adapter_type == "postgres":
+        import dbt.adapters.fal_experimental.support.postgres as support_postgres
+
+        return support_postgres.write_df_to_relation(adapter, dataframe, relation)
+    elif adapter.type() == "athena":
+        import dbt.adapters.fal_experimental.support.athena as support_athena
+
+        return support_athena.write_df_to_relation(
+            adapter, dataframe, relation, if_exists
+        )
+    elif adapter_type == "redshift":
+        import dbt.adapters.fal_experimental.support.redshift as support_redshift
+
+        return support_redshift.write_df_to_relation(adapter, dataframe, relation)
+
+    else:
+        with new_connection(adapter, "fal:write_df_to_relation") as connection:
+            # TODO: this should probably live in the materialization macro.
+            temp_relation = relation.replace_path(
+                identifier=f"__dbt_fal_temp_{relation.identifier}"
+            )
+
+            drop_relation_if_it_exists(adapter, temp_relation)
+
+            alchemy_engine = _get_alchemy_engine(adapter, connection)
+
+            # TODO: probably worth handling errors here an returning
+            # a proper adapter response.
+            rows_affected = dataframe.to_sql(
+                con=alchemy_engine,
+                name=temp_relation.identifier,
+                schema=temp_relation.schema,
+                if_exists=if_exists,
+                index=False,
+            )
+            adapter.cache.add(temp_relation)
+            drop_relation_if_it_exists(adapter, relation)
+
+            adapter.rename_relation(temp_relation, relation)
+            adapter.commit_if_has_connection()
+
+            return AdapterResponse("OK", rows_affected=rows_affected)
+
+
+def read_relation_as_df(adapter: BaseAdapter, relation: BaseRelation) -> pd.DataFrame:
+    """Generic version of the read_df_from_relation."""
+
+    adapter_type = adapter.type()
+
+    if adapter_type == "snowflake":
+        import dbt.adapters.fal_experimental.support.snowflake as support_snowflake
+
+        return support_snowflake.read_relation_as_df(adapter, relation)
+
+    elif adapter_type == "bigquery":
+        import dbt.adapters.fal_experimental.support.bigquery as support_bq
+
+        return support_bq.read_relation_as_df(adapter, relation)
+
+    elif adapter_type == "duckdb":
+        import dbt.adapters.fal_experimental.support.duckdb as support_duckdb
+
+        return support_duckdb.read_relation_as_df(adapter, relation)
+
+    elif adapter_type == "postgres":
+        import dbt.adapters.fal_experimental.support.postgres as support_postgres
+
+        return support_postgres.read_relation_as_df(adapter, relation)
+
+    elif adapter.type() == "athena":
+        import dbt.adapters.fal_experimental.support.athena as support_athena
+
+        return support_athena.read_relation_as_df(adapter, relation)
+
+    elif adapter.type() == "redshift":
+        import dbt.adapters.fal_experimental.support.redshift as support_redshift
+
+        return support_redshift.read_relation_as_df(adapter, relation)
+
+    else:
+        with new_connection(adapter, "fal:read_relation_as_df") as connection:
+            alchemy_engine = _get_alchemy_engine(adapter, connection)
+
+            return pd.read_sql_table(
+                con=alchemy_engine,
+                table_name=relation.identifier,
+                schema=relation.schema,
+            )
+
+
+def prepare_for_adapter(adapter: BaseAdapter, function: Any) -> Any:
+    """Prepare the given function to be used with string-like inputs
+    (for relations) on the given adapter."""
+
+    @functools.wraps(function)
+    def wrapped(quoted_relation: str, *args, **kwargs) -> Any:
+        # HACK: we need to drop the quotes from the relation parts
+        # This was introduced in https://github.com/dbt-labs/dbt-core/pull/7115
+        # and the recommended solution would be to create a macro `fal__resolve_model_name`
+        # but it is not possible thanks a macro resolution error we get by returning the db_adapter type.
+        # The overall solution could be to avoid creating a Relation and just passing the string as is to the read/write functions.
+        parts = map(
+            lambda part: part.strip(adapter.Relation.quote_character),
+            [*quoted_relation.split(".")],
+        )
+
+        relation = adapter.Relation.create(*parts, type=RelationType.Table)
+        return function(adapter, relation, *args, **kwargs)
+
+    return wrapped
+
+
+def reconstruct_adapter(
+    flags: Namespace,
+    config: RuntimeConfig,
+    manifest: Manifest,
+    macro_manifest: MacroManifest,
+) -> BaseAdapter:
+    from dbt.flags import set_flags
+    from dbt.tracking import do_not_track
+
+    # Avoid dbt tracking
+    do_not_track()
+
+    # Flags need to be set before any plugin is loaded
+    set_flags(flags)
+
+    # Prepare the plugin loading system to handle the adapter
+    factory.load_plugin(config.credentials.type)
+    factory.load_plugin(config.python_adapter_credentials.type)
+    factory.register_adapter(config)
+
+    # Initialize the adapter
+    db_adapter = factory.get_adapter(config)
+    db_adapter._macro_manifest_lazy = macro_manifest
+    reload_adapter_cache(db_adapter, manifest)
+
+    return db_adapter
+
+
+def reload_adapter_cache(adapter: BaseAdapter, manifest: Manifest) -> None:
+    with new_connection(adapter, "fal:reload_adapter_cache"):
+        adapter.set_relations_cache(manifest, True)
+
+
+@contextmanager
+def new_connection(adapter: BaseAdapter, connection_name: str) -> Connection:
+    with adapter.connection_named(connection_name):
+        yield adapter.connections.get_thread_connection()

dbt/adapters/fal_experimental/connections.py

@@ -0,0 +1,72 @@
+from dataclasses import dataclass
+from typing import Optional
+import os
+
+from dbt.adapters.base import Credentials
+from dbt.dataclass_schema import StrEnum, ExtensibleDbtClassMixin
+
+from dbt.fal.adapters.python import PythonConnectionManager
+
+
+DEFAULT_HOSTS = {
+    "cloud": "api.alpha.fal.ai",
+    "cloud-eu": "api.eu.fal.ai",
+}
+
+
+class TeleportTypeEnum(StrEnum):
+    LOCAL = "local"
+    REMOTE_S3 = "s3"
+
+
+@dataclass
+class TeleportCredentials(ExtensibleDbtClassMixin):
+    type: TeleportTypeEnum
+
+    # local
+    local_path: Optional[str] = os.getcwd()
+
+    # s3
+    s3_bucket: Optional[str] = None
+    s3_region: Optional[str] = None
+    s3_access_key_id: Optional[str] = None
+    s3_access_key: Optional[str] = None
+
+
+class FalConnectionManager(PythonConnectionManager):
+    TYPE = "fal_experimental"
+
+    @classmethod
+    def open(cls, connection):
+        raise NotImplementedError
+
+    def execute(self, compiled_code: str):
+        raise NotImplementedError
+
+    def cancel(self, connection):
+        raise NotImplementedError
+
+
+@dataclass
+class FalCredentials(Credentials):
+    default_environment: str = "local"
+    teleport: Optional[TeleportCredentials] = None
+    host: str = ""
+    key_secret: str = ""
+    key_id: str = ""
+
+    # NOTE: So we are allowed to not set them in profiles.yml
+    # they are ignored for now
+    database: str = ""
+    schema: str = ""
+
+    def __post_init__(self):
+        if self.host in list(DEFAULT_HOSTS.keys()):
+            self.host = DEFAULT_HOSTS[self.host]
+
+    @property
+    def type(self):
+        return "fal_experimental"
+
+    def _connection_keys(self):
+        return ()

dbt/adapters/fal_experimental/impl.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Iterator
+
+from dbt.adapters.base.impl import BaseAdapter
+from dbt.adapters.base.meta import AdapterMeta, available
+from dbt.adapters.base.relation import BaseRelation
+from dbt.contracts.connection import AdapterResponse
+
+from dbt.fal.adapters.teleport.info import (
+    TeleportInfo,
+    S3TeleportInfo,
+    LocalTeleportInfo,
+)
+from dbt.fal.adapters.teleport.impl import TeleportAdapter
+from dbt.fal.adapters.python.impl import PythonAdapter
+from dbt.parser.manifest import MacroManifest, Manifest, ManifestLoader
+
+from . import telemetry
+
+from .connections import FalConnectionManager, FalCredentials, TeleportTypeEnum
+
+from .teleport_adapter_support import wrap_db_adapter
+from .teleport import DataLocation, run_in_environment_with_teleport, run_with_teleport
+
+from .adapter_support import reload_adapter_cache
+from .adapter import run_in_environment_with_adapter, run_with_adapter
+
+from .utils.environments import fetch_environment, db_adapter_config
+
+
+class FalAdapterMixin(TeleportAdapter, metaclass=AdapterMeta):
+    ConnectionManager = FalConnectionManager
+
+    def __init__(self, config, db_adapter: BaseAdapter):
+        self.config = config
+        self._db_adapter = db_adapter
+
+        self._relation_data_location_cache: DataLocation = DataLocation({})
+        if self.is_teleport():
+            self._wrapper = wrap_db_adapter(self._db_adapter, self.credentials.teleport)
+
+    @classmethod
+    def type(cls):
+        return "fal_experimental"
+
+    @classmethod
+    def storage_formats(cls):
+        return ["csv", "parquet"]
+
+    @available
+    def is_teleport(self) -> bool:
+        return getattr(self.credentials, "teleport", None) is not None
+
+    @property
+    def manifest(self) -> Manifest:
+        return ManifestLoader.get_full_manifest(self.config)
+
+    @property
+    def macro_manifest(self) -> MacroManifest:
+        return self._db_adapter.load_macro_manifest()
+
+    @telemetry.log_call("experimental_submit_python_job", config=True)
+    def submit_python_job(
+        self, parsed_model: dict, compiled_code: str
+    ) -> AdapterResponse:
+        """Execute the given `compiled_code` in the target environment."""
+        config_dict = parsed_model["config"]
+
+        environment_name = config_dict.get(
+            "fal_environment",
+            self.credentials.default_environment,
+        )
+
+        machine_type = config_dict.get(
+            "fal_machine",
+            "S",
+        )
+
+        environment, is_local = fetch_environment(
+            self.config.project_root,
+            environment_name,
+            machine_type,
+            self.credentials
+        )
+
+        telemetry.log_api(
+            "experimental_submit_python_job_config",
+            config=self.config,
+            additional_props={
+                "is_teleport": self.is_teleport(),
+                "environment_is_local": is_local,
+            },
+        )
+
+        if self.is_teleport():
+            # We need to build teleport_info because we read from the external storage,
+            # we did not _localize_ the data in `teleport_from_external_storage`
+            teleport_info = self._build_teleport_info()
+            if is_local:
+                result_table_path = run_with_teleport(
+                    code=compiled_code,
+                    teleport_info=teleport_info,
+                    locations=self._relation_data_location_cache,
+                    config=db_adapter_config(self.config)
+                )
+            else:
+                result_table_path = run_in_environment_with_teleport(
+                    environment,
+                    compiled_code,
+                    teleport_info=teleport_info,
+                    locations=self._relation_data_location_cache,
+                    config=db_adapter_config(self.config),
+                    adapter_type=self._db_adapter.type()
+                )
+
+            relation = self._db_adapter.Relation.create(
+                parsed_model["database"], parsed_model["schema"], parsed_model["alias"]
+            )
+            self._sync_result_table(relation)
+
+            return AdapterResponse("OK")
+
+        else:
+            if is_local:
+                return run_with_adapter(compiled_code, self._db_adapter, self.config)
+
+            with self._invalidate_db_cache():
+                return run_in_environment_with_adapter(
+                    environment,
+                    compiled_code,
+                    db_adapter_config(self.config),
+                    self.manifest,
+                    self.macro_manifest,
+                    self._db_adapter.type()
+                )
+
+    @contextmanager
+    def _invalidate_db_cache(self) -> Iterator[None]:
+        try:
+            yield
+        finally:
+            # Since executed Python code might alter the database
+            # layout, we need to regenerate the relations cache
+            # after every time we execute a Python model.
+            #
+            # TODO: maybe propagate a list of tuples with the changes
+            # from the Python runner, so that we can tell the cache
+            # manager about what is going on instead of hard-resetting
+            # the cache-db.
+            reload_adapter_cache(self._db_adapter, self.manifest)
+
+    @property
+    def credentials(self):
+        python_creds: FalCredentials = self.config.python_adapter_credentials
+        # dbt-fal is not configured as a Python adapter,
+        # maybe we should raise an error?
+        assert python_creds is not None
+        return python_creds
+
+    def teleport_from_external_storage(
+        self, relation: BaseRelation, relation_path: str, teleport_info: TeleportInfo
+    ):
+        """
+        Store the teleport urls for later use
+        """
+
+        rel_name = teleport_info.relation_name(relation)
+        self._relation_data_location_cache[rel_name] = relation_path
+
+    def teleport_to_external_storage(
+        self, relation: BaseRelation, teleport_info: TeleportInfo
+    ):
+        # Already in external_storage, we do not have local storage
+        # Just return the path
+        return teleport_info.build_relation_path(relation)
+
+    # TODO: cache this?
+    def _build_teleport_info(self):
+        teleport_creds = self.credentials.teleport
+        assert teleport_creds
+
+        teleport_format = TeleportAdapter.find_format(self, self._wrapper)
+
+        if teleport_creds.type == TeleportTypeEnum.LOCAL:
+            assert teleport_creds.local_path
+            return LocalTeleportInfo(
+                teleport_format, teleport_creds, teleport_creds.local_path
+            )
+        elif teleport_creds.type == TeleportTypeEnum.REMOTE_S3:
+            assert teleport_creds.s3_bucket
+            return S3TeleportInfo(
+                teleport_format, teleport_creds, teleport_creds.s3_bucket, "teleport"
+            )
+        else:
+            raise NotImplementedError(
+                f"Teleport credentials of type {teleport_creds.type} not supported"
+            )
+
+    ######
+    # HACK: Following implementations only necessary until dbt-core adds Teleport.
+    #####
+    @available
+    def sync_teleport_relation(self, relation: BaseRelation):
+        """
+        Internal implementation of sync to avoid dbt-core changes
+        """
+        teleport_info = self._build_teleport_info()
+        data_path = self._wrapper.teleport_to_external_storage(relation, teleport_info)
+        self.teleport_from_external_storage(relation, data_path, teleport_info)
+
+    def _sync_result_table(self, relation: BaseRelation):
+        """
+        Internal implementation of sync to put data back into datawarehouse.
+        This is necessary because Teleport is not part of dbt-core.
+        Once it is and adapters implement it, we will sync the result table back.
+        Instead the other adapter will call `sync_teleport` and it will automatically call
+        FalAdapter's `teleport_to_external_storage` and the adapter's `teleport_from_external_storage`.
+        """
+        teleport_info = self._build_teleport_info()
+        data_path = self.teleport_to_external_storage(relation, teleport_info)
+        self._wrapper.teleport_from_external_storage(relation, data_path, teleport_info)
+
+
+class FalAdapter(FalAdapterMixin, PythonAdapter):
+    def __init__(self, config):
+        PythonAdapter.__init__(self, config)
+        FalAdapterMixin.__init__(self, config, self._db_adapter)
+
+        telemetry.log_api(
+            "experimental_init",
+            config=config,
+            additional_props={"is_teleport": self.is_teleport()},
+        )
+
+    @classmethod
+    def is_cancelable(cls) -> bool:
+        # TODO: maybe it is?
+        return False

dbt/adapters/fal_experimental/support/athena.py

@@ -0,0 +1,92 @@
+from typing import Any
+import six
+from dbt.adapters.base.relation import BaseRelation
+from dbt.contracts.connection import AdapterResponse
+from dbt.adapters.fal.wrappers import FalCredentialsWrapper
+import sqlalchemy
+import pandas as pd
+from dbt.adapters.base import BaseAdapter
+from urllib.parse import quote_plus
+
+
+def create_engine(adapter: BaseAdapter) -> Any:
+    if isinstance(adapter.config.credentials, FalCredentialsWrapper):
+        creds = adapter.config.credentials._db_creds
+    else:
+        # In isolated environment, credentials are AthenaCredentials
+        creds = adapter.config.credentials
+    conn_str = ("awsathena+rest://:@athena.{region_name}.amazonaws.com:443/"
+                "{schema_name}?s3_staging_dir={s3_staging_dir}"
+                "&location={location}&compression=snappy")
+    return sqlalchemy.create_engine(conn_str.format(
+        region_name=creds.region_name,
+        schema_name=creds.schema,
+        s3_staging_dir=quote_plus(creds.s3_staging_dir),
+        location=quote_plus(creds.s3_staging_dir)))
+
+
+def drop_relation_if_it_exists(adapter: BaseAdapter, relation: BaseRelation) -> None:
+    if adapter.get_relation(
+        database=relation.database,
+        schema=relation.schema,
+        identifier=relation.identifier,
+    ):
+        adapter.drop_relation(relation)
+
+
+def write_df_to_relation(adapter, dataframe, relation, if_exists) -> AdapterResponse:
+
+    assert adapter.type() == "athena"
+    if isinstance(adapter.config.credentials, FalCredentialsWrapper):
+        creds = adapter.config.credentials._db_creds
+    else:
+        # In isolated environment, credentials are AthenaCredentials
+        creds = adapter.config.credentials
+
+    # This is a quirk of dbt-athena-community, where they set
+    # relation.schema = relation.identifier
+    temp_relation = relation.replace_path(
+        schema=relation.database,
+        database=creds.database,
+        # athena complanes when table location has x.__y
+        identifier=f"dbt_fal_temp_{relation.schema}"
+    )
+
+    relation = temp_relation.replace_path(identifier=relation.schema)
+
+    drop_relation_if_it_exists(adapter, temp_relation)
+
+    alchemy_engine = create_engine(adapter)
+
+    rows_affected = dataframe.to_sql(
+        con=alchemy_engine,
+        name=temp_relation.identifier,
+        schema=temp_relation.schema,
+        if_exists=if_exists,
+        index=False,
+    )
+
+    adapter.cache.add(temp_relation)
+
+    drop_relation_if_it_exists(adapter, relation)
+
+
+    # athena doesn't let us rename relations, so we do it by hand
+    stmt = f"create table {relation} as select * from {temp_relation} with data"
+    adapter.execute(six.text_type(stmt).strip())
+    adapter.cache.add(relation)
+    adapter.drop_relation(temp_relation)
+
+    adapter.commit_if_has_connection()
+    return AdapterResponse("OK", rows_affected=rows_affected)
+
+def read_relation_as_df(adapter: BaseAdapter, relation: BaseRelation) -> pd.DataFrame:
+    alchemy_engine = create_engine(adapter)
+
+    # This is dbt-athena-community quirk, table_name=relation.schema
+
+    return pd.read_sql_table(
+        con=alchemy_engine,
+        table_name=relation.schema,
+        schema=relation.database,
+    )