Skip to content

Commit

Permalink
feat: add table description and automatically list projects
Browse files Browse the repository at this point in the history
  • Loading branch information
gabriel-milan committed Oct 20, 2023
1 parent 9c8a092 commit 3562abb
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 4 deletions.
7 changes: 3 additions & 4 deletions pipelines/rj_escritorio/data_catalog/flows.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
from pipelines.rj_escritorio.data_catalog.schedules import update_data_catalog_schedule
from pipelines.rj_escritorio.data_catalog.tasks import (
generate_dataframe_from_list_of_tables,
list_projects,
list_tables,
merge_list_of_list_of_tables,
update_gsheets_data_catalog,
Expand All @@ -28,18 +29,16 @@
],
) as rj_escritorio_data_catalog_flow:
# Parameters
project_ids = Parameter("project_ids")
spreadsheet_url = Parameter("spreadsheet_url")
sheet_name = Parameter("sheet_name")
bq_client_mode = Parameter("bq_client_mode", default="prod")
exclude_dev_projects = Parameter("exclude_dev_projects", default=True)
exclude_staging = Parameter("exclude_staging", default=True)
exclude_test = Parameter("exclude_test", default=True)
exclude_logs = Parameter("exclude_logs", default=True)

# Flow
project_ids = parse_comma_separated_string_to_list(
input_text=project_ids, output_type=str
)
project_ids = list_projects(mode=bq_client_mode, exclude_dev=exclude_dev_projects)
list_of_list_of_tables = list_tables.map(
project_id=project_ids,
mode=unmapped(bq_client_mode),
Expand Down
40 changes: 40 additions & 0 deletions pipelines/rj_escritorio/data_catalog/tasks.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,10 @@
"""
Tasks for generating a data catalog from BigQuery.
"""
from typing import List

from google.cloud import bigquery
from googleapiclient import discovery
import gspread
import pandas as pd
from prefect import task
Expand All @@ -15,6 +18,41 @@
from pipelines.utils.utils import get_credentials_from_env, log


@task
def list_projects(
mode: str = "prod",
exclude_dev: bool = True,
) -> List[str]:
"""
Lists all GCP projects that we have access to.
Args:
mode: Credentials mode.
exclude_dev: Exclude projects that ends with "-dev".
Returns:
List of project IDs.
"""
credentials = get_credentials_from_env(mode=mode)
service = discovery.build("cloudresourcemanager", "v1", credentials=credentials)
request = service.projects().list()
projects = []
while request is not None:
response = request.execute()
for project in response.get("projects", []):
project_id = project["projectId"]
if exclude_dev and project_id.endswith("-dev"):
log(f"Excluding dev project {project_id}.")
continue
log(f"Found project {project_id}.")
projects.append(project_id)
request = service.projects().list_next(
previous_request=request, previous_response=response
)
log(f"Found {len(projects)} projects.")
return projects


@task
def list_tables( # pylint: disable=too-many-arguments
project_id: str,
Expand Down Expand Up @@ -68,10 +106,12 @@ def list_tables( # pylint: disable=too-many-arguments
if exclude_test and "test" in table_id:
log(f"Excluding test table {table_id}.")
continue
table_description = table.description
table_info = {
"project_id": project_id,
"dataset_id": dataset_id,
"table_id": table_id,
"description": table_description,
"url": f"https://console.cloud.google.com/bigquery?p={project_id}&d={dataset_id}&t={table_id}&page=table",
"private": not project_id == "datario",
}
Expand Down

0 comments on commit 3562abb

Please sign in to comment.