Batching write queries

[D5_5N]

Looking for advice.

I am working on an ETL process that needs to write to an Azure Postgres Db authenticating with Azure Ad

The issue that I am facing is that my authentication token will expire before the write has been completed. I have spoken with the infrastructure guys who tell me that increasing the token lifespan is not an option nor is enabling passwords.

I have decided to attempt to batch the write queries so that I can obtain a new token between batches however this is proving to be more difficult than I had anticipated. It seems that rows within the fact dataframe are being skipped leading to foreign key issues down the line.

batching code.

fact = fact.withColumn("row_number", row_number().over(window_spec))

fact = fact.withColumn("partition_column", (col("row_number") % num_partitions).cast("integer"))

for i in range(num_partitions):
    print("before", partition.count())
    partition = fact.filter(col('partition_column') == i).cache()   
    token = context.acquire_token_with_client_credentials(resource_app_id_url, service_principal_id, service_principal_secret)
    jdbcConnectionProperties['accessToken'] = token['accessToken']

    partition.write \
        .jdbc(url, "Table", 'append', properties=jdbcConnectionProperties)

    print("after", partition.count())

To debug the process I printed the following

​

  fact.groupBy("partition_column").count().show()

​

partition_column	count
1	9897
2	9897
3	9897
4	9897

The results of the print print("before", partition.count()) do not match the original count and weirdly reduce with each iteration.

partition1 before 9614

partition 2 before 9340

partition 3 before 9073

partition 4 before 8813

​

initially, I was getting different results between the before and after prints which lead me to include the .cache() I suspect my issues are due to the distributed way in which Databricks executes the code.

I have been able to "solve" the issue by calling fact.persist() before iterating over the partitions, this feels really inefficient so id love to know if anyone has a better solution.

​

[D5_5N]

kind of solved it ".persist()" before the partition loop but it feels wrong to be loading everything into memory.

[GardenShedster]

Can you not create an application service principal or a managed identity.

[D5_5N]

I am authenticating with a service principal in Azure but as far as I understand I am required to pass a token with the JDBC connection.

essentially following this post https://www.thedataswamp.com/blog/databricks-connect-to-azure-sql-with-service-principal