google drive factory update

Author: C00ldudeNoonan

.env.example

@@ -4,7 +4,7 @@ MOTHERDUCK_TOKEN=
 MOTHERDUCK_DATABASE=
 MOTHERDUCK_PROD_SCHEMA=
 GOOGLE_DRIVE_FOLDER_ID=
-GOOGLE_DRIVE_CREDENTIALS_PATH=
+GOOGLE_SERVICE_ACCOUNT_JSON=
 ENVIRONMENT=
 DBT_TARGET=
 PYTHONLEGACYWINDOWSSTDIO=

econ_data_platform/econ_data_platform/assets/ingestion/realtor.py

@@ -1,12 +1,14 @@
 import dagster as dg
 from econ_data_platform.resources.motherduck import MotherDuckResource
+from econ_data_platform.resources.google_drive import (
+    GoogleDriveResource,
+    google_drive_resource,
+)
 
 import polars as pl
 from io import StringIO
-from google.oauth2 import service_account
-from googleapiclient.discovery import build
+
 from dataclasses import dataclass
-from collections.abc import Sequence
 
 from datetime import datetime
 import os
@@ -20,7 +22,9 @@ class DriveFile:
     modifiedTime: str
 
 
-def realtor_asset_factory(file_definition: DriveFile) -> dg.Definitions:
+def realtor_asset_factory(
+    file_definition: DriveFile, google_drive: GoogleDriveResource
+) -> dg.Definitions:
     file_name, _ = os.path.splitext(file_definition.name)
     file_id = file_definition.id
 
@@ -30,13 +34,14 @@ def realtor_asset_factory(file_definition: DriveFile) -> dg.Definitions:
         kinds={"polars", "duckdb", "google_drive"},
     )
     def read_csv_from_drive(
-        context: dg.AssetExecutionContext, md: MotherDuckResource
+        context: dg.AssetExecutionContext,
+        md: MotherDuckResource,
+        google_drive: GoogleDriveResource,
     ) -> dg.MaterializeResult:
         """Read CSV directly from Google Drive file ID into a polars DataFrame"""
         context.log.info(f"Reading file {file_name} from Google Drive")
-        request = service.files().get_media(fileId=file_id)
-        content = request.execute()
-        csv_string = content.decode("utf-8")
+        request = google_drive.request_content(file_id)
+        csv_string = request.decode("utf-8")
         df = pl.read_csv(StringIO(csv_string))
         md.drop_create_duck_db_table(file_name, df)
 
@@ -57,14 +62,12 @@ def read_csv_from_drive(
         job_name=f"{file_name}_job",
         minimum_interval_seconds=15,
     )
-    def file_sensor(context):
+    def file_sensor(context, google_drive: GoogleDriveResource):
         # Get current modification time from cursor
         last_mtime = float(context.cursor) if context.cursor else 0
 
         # Get file details from Drive
-        file_metadata = (
-            service.files().get(fileId=file_id, fields="modifiedTime").execute()
-        )
+        file_metadata = google_drive.get_file_metadata(file_id)
         context.log.info(f"File metadata: {file_metadata}")
 
         current_mtime = datetime.strptime(
@@ -83,24 +86,16 @@ def file_sensor(context):
         assets=[read_csv_from_drive],
         jobs=[file_job],
         sensors=[file_sensor],
+        resources={"google_drive": google_drive},
     )
 
 
-SCOPES = ["https://www.googleapis.com/auth/drive.readonly"]
-credentials = service_account.Credentials.from_service_account_file(
-    "creds.json", scopes=SCOPES
-)
-service = build("drive", "v3", credentials=credentials)
+# Fetch files from the Google Drive folder using properly initialized _client
 folder_id = os.environ.get("GOOGLE_DRIVE_FOLDER_ID", "")
+file_results = google_drive_resource.retrieve_files(folder_id).get("files", [])
 
-# Get files from folder
-query = f"'{folder_id}' in parents and mimeType='text/csv'"
-results = (
-    service.files()
-    .list(q=query, fields="files(id, name, createdTime, modifiedTime)")
-    .execute()
-)
-
+# Create realtor definitions dynamically
 realtor_definitions = [
-    realtor_asset_factory(DriveFile(**file)) for file in results.get("files", [])
+    realtor_asset_factory(DriveFile(**file), google_drive_resource)
+    for file in file_results
 ]

econ_data_platform/econ_data_platform/resources/google_drive.py

@@ -0,0 +1,77 @@
+import json
+import os
+from datetime import datetime
+from io import StringIO
+
+import dagster as dg
+import polars as pl
+from google.oauth2 import service_account
+from googleapiclient.discovery import build
+from pydantic import BaseModel, PrivateAttr
+
+
+class DriveFile(BaseModel):
+    id: str
+    name: str
+    createdTime: str
+    modifiedTime: str
+
+
+class GoogleDriveClient:
+    """Handles the Google Drive client creation and API interactions."""
+
+    def __init__(self, credentials_json: dict):
+        self.credentials = service_account.Credentials.from_service_account_info(
+            credentials_json, scopes=["https://www.googleapis.com/auth/drive.readonly"]
+        )
+        self.service = build("drive", "v3", credentials=self.credentials)
+
+    def retrieve_files(self, folder_id: str):
+        """Query for files in a Google Drive folder."""
+        query = f"'{folder_id}' in parents and mimeType='text/csv'"
+        return (
+            self.service.files()
+            .list(q=query, fields="files(id, name, createdTime, modifiedTime)")
+            .execute()
+        )
+
+    def request_content(self, file_id: str):
+        """Fetch file content from Google Drive using file_id."""
+        request = self.service.files().get_media(fileId=file_id)
+        return request.execute()
+
+    def get_file_metadata(self, file_id: str, fields: str = "modifiedTime"):
+        """Get metadata for a specific file."""
+        return self.service.files().get(fileId=file_id, fields=fields).execute()
+
+
+class GoogleDriveResource(dg.ConfigurableResource):
+    """Resource configuration for Google Drive credentials."""
+
+    json_data: str
+
+    _client: GoogleDriveClient = PrivateAttr()
+
+    def setup_for_execution(self, context: dg.InitResourceContext):
+        """Initialize the Google Drive client using the credentials."""
+        credentials_json = json.loads(self.json_data)
+        self._client = GoogleDriveClient(credentials_json)
+
+    def retrieve_files(self, folder_id: str):
+        """Delegates to the client to retrieve files."""
+        return self._client.retrieve_files(folder_id)
+
+    def request_content(self, file_id: str):
+        """Delegates to the client to fetch content."""
+        return self._client.request_content(file_id)
+
+    def get_file_metadata(self, file_id: str, fields: str = "modifiedTime"):
+        """Delegates to the client to get file metadata."""
+        return self._client.get_file_metadata(file_id, fields)
+
+
+google_drive_resource = GoogleDriveResource(
+    json_data=os.environ["GOOGLE_SERVICE_ACCOUNT_JSON"]
+)
+
+google_drive_resource.setup_for_execution(dg.build_init_resource_context())