# Copyright 2022 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 # # https://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. from typing import Dict, Sequence, Optional, Any from logging import Logger import argparse import pprint import sys from pyspark.sql import SparkSession, DataFrame from dataproc_templates import BaseTemplate import dataproc_templates.util.template_constants as constants import dataproc_templates.util.secret_manager as secret_manager __all__ = ['JDBCToJDBCTemplate'] class JDBCToJDBCTemplate(BaseTemplate): """ Dataproc template implementing loads from JDBC into JDBC """ @staticmethod def parse_args(args: Optional[Sequence[str]] = None) -> Dict[str, Any]: parser: argparse.ArgumentParser = argparse.ArgumentParser() groupinput = parser.add_mutually_exclusive_group(required=True) groupinput.add_argument( f'--{constants.JDBCTOJDBC_INPUT_URL}', dest=constants.JDBCTOJDBC_INPUT_URL, required=False, default="", help='JDBC input URL' ) groupinput.add_argument( f'--{constants.JDBCTOJDBC_INPUT_URL_SECRET}', dest=constants.JDBCTOJDBC_INPUT_URL_SECRET, required=False, default="", help='JDBC input URL secret name' ) parser.add_argument( f'--{constants.JDBCTOJDBC_INPUT_DRIVER}', dest=constants.JDBCTOJDBC_INPUT_DRIVER, required=True, help='JDBC input driver name' ) parser.add_argument( f'--{constants.JDBCTOJDBC_INPUT_TABLE}', dest=constants.JDBCTOJDBC_INPUT_TABLE, required=True, help='JDBC input table name' ) parser.add_argument( f'--{constants.JDBCTOJDBC_INPUT_PARTITIONCOLUMN}', dest=constants.JDBCTOJDBC_INPUT_PARTITIONCOLUMN, required=False, default="", help='JDBC input table partition column name' ) parser.add_argument( f'--{constants.JDBCTOJDBC_INPUT_LOWERBOUND}', dest=constants.JDBCTOJDBC_INPUT_LOWERBOUND, required=False, default="", help='JDBC input table partition column lower bound which is used to decide the partition stride' ) parser.add_argument( f'--{constants.JDBCTOJDBC_INPUT_UPPERBOUND}', dest=constants.JDBCTOJDBC_INPUT_UPPERBOUND, required=False, default="", help='JDBC input table partition column upper bound which is used to decide the partition stride' ) parser.add_argument( f'--{constants.JDBCTOJDBC_NUMPARTITIONS}', dest=constants.JDBCTOJDBC_NUMPARTITIONS, required=False, default="10", help='The maximum number of partitions that can be used for parallelism in table reading and writing. Default set to 10' ) parser.add_argument( f'--{constants.JDBCTOJDBC_INPUT_FETCHSIZE}', dest=constants.JDBCTOJDBC_INPUT_FETCHSIZE, required=False, default=0, type=int, help='Determines how many rows to fetch per round trip' ) parser.add_argument( f'--{constants.JDBCTOJDBC_SESSIONINITSTATEMENT}', dest=constants.JDBCTOJDBC_SESSIONINITSTATEMENT, required=False, default="", help='Custom SQL statement to execute in each reader database session' ) groupoutput = parser.add_mutually_exclusive_group(required=True) groupoutput.add_argument( f'--{constants.JDBCTOJDBC_OUTPUT_URL}', dest=constants.JDBCTOJDBC_OUTPUT_URL, required=False, default="", help='JDBC input URL' ) groupoutput.add_argument( f'--{constants.JDBCTOJDBC_OUTPUT_URL_SECRET}', dest=constants.JDBCTOJDBC_OUTPUT_URL_SECRET, required=False, default="", help='JDBC input URL secret name' ) parser.add_argument( f'--{constants.JDBCTOJDBC_OUTPUT_DRIVER}', dest=constants.JDBCTOJDBC_OUTPUT_DRIVER, required=True, help='JDBC output driver name' ) parser.add_argument( f'--{constants.JDBCTOJDBC_OUTPUT_TABLE}', dest=constants.JDBCTOJDBC_OUTPUT_TABLE, required=True, help='JDBC output table name' ) parser.add_argument( f'--{constants.JDBCTOJDBC_OUTPUT_CREATE_TABLE_OPTION}', dest=constants.JDBCTOJDBC_OUTPUT_CREATE_TABLE_OPTION, required=False, default="", help='This option allows setting of database-specific table and partition options when creating a output table' ) parser.add_argument( f'--{constants.JDBCTOJDBC_OUTPUT_MODE}', dest=constants.JDBCTOJDBC_OUTPUT_MODE, required=False, default=constants.OUTPUT_MODE_APPEND, help=( 'Output write mode ' '(one of: append,overwrite,ignore,errorifexists) ' '(Defaults to append)' ), choices=[ constants.OUTPUT_MODE_OVERWRITE, constants.OUTPUT_MODE_APPEND, constants.OUTPUT_MODE_IGNORE, constants.OUTPUT_MODE_ERRORIFEXISTS ] ) parser.add_argument( f'--{constants.JDBCTOJDBC_OUTPUT_BATCH_SIZE}', dest=constants.JDBCTOJDBC_OUTPUT_BATCH_SIZE, required=False, default="1000", help='JDBC output batch size. Default set to 1000' ) parser.add_argument( f'--{constants.JDBCTOGCS_TEMP_VIEW_NAME}', dest=constants.JDBCTOGCS_TEMP_VIEW_NAME, required=False, default="", help='Temp view name for creating a spark sql view on source data. This name has to match with the table name that will be used in the SQL query' ) parser.add_argument( f'--{constants.JDBCTOJDBC_SQL_QUERY}', dest=constants.JDBCTOJDBC_SQL_QUERY, required=False, default="", help='SQL query for data transformation. This must use the temp view name as the table to query from.' ) known_args: argparse.Namespace known_args, _ = parser.parse_known_args(args) if getattr(known_args, constants.JDBCTOJDBC_SQL_QUERY) and not getattr(known_args, constants.JDBCTOGCS_TEMP_VIEW_NAME): sys.exit('ArgumentParser Error: Temp view name cannot be null if you want to do data transformations with query') return vars(known_args) def run(self, spark: SparkSession, args: Dict[str, Any]) -> None: logger: Logger = self.get_logger(spark=spark) # Arguments #check if secret is passed or the connection string in URL if str(args[constants.JDBCTOJDBC_INPUT_URL])=="": input_jdbc_url: str = secret_manager.access_secret_version(args[constants.JDBCTOJDBC_INPUT_URL_SECRET]) else: input_jdbc_url: str = args[constants.JDBCTOJDBC_INPUT_URL] input_jdbc_driver: str = args[constants.JDBCTOJDBC_INPUT_DRIVER] input_jdbc_table: str = args[constants.JDBCTOJDBC_INPUT_TABLE] input_jdbc_partitioncolumn: str = args[constants.JDBCTOJDBC_INPUT_PARTITIONCOLUMN] input_jdbc_lowerbound: str = args[constants.JDBCTOJDBC_INPUT_LOWERBOUND] input_jdbc_upperbound: str = args[constants.JDBCTOJDBC_INPUT_UPPERBOUND] jdbc_numpartitions: str = args[constants.JDBCTOJDBC_NUMPARTITIONS] input_jdbc_fetchsize: int = args[constants.JDBCTOJDBC_INPUT_FETCHSIZE] input_jdbc_sessioninitstatement: str = args[constants.JDBCTOJDBC_SESSIONINITSTATEMENT] #check if secret is passed or the connection string in URL if str(args[constants.JDBCTOJDBC_OUTPUT_URL])=="": output_jdbc_url: str = secret_manager.access_secret_version(args[constants.JDBCTOJDBC_OUTPUT_URL_SECRET]) else: output_jdbc_url: str = args[constants.JDBCTOJDBC_OUTPUT_URL] output_jdbc_driver: str = args[constants.JDBCTOJDBC_OUTPUT_DRIVER] output_jdbc_table: str = args[constants.JDBCTOJDBC_OUTPUT_TABLE] output_jdbc_create_table_option: str = args[constants.JDBCTOJDBC_OUTPUT_CREATE_TABLE_OPTION] output_jdbc_mode: str = args[constants.JDBCTOJDBC_OUTPUT_MODE] output_jdbc_batch_size: int = args[constants.JDBCTOJDBC_OUTPUT_BATCH_SIZE] temp_view: str = args[constants.JDBCTOGCS_TEMP_VIEW_NAME] sql_query: str = args[constants.JDBCTOJDBC_SQL_QUERY] ignore_keys = {constants.JDBCTOJDBC_INPUT_URL, constants.JDBCTOJDBC_OUTPUT_URL} filtered_args = {key:val for key,val in args.items() if key not in ignore_keys} logger.info( "Starting JDBC to JDBC spark job with parameters:\n" f"{pprint.pformat(filtered_args)}" ) # Read input_data: DataFrame partition_parameters = str(input_jdbc_partitioncolumn) + str(input_jdbc_lowerbound) + str(input_jdbc_upperbound) if ((partition_parameters != "") & ((input_jdbc_partitioncolumn == "") | (input_jdbc_lowerbound == "") | (input_jdbc_upperbound == ""))): logger.error("Set all the sql partitioning parameters together-jdbctojdbc.input.partitioncolumn,jdbctojdbc.input.lowerbound,jdbctojdbc.input.upperbound. Refer to README.md for more instructions.") exit (1) properties = {constants.JDBC_URL: input_jdbc_url, constants.JDBC_DRIVER: input_jdbc_driver, constants.JDBC_TABLE: input_jdbc_table, constants.JDBC_NUMPARTITIONS: jdbc_numpartitions, constants.JDBC_FETCHSIZE: input_jdbc_fetchsize} if input_jdbc_sessioninitstatement: properties[constants.JDBC_SESSIONINITSTATEMENT] = input_jdbc_sessioninitstatement if partition_parameters: properties.update({constants.JDBC_PARTITIONCOLUMN: input_jdbc_partitioncolumn, constants.JDBC_LOWERBOUND: input_jdbc_lowerbound, constants.JDBC_UPPERBOUND: input_jdbc_upperbound}) input_data = spark.read \ .format(constants.FORMAT_JDBC) \ .options(**properties) \ .load() if sql_query: # Create temp view on source data input_data.createGlobalTempView(temp_view) # Execute SQL output_data = spark.sql(sql_query) else: output_data = input_data # Write output_data.write \ .format(constants.FORMAT_JDBC) \ .option(constants.JDBC_URL, output_jdbc_url) \ .option(constants.JDBC_DRIVER, output_jdbc_driver) \ .option(constants.JDBC_TABLE, output_jdbc_table) \ .option(constants.JDBC_CREATE_TABLE_OPTIONS, output_jdbc_create_table_option) \ .option(constants.JDBC_BATCH_SIZE, output_jdbc_batch_size) \ .option(constants.JDBC_NUMPARTITIONS, jdbc_numpartitions) \ .mode(output_jdbc_mode) \ .save()