# Copyright 2024 Google LLC # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Processes input data to fit to DLP inspection standards.""" import dataclasses from enum import Enum from typing import List, Tuple, Dict from google.api_core.exceptions import NotFound from google.cloud import bigquery, dlp_v2 from google.cloud.sql.connector import Connector from sqlalchemy import create_engine, MetaData, Table, func, select, inspect @dataclasses.dataclass class Bigquery: """Represents a connection to a Google BigQuery dataset and table.""" bq_client: bigquery.Client dataset: str table: str @dataclasses.dataclass class CloudSQL: """Represents a connection to a Google CloudSQL.""" connector: Connector connection_name: str service_account: str db_name: str table: str driver: str connection_type: str class Database(Enum): """Represents available sources for database connections.""" BIGQUERY = "bigquery" CLOUDSQL = "cloudsql" class Preprocessing: """Converts input data into Data Loss Prevention tables.""" def __init__( self, source: str, project: str, zone: str, **preprocess_args ): """Initializes `Preprocessing` class with arguments. Args: source (str): The name of the source of data used. project (str): The name of the Google Cloud Platform project. zone (str): The name of the zone. **preprocess_args: Additional arguments for preprocessing. Supported arguments are: - bigquery_args(Dict): - dataset (str): The name of the BigQuery dataset. - table (str, optional): The name of the BigQuery table. If not provided, the entire dataset is scanned. Optional. Defaults to None. - cloudsql_args(Dict): - instance (str): Name of the database instance. - service_account(str): Service account email to be used. - db_name(str): The name of the database. - table (str): The name of the table. - db_type(str): The type of the database. e.g. postgres, mysql. """ self.source = Database(source) self.project = project self.zone = zone if self.source == Database.BIGQUERY: # Handle BigQuery source. bigquery_args = preprocess_args.get("bigquery_args", {}) self.bigquery = Bigquery(bigquery.Client(project=project), bigquery_args["dataset"], bigquery_args["table"]) elif self.source == Database.CLOUDSQL: # Handle Cloud SQL source. cloudsql_args = preprocess_args.get("cloudsql_args", {}) instance = cloudsql_args["instance"] db_type = cloudsql_args["db_type"] # Determine the appropriate database driver and connection # name based on db_type. if db_type == "mysql": driver = "pymysql" connection_name = f"mysql+{driver}" elif db_type == "postgres": driver = "pg8000" connection_name = f"postgresql+{driver}" self.cloudsql = CloudSQL( Connector(), f"{project}:{zone}:{instance}", cloudsql_args["service_account"], cloudsql_args["db_name"], cloudsql_args["table"], driver, connection_name) def get_connection(self): """Returns a connection to the database. Returns: A connection object that can be used to execute queries. """ connector = self.cloudsql.connector.connect( self.cloudsql.connection_name, self.cloudsql.driver, enable_iam_auth=True, user=self.cloudsql.service_account, db=self.cloudsql.db_name ) return connector def get_cloudsql_tables(self): """Returns a list of all tables in the CloudSQL database. Returns: A list of table names. """ # Create a database engine instance. engine = create_engine( f'{self.cloudsql.connection_type}://', creator=self.get_connection) table_names = inspect(engine).get_table_names() return table_names def get_cloudsql_data( self, table: str, batch_size: int, start_index: int ) -> Tuple[List, List]: """Retrieves the schema and content of a table from CloudSQL. Args: table (str): The name of the table. batch_size (int): The block of cells to be analyzed. start_index (int): The starting index of each block to be analyzed. Returns: Tuple(List, List): A tuple containing the schema and content as a List. """ # Create a database engine instance. engine = create_engine( f'{self.cloudsql.connection_type}://', creator=self.get_connection) # Create a Metadata and Table instance. metadata = MetaData() table = Table(table, metadata, extend_existing=True, autoload_with=engine) num_columns = len(table.columns.keys()) # Get table schema. schema = [column.name for column in table.columns] # Get table contents. with engine.connect() as connection: query = table.select().with_only_columns(table.columns) \ .limit(int(batch_size/num_columns)) \ .offset(int(start_index/num_columns)) content = list(connection.execute(query).fetchall()) return schema, content def get_bigquery_tables(self, dataset: str) -> List[str]: """Constructs a list of table names from a BigQuery dataset. Args: Dataset (str): Name of the dataset in BigQuery. Returns: List of tablenames. """ dataset_tables = list(self.bigquery.bq_client.list_tables(dataset)) table_names = [table.table_id for table in dataset_tables] return table_names def fetch_rows( self, table_bq: bigquery.table.Table, start_index: int, batch_size: int ) -> List[Dict]: """Fetches a batch of rows from a BigQuery table. Args: table_bq (bigquery.table.Table) : The path of the table where the data is fetched. start_index (int) : The starting index of each block to be analyzed. batch_size (int) : The block of cells to be analyzed. Returns: List[Dict]: A list of rows, where each row is a tuple containing the values for each field in the table schema. """ content = [] num_columns = len(table_bq.schema) rows_iter = self.bigquery.bq_client.list_rows( table=table_bq,start_index=int(start_index/num_columns), max_results=int(batch_size/num_columns)) if not rows_iter.total_rows: print(f"""The Table {table_bq.table_id} is empty. Please populate the table and try again.""") else: for row in rows_iter: content.append(tuple(row)) return content def get_table_schema(self, table_id: str) -> Tuple[List, List, List]: """Generates a schema for a given table ID. Args: table_id (str): The ID of the table for which the schema needs to be generated. Returns: tuple: A tuple containing three lists - schema, nested_columns, and record_columns. - schema (list): The list of fields in the schema. - nested_columns (list): A list with the columns of the nested columns. - record_columns (list): The columns with the record type. """ schema = [] nested_columns = [] record_columns = [] fields = table_id.schema for field in fields: record, nested, main_field = self.get_field(field) if nested: record_columns.append(main_field) nested_columns.append(record) else: schema.append(record) return schema, nested_columns, record_columns def get_field(self, field): """Generates a field for the given field object. Args: field: The field object for which the field needs to be generated. Returns: tuple: A tuple containing three values - record, nested, and main_cell. - record: The generated field or list of nested fields. - nested (bool): Indicates if the field is nested or not. - main_cell: The main field associated with the nested fields. """ # Checks if the field has nested columns. if field.field_type == "RECORD": field_names = [] for subfield in field.fields: main_cell = field.name cell = f"{field.name}.{subfield.name}" # Checks if the field has nested fields within. if subfield.field_type == "RECORD": field_names.append(self.get_field(subfield)) else: field_names.append(cell) return field_names, True, main_cell return field.name, False, False def get_query( self, table_bq: bigquery.table.Table, query_args: Dict) -> str: """Creates a SQL query as a string. Args: table_bq (bigquery.table.Table): The fully qualified name of the BigQuery table. query_args (Dict) : columns_selected (str): The string with the selected columns. unnest (str): The unnest string for the query. limit (int): The maximum number of rows to retrieve in each block. offset(int): The starting index of the rows to retrieve. Returns: str: SQL query as a string. """ query = f"""SELECT {query_args['columns_selected']} FROM `{table_bq}`, {query_args['unnest']} LIMIT {query_args['limit']} OFFSET {query_args['offset']}""" return query def get_nested_types(self, table_bq: bigquery.table.Table) -> List[str]: """ Gets the field modes of the selected table. Args: table_bq (bigquery.table.Table): The fully qualified name of the BigQuery table. Returns: List: A complete list with the field modes of the columns. """ nested_types = [field.mode for field in table_bq.schema] return nested_types def get_rows_query( self, nested_args: Dict, table_bq: bigquery.table.Table, batch_size: int, start_index: int ) -> List[Dict]: """Retrieves the content of the table. Args: nested_args (Dict): table_schema (List[Dict]): The schema of a BigQuery table. nested_columns (List[Dict]): A list with the columns of the nested columns. record_columns (List[Dict]): The columns with the record type. table_bq (bigquery.table.Table): The fully qualified name of the BigQuery table. batch_size (int): The block of cells to be analyzed. start_index (int): The starting index of each block to be analyzed. Returns: List[Dict]: The content of the BigQuery table. """ nested_types = self.get_nested_types(table_bq) if "REPEATED" in nested_types: bq_schema = nested_args["table_schema"] \ + nested_args["record_columns"] num_columns = len(bq_schema) columns_selected = ", ".join(str(column) for column in bq_schema) unnest = f"UNNEST ({nested_args['record_columns'][0]})" else: bq_schema = nested_args["table_schema"] \ + nested_args["nested_columns"] num_columns = len(bq_schema) columns_selected = ", ".join(str(column) for column in bq_schema) unnest = ( f"""UNNEST ([{nested_args['record_columns'][0]}]) as {nested_args['record_columns'][0]} """ ) # Generate the SQL query using the selected columns, # table, unnest, limit, and offset. Calculate the limit and offset for # the SQL query based on the block size. sql_query = self.get_query( table_bq, { "columns_selected": columns_selected, "unnest": unnest, "offset": int(start_index/num_columns), "limit": int(batch_size/num_columns) }) query_job = self.bigquery.bq_client.query(sql_query) query_results = query_job.result() bq_rows_content = [tuple(dict(row).values()) for row in query_results] return bq_rows_content def get_data_types(self, table_id: str) -> List: """ Gets the data types of the selected table. Args: table_id (str): The fully qualified name of the BigQuery table. Returns: List: A complete list with the data types of the columns. """ return [field.field_type for field in table_id.schema] def flatten_list(self, unflattened_list: List) -> List: """ Recursively flattens a nested list and returns a flattened list. Args: list (list): The input list that needs to be flattened. Returns: list: The flattened list. """ # Create an empty list to store the flattened elements. flattened = [] # Iterate through each element in the list. for element in unflattened_list: # If the element is a list, recursively flatten the list. if isinstance(element, list): flattened.extend(self.flatten_list(element)) else: # If the element is not a list, add it to the flattened list. flattened.append(element) # Return the flattened list. return flattened def get_bigquery_data( self, table_id: str, start_index: int, batch_size: int ) -> Tuple[List[Dict], List[Dict]]: """Retrieves the schema and content of a BigQuery table. Args: table_id (str): The fully qualified name of the BigQuery table. star_index (int): The starting index of each block to be analyzed. batch_size (int): The block of cells to be analyzed. Returns: Tuple[List[Dict], List[Dict]]: A tuple containing the BigQuery schema and content as a List of Dictionaries. """ try: table_bq = self.bigquery.bq_client.get_table(table_id) except NotFound as exc: raise ValueError(f"Error retrieving table {table_id}.") from exc dtypes = self.get_data_types(table_bq) # Checks if there are nested fields in the schema. if "RECORD" in dtypes: table_schema, nested_columns, record_columns = ( self.get_table_schema(table_bq)) table_schema = self.flatten_list(table_schema) nested_columns = self.flatten_list(nested_columns) bq_schema = table_schema + nested_columns bq_rows_content = self.get_rows_query( { "table_schema": table_schema, "nested_columns": nested_columns, "record_columns": record_columns }, table_bq, batch_size, start_index) else: table_schema = table_bq.schema bq_schema = [field.to_api_repr()["name"] for field in table_schema] bq_rows_content = self.fetch_rows( table_bq, start_index, batch_size) return bq_schema, bq_rows_content def convert_to_dlp_table(self, schema: List[Dict], content: List[Dict]) -> dlp_v2.Table: """Converts a BigQuery table into a DLP table. Converts a BigQuery table into a Data Loss Prevention table, an object that can be inspected by Data Loss Prevention. Args: schema (List[Dict]): The schema of a BigQuery table. content (List[Dict]): The content of a BigQuery table. Returns: A table object that can be inspected by Data Loss Prevention. """ table_dlp = dlp_v2.Table() table_dlp.headers = [ {"name": schema_object} for schema_object in schema ] rows = [] for row in content: rows.append(dlp_v2.Table.Row( values=[dlp_v2.Value( string_value=str(cell_val)) for cell_val in row])) table_dlp.rows = rows return table_dlp def get_tables_info(self) -> List[Tuple]: """Retrieves information about tables in a dataset from BigQuery or CloudSQL. This function is used to retrieve information necessary for running the program with dataflow, as it enables subsequent table fragmentation for parallelization. Returns: List[Tuple]: A list of tuples containing the table name and the total number of cells. """ # Retrieve table names from either a specific table # or all tables in a dataset. if self.source == Database.BIGQUERY: table_names = [self.bigquery.table] if self.bigquery.table \ else self.get_bigquery_tables(self.bigquery.dataset) elif self.source == Database.CLOUDSQL: table_names = [self.cloudsql.table] if self.cloudsql.table \ else self.get_cloudsql_tables() tables = [] if self.source == Database.BIGQUERY: for table_name in table_names: # Get the table object from BigQuery. table_bq = self.bigquery.bq_client.get_table( f"{self.bigquery.dataset}.{table_name}") # Calculate the total number of rows and columns in the table. num_rows = table_bq.num_rows dtypes = self.get_data_types(table_bq) # Checks if there are nested fields in the schema. if "RECORD" in dtypes: table_schema, _, record_columns = ( self.get_table_schema(table_bq)) num_columns = len(table_schema + record_columns) else: num_columns = len(table_bq.schema) # Append a tuple with the table name and the total number # of cells to the list. tables.append((table_name,num_rows*num_columns)) elif self.source == Database.CLOUDSQL: # Create a database engine instance. engine = create_engine( f'{self.cloudsql.connection_type}://', creator=self.get_connection ) for table_name in table_names: # Create a Metadata and Table instance. metadata = MetaData() table = Table(table_name, metadata, extend_existing=True, autoload_with=engine) num_columns = len(table.columns.keys()) # Get table contents. with engine.connect() as connection: count_query = select( # pylint: disable=E1102 func.count("*")).select_from(table) num_rows = connection.execute(count_query).scalar() tables.append((table_name,num_rows*num_columns)) return tables def get_dlp_table_per_block( self, batch_size: int, table_name: str, start_index: int ) -> dlp_v2.Table: """Constructs a DLP Table object, a partir de cada bloque de celdas. Args: batch_size (int): The block of cells to be analyzed. table_name (str): The name of the table to be analyzed. start_index (int): The starting index of each block to be analyzed. Returns: A Data Loss Prevention table object. """ if self.source == Database.BIGQUERY: schema, content = self.get_bigquery_data( f"{self.bigquery.dataset}.{table_name}", start_index, batch_size ) elif self.source == Database.CLOUDSQL: schema, content = self.get_cloudsql_data( table_name, batch_size, start_index ) return self.convert_to_dlp_table(schema, content)