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

dbt/adapters/fal_experimental/support/bigquery.py

@@ -0,0 +1,74 @@
+from dbt.adapters.base import BaseAdapter, BaseRelation
+from dbt.adapters.base.connections import AdapterResponse
+from dbt.adapters.fal_experimental.adapter_support import new_connection
+import pandas as pd
+
+# [bigquery] extras dependencies
+import google.cloud.bigquery as bigquery
+from google.cloud.bigquery.job import WriteDisposition
+
+from dbt.adapters.bigquery import BigQueryAdapter, BigQueryConnectionManager
+
+def read_relation_as_df(adapter: BaseAdapter, relation: BaseRelation) -> pd.DataFrame:
+    sql = f"SELECT * FROM {relation}"
+
+    assert adapter.type() == "bigquery"
+
+    with new_connection(adapter, "fal-bigquery:read_relation_as_df") as conn:
+
+        connection_manager: BaseConnectionManager = adapter.connections  # type: ignore
+        client: bigquery.Client = connection_manager.get_thread_connection().handle  # type: ignore
+
+        job = client.query(sql)
+        df = job.to_dataframe()
+
+    return df
+
+
+def write_df_to_relation(
+    adapter: BigQueryAdapter,
+    data: pd.DataFrame,
+    relation: BaseRelation,
+) -> AdapterResponse:
+
+    assert adapter.type() == "bigquery"
+
+    project: str = relation.database  # type: ignore
+    dataset: str = relation.schema  # type: ignore
+    table: str = relation.identifier  # type: ignore
+
+    with new_connection(adapter, "fal-bigquery:write_df_to_relation") as conn:
+        connection_manager: BigQueryConnectionManager = adapter.connections
+        client: bigquery.Client = conn.handle
+
+        table_ref = bigquery.TableReference(
+            bigquery.DatasetReference(project, dataset), table
+        )
+
+        job_config = bigquery.LoadJobConfig(
+            # Specify a (partial) schema. All columns are always written to the
+            # table. The schema is used to assist in data type definitions.
+            schema=[
+                # TODO: offer as a dbt.config parameter?
+                # bigquery.SchemaField.from_api_repr(field)
+                # for field in (fields_schema or [])
+            ],
+            source_format="PARQUET",
+            write_disposition=WriteDisposition.WRITE_TRUNCATE,
+        )
+
+        with connection_manager.exception_handler("START JOB"):
+            job = client.load_table_from_dataframe(
+                data, table_ref, job_config=job_config
+            )
+
+        timeout = connection_manager.get_job_execution_timeout_seconds(conn) or 300
+
+        with connection_manager.exception_handler("LOAD TABLE"):
+            adapter.poll_until_job_completes(job, timeout)
+
+        query_table = client.get_table(job.destination)
+        num_rows = query_table.num_rows
+
+    # TODO: better AdapterResponse
+    return AdapterResponse("OK", rows_affected=num_rows)

dbt/adapters/fal_experimental/support/duckdb.py

@@ -0,0 +1,28 @@
+from dbt.adapters.base import BaseAdapter, BaseRelation
+from dbt.adapters.base.connections import AdapterResponse
+from dbt.adapters.fal_experimental.adapter_support import new_connection
+import pandas as pd
+from dbt.adapters.sql import SQLAdapter
+import duckdb
+
+
+def read_relation_as_df(adapter: BaseAdapter, relation: BaseRelation) -> pd.DataFrame:
+    db_path = adapter.config.credentials.path
+
+    con = duckdb.connect(database=db_path)
+    df = con.execute(f"SELECT * FROM {relation.schema}.{relation.identifier}").fetchdf()
+    return df
+
+
+def write_df_to_relation(
+    adapter: SQLAdapter,
+    data: pd.DataFrame,
+    relation: BaseRelation,
+) -> AdapterResponse:
+
+    db_path = adapter.config.credentials.path    
+    con = duckdb.connect(database=db_path)
+    rows_affected = con.execute(
+        f"CREATE OR REPLACE TABLE {relation.schema}.{relation.identifier} AS SELECT * FROM data;"
+    ).fetchall()[0][0]
+    return AdapterResponse("OK", rows_affected=rows_affected)

dbt/adapters/fal_experimental/support/postgres.py

@@ -0,0 +1,88 @@
+import csv
+from io import StringIO
+
+import pandas as pd
+import sqlalchemy
+
+from dbt.adapters.base import BaseRelation
+from dbt.adapters.base.connections import AdapterResponse
+from dbt.adapters.fal_experimental.adapter_support import drop_relation_if_it_exists, new_connection
+from dbt.adapters.postgres import PostgresAdapter
+
+
+def read_relation_as_df(
+    adapter: PostgresAdapter, relation: BaseRelation
+) -> pd.DataFrame:
+    assert adapter.type() == "postgres"
+
+    with new_connection(adapter, "fal-postgres:read_relation_as_df") as connection:
+        # If the given adapter supports the DBAPI (PEP 249), we can
+        # use its connection directly for the engine.
+        alchemy_engine = sqlalchemy.create_engine(
+            "postgresql+psycopg2://",
+            creator=lambda *args, **kwargs: connection.handle,
+        )
+
+        return pd.read_sql_table(
+            con=alchemy_engine,
+            table_name=relation.identifier,
+            schema=relation.schema,
+        )
+
+
+def write_df_to_relation(
+    adapter: PostgresAdapter,
+    data: pd.DataFrame,
+    relation: BaseRelation,
+    *,
+    if_exists: str = "replace",
+) -> AdapterResponse:
+    assert adapter.type() == "postgres"
+
+    with new_connection(adapter, "fal-postgres: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 = sqlalchemy.create_engine(
+            "postgresql+psycopg2://",
+            creator=lambda *args, **kwargs: connection.handle,
+        )
+
+        # TODO: probably worth handling errors here an returning
+        # a proper adapter response.
+        rows_affected = data.to_sql(
+            con=alchemy_engine,
+            name=temp_relation.identifier,
+            schema=temp_relation.schema,
+            if_exists=if_exists,
+            index=False,
+            method=_psql_insert_copy,
+        )
+        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 _psql_insert_copy(table, conn, keys, data_iter):
+    """Alternative to_sql method for PostgreSQL.
+
+    Adapted from https://pandas.pydata.org/pandas-docs/stable/user_guide/io.html#io-sql-method
+    """
+    dbapi_conn = conn.connection
+    with dbapi_conn.cursor() as cur:
+        s_buf = StringIO()
+        writer = csv.writer(s_buf)
+        writer.writerows(data_iter)
+        s_buf.seek(0)
+
+        columns = ", ".join((f'"{k}"' for k in keys))
+        table_name = f"{table.schema}.{table.name}" if table.schema else table.name
+
+        sql = f"COPY {table_name} ({columns}) FROM STDIN WITH CSV"
+        cur.copy_expert(sql=sql, file=s_buf)

dbt/adapters/fal_experimental/support/redshift.py

@@ -0,0 +1,56 @@
+import csv
+from io import StringIO
+
+import pandas as pd
+import sqlalchemy
+import awswrangler as wr
+
+from dbt.adapters.base import BaseRelation
+from dbt.adapters.base.connections import AdapterResponse
+
+from dbt.adapters.fal_experimental.adapter_support import new_connection, drop_relation_if_it_exists
+
+from dbt.adapters.redshift import RedshiftAdapter
+
+
+def read_relation_as_df(
+    adapter: RedshiftAdapter, relation: BaseRelation
+) -> pd.DataFrame:
+    sql = f"SELECT * FROM {relation}"
+
+    assert adapter.type() == "redshift"
+
+    with new_connection(adapter, "fal-redshift:read_relation_as_df") as conn:
+        df = wr.redshift.read_sql_query(sql, con=conn.handle)
+        return df
+
+
+def write_df_to_relation(
+    adapter: RedshiftAdapter,
+    data: pd.DataFrame,
+    relation: BaseRelation,
+) -> AdapterResponse:
+
+    assert adapter.type() == "redshift"
+
+    with new_connection(adapter, "fal-redshift: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)
+
+        wr.redshift.to_sql(
+            data,
+            connection.handle,
+            table=temp_relation.identifier,
+            schema=temp_relation.schema,
+            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")

dbt/adapters/fal_experimental/support/snowflake.py

@@ -0,0 +1,76 @@
+import pandas as pd
+
+from dbt.adapters.base import BaseRelation
+from dbt.adapters.base.connections import AdapterResponse
+
+from dbt.adapters.fal_experimental.adapter_support import new_connection
+
+# [snowflake] extras dependencies
+import snowflake.connector as snowflake
+
+from dbt.adapters.snowflake import SnowflakeAdapter, SnowflakeConnectionManager
+
+
+def read_relation_as_df(
+    adapter: SnowflakeAdapter, relation: BaseRelation
+) -> pd.DataFrame:
+    sql = f"SELECT * FROM {relation}"
+
+    assert adapter.type() == "snowflake"
+
+    with new_connection(adapter, "fal-snowflake:read_relation_as_df") as conn:
+        handle: snowflake.SnowflakeConnection = conn.handle
+        cur = handle.cursor()
+
+        cur.execute(sql)
+        df: pd.DataFrame = cur.fetch_pandas_all()
+
+        # HACK: manually parse ARRAY and VARIANT since they are returned as strings right now
+        # Related issue: https://github.com/snowflakedb/snowflake-connector-python/issues/544
+        for desc in cur.description:
+            # 5=VARIANT, 10=ARRAY -- https://docs.snowflake.com/en/user-guide/python-connector-api.html#type-codes
+            if desc.type_code in [5, 10]:
+                import json
+
+                df[desc.name] = df[desc.name].map(lambda v: json.loads(v))
+
+        return df
+
+
+def write_df_to_relation(
+    adapter: SnowflakeAdapter,
+    data: pd.DataFrame,
+    relation: BaseRelation,
+) -> AdapterResponse:
+    import snowflake.connector.pandas_tools as snowflake_pandas
+
+    assert adapter.type() == "snowflake"
+
+    database: str = relation.database  # type: ignore
+    schema: str = relation.schema  # type: ignore
+    table: str = relation.identifier  # type: ignore
+
+    with new_connection(adapter, "fal-snowflake:write_df_to_relation") as conn:
+        connection_manager: SnowflakeConnectionManager = adapter.connections  # type: ignore
+        handle: snowflake.SnowflakeConnection = conn.handle
+
+        with connection_manager.exception_handler("LOAD TABLE"):
+            success, _, num_rows, output = snowflake_pandas.write_pandas(
+                handle,
+                data,
+                table_name=table,
+                database=database,
+                schema=schema,
+                overwrite=True,  # TODO: This helps when table schema changes, but it is not atomic
+                auto_create_table=True,
+                quote_identifiers=False,
+            )
+            if not success:
+                # In case the failure does not raise by itself
+                # I have not been able to reproduce such a case
+                from dbt.exceptions import DbtDatabaseError
+
+                raise DbtDatabaseError(output)
+
+            # TODO: better AdapterResponse
+            return AdapterResponse(str(output[0][1]), rows_affected=num_rows)

dbt/adapters/fal_experimental/support/trino.py

@@ -0,0 +1,26 @@
+from typing import Any, Dict
+from dbt.adapters.base import BaseAdapter, Credentials
+from trino.sqlalchemy import URL
+from dbt.adapters.trino.connections import TrinoCredentials
+import sqlalchemy
+
+def create_engine(adapter: BaseAdapter) -> Any:
+    creds = adapter.config.credentials
+
+    connect_args = _build_connect_args(creds)
+
+    url = URL(
+        host=creds.host,
+        port=creds.port,
+        catalog=creds.database,
+        user=creds.user
+    )
+    return sqlalchemy.create_engine(url, connect_args=connect_args)
+
+def _build_connect_args(credentials: TrinoCredentials) -> Dict[str, Any]:
+    # See:
+    # https://github.com/starburstdata/dbt-trino/blob/master/dbt/adapters/trino/connections.py
+    return {
+        "auth": credentials.trino_auth(),
+        "http_scheme": credentials.http_scheme.value
+    }

dbt/adapters/fal_experimental/telemetry/__init__.py

@@ -0,0 +1 @@
+from .telemetry import *