Skip to content

Commit 199cc70

Browse files
nsheilsdev2049
andauthored
Ability to specify credentials wihen using Google BigQuery as a data loader (langchain-ai#5466)
# Adds ability to specify credentials when using Google BigQuery as a data loader Fixes langchain-ai#5465 . Adds ability to set credentials which must be of the `google.auth.credentials.Credentials` type. This argument is optional and will default to `None. Co-authored-by: Dev 2049 <[email protected]>
1 parent eab4b4c commit 199cc70

File tree

3 files changed

+30
-6
lines changed

3 files changed

+30
-6
lines changed

langchain/document_loaders/bigquery.py

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,13 @@
1-
from typing import List, Optional
1+
from __future__ import annotations
2+
3+
from typing import TYPE_CHECKING, List, Optional
24

35
from langchain.docstore.document import Document
46
from langchain.document_loaders.base import BaseLoader
57

8+
if TYPE_CHECKING:
9+
from google.auth.credentials import Credentials
10+
611

712
class BigQueryLoader(BaseLoader):
813
"""Loads a query result from BigQuery into a list of documents.
@@ -11,6 +16,7 @@ class BigQueryLoader(BaseLoader):
1116
are written into the `page_content` of the document. The `metadata_columns`
1217
are written into the `metadata` of the document. By default, all columns
1318
are written into the `page_content` and none into the `metadata`.
19+
1420
"""
1521

1622
def __init__(
@@ -19,11 +25,28 @@ def __init__(
1925
project: Optional[str] = None,
2026
page_content_columns: Optional[List[str]] = None,
2127
metadata_columns: Optional[List[str]] = None,
28+
credentials: Optional[Credentials] = None,
2229
):
30+
"""Initialize BigQuery document loader.
31+
32+
Args:
33+
query: The query to run in BigQuery.
34+
project: Optional. The project to run the query in.
35+
page_content_columns: Optional. The columns to write into the `page_content`
36+
of the document.
37+
metadata_columns: Optional. The columns to write into the `metadata` of the
38+
document.
39+
credentials : google.auth.credentials.Credentials, optional
40+
Credentials for accessing Google APIs. Use this parameter to override
41+
default credentials, such as to use Compute Engine
42+
(`google.auth.compute_engine.Credentials`) or Service Account
43+
(`google.oauth2.service_account.Credentials`) credentials directly.
44+
"""
2345
self.query = query
2446
self.project = project
2547
self.page_content_columns = page_content_columns
2648
self.metadata_columns = metadata_columns
49+
self.credentials = credentials
2750

2851
def load(self) -> List[Document]:
2952
try:
@@ -34,7 +57,7 @@ def load(self) -> List[Document]:
3457
"Please install it with `pip install google-cloud-bigquery`."
3558
) from ex
3659

37-
bq_client = bigquery.Client(self.project)
60+
bq_client = bigquery.Client(credentials=self.credentials, project=self.project)
3861
query_result = bq_client.query(self.query).result()
3962
docs: List[Document] = []
4063

poetry.lock

Lines changed: 3 additions & 4 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

pyproject.toml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ pymongo = {version = "^4.3.3", optional = true}
4040
clickhouse-connect = {version="^0.5.14", optional=true}
4141
weaviate-client = {version = "^3", optional = true}
4242
google-api-python-client = {version = "2.70.0", optional = true}
43+
google-auth = {version = "^2.18.1", optional = true}
4344
wolframalpha = {version = "5.0.0", optional = true}
4445
anthropic = {version = "^0.2.6", optional = true}
4546
qdrant-client = {version = "^1.1.2", optional = true, python = ">=3.8.1,<3.12"}
@@ -239,6 +240,7 @@ all = [
239240
"weaviate-client",
240241
"redis",
241242
"google-api-python-client",
243+
"google-auth",
242244
"wolframalpha",
243245
"qdrant-client",
244246
"tensorflow-text",

0 commit comments

Comments
 (0)