airbyte.caches.bigquery

39class BigQueryCache(BigQueryConfig, CacheBase):
40    """The BigQuery cache implementation."""
41
42    _sql_processor_class: ClassVar[type[SqlProcessorBase]] = BigQuerySqlProcessor
43
44    paired_destination_name: ClassVar[str | None] = "destination-bigquery"
45    paired_destination_config_class: ClassVar[type | None] = DestinationBigquery
46
47    @property
48    def paired_destination_config(self) -> DestinationBigquery:
49        """Return a dictionary of destination configuration values."""
50        return bigquery_cache_to_destination_configuration(cache=self)
51
52    def get_arrow_dataset(
53        self,
54        stream_name: str,
55        *,
56        max_chunk_size: int = DEFAULT_ARROW_MAX_CHUNK_SIZE,
57    ) -> NoReturn:
58        """Raises NotImplementedError; BigQuery doesn't support `pd.read_sql_table`.
59
60        See: https://github.com/airbytehq/PyAirbyte/issues/165
61        """
62        raise NotImplementedError(
63            "BigQuery doesn't currently support to_arrow"
64            "Please consider using a different cache implementation for these functionalities."
65        )

38class BigQueryConfig(SqlConfig):
39    """Configuration for BigQuery."""
40
41    database_name: str = Field(alias="project_name")
42    """The name of the project to use. In BigQuery, this is equivalent to the database name."""
43
44    schema_name: str = Field(alias="dataset_name", default=DEFAULT_CACHE_SCHEMA_NAME)
45    """The name of the dataset to use. In BigQuery, this is equivalent to the schema name."""
46
47    credentials_path: str | None = None
48    """The path to the credentials file to use.
49    If not passed, falls back to the default inferred from the environment."""
50
51    @property
52    def project_name(self) -> str:
53        """Return the project name (alias of self.database_name)."""
54        return self.database_name
55
56    @property
57    def dataset_name(self) -> str:
58        """Return the dataset name (alias of self.schema_name)."""
59        return self.schema_name
60
61    @overrides
62    def get_sql_alchemy_url(self) -> SecretString:
63        """Return the SQLAlchemy URL to use.
64
65        We suppress warnings about unrecognized JSON type. More info on that here:
66        - https://github.com/airbytehq/PyAirbyte/issues/254
67        """
68        warnings.filterwarnings(
69            "ignore",
70            message="Did not recognize type 'JSON' of column",
71            category=sqlalchemy.exc.SAWarning,
72        )
73
74        url: URL = make_url(f"bigquery://{self.project_name!s}")
75        if self.credentials_path:
76            url = url.update_query_dict({"credentials_path": self.credentials_path})
77
78        return SecretString(url)
79
80    @overrides
81    def get_database_name(self) -> str:
82        """Return the name of the database. For BigQuery, this is the project name."""
83        return self.project_name
84
85    def get_vendor_client(self) -> bigquery.Client:
86        """Return a BigQuery python client."""
87        if self.credentials_path:
88            credentials = service_account.Credentials.from_service_account_file(
89                self.credentials_path
90            )
91        else:
92            credentials, _ = google.auth.default()
93
94        return bigquery.Client(credentials=credentials)