# This file is auto-generated by tools/generate_typing_context.py, please
# don't manually alter the contents.

from __future__ import annotations
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import pandas as pd
    from typing import Optional, Dict, List, Protocol, Any
    from faldbt.project import DbtModel, DbtTest, DbtSource, Feature
    from fal.fal_script import Context, CurrentModel

    class _List_Sources(Protocol):
        def __call__(self) -> List[DbtSource]:
            """
            List tables available for `source` usage
            """
            ...

    class _List_Models_Ids(Protocol):
        def __call__(self) -> Dict[str, str]:
            """
            List model ids available for `ref` usage, formatting like `[ref_name, ...]`
            """
            ...

    class _List_Models(Protocol):
        def __call__(self) -> List[DbtModel]:
            """
            List models
            """
            ...

    class _List_Tests(Protocol):
        def __call__(self) -> List[DbtTest]:
            """
            List tests
            """
            ...

    class _List_Features(Protocol):
        def __call__(self) -> List[Feature]:
            ...

    class _Ref(Protocol):
        def __call__(
            self, target_1: str, target_2: Optional[str] = None
        ) -> pd.DataFrame:
            """
            Download a dbt model as a pandas.DataFrame automagically.
            """
            ...

    class _Source(Protocol):
        def __call__(
            self, target_source_name: str, target_table_name: str
        ) -> pd.DataFrame:
            """
            Download a dbt source as a pandas.DataFrame automagically.
            """
            ...

    class _Write_To_Source(Protocol):
        def __call__(
            self,
            data: pd.DataFrame,
            target_source_name: str,
            target_table_name: str,
            *,
            dtype: Any = None,
            mode: str = "append",
        ):
            """
            Write a pandas.DataFrame to a dbt source automagically.
            """
            ...

    class _Write_To_Firestore(Protocol):
        def __call__(self, df: pd.DataFrame, collection: str, key_column: str):
            """
            Write a pandas.DataFrame to a GCP Firestore collection. You must specify the column to use as key.
            """
            ...

    class _Execute_Sql(Protocol):
        def __call__(self, sql: str) -> pd.DataFrame:
            """Execute a sql query."""
            ...

    # Manually introduced annotations, update manually in tools/generate_typing_context.py template.
    class _Write_To_Model(Protocol):
        def __call__(
            self,
            data: pd.DataFrame,
            *,
            dtype: Any = None,
            mode: str = "overwrite",
            target_1: str = ...,
            target_2: Optional[str] = ...,
        ):
            """
            Write a pandas.DataFrame to a dbt model automagically.
            """
            ...


context: Context
write_to_model: _Write_To_Model

list_sources: _List_Sources
list_models_ids: _List_Models_Ids
list_models: _List_Models
list_tests: _List_Tests
list_features: _List_Features
ref: _Ref
source: _Source
write_to_source: _Write_To_Source
write_to_firestore: _Write_To_Firestore
execute_sql: _Execute_Sql
