in python/dataproc_templates/jdbc/jdbc_to_jdbc.py [0:0]
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()