Add direct GCS export to DatabricksSqlOperator with Parquet/Avro support

[haseebmalik18]

Adds direct GCS export capability to DatabricksSqlOperator with Parquet and Avro format support.

closes: #55128

Changes

• Added parquet and avro to supported output_format values
• Added GCS URI support (gs://bucket/path) in output_path parameter
• New parameters: gcp_conn_id, gcs_impersonation_chain
• Added optional [gcs] dependency for Google provider

[jason810496]

Nice! Thanks for the PR, LGTM overall.

[jason810496]

After this PR, the Databricks provider will depend on the GCP provider. Eventually, the Databricks provider will depend on all three cloud providers (AWS, Azure, and GCP), right?

https://www.databricks.com/company/partners/cloud-partners

[jason810496]

I'm wondering we could move this kind of common serialization logic to common.io package In the long term. There are similar utils for serializing different output file format in both GCS and S3 Operators:

airflow/providers/google/src/airflow/providers/google/cloud/transfers/sql_to_gcs.py

Lines 377 to 385 in 26a9d3b

	def _configure_csv_file(self, file_handle, schema):
	"""Configure a csv writer with the file_handle and write schema as headers for the new file."""
	csv_writer = csv.writer(file_handle, delimiter=self.field_delimiter)
	csv_writer.writerow(schema)
	return csv_writer
	def _configure_parquet_file(self, file_handle, parquet_schema) -> pq.ParquetWriter:
	parquet_writer = pq.ParquetWriter(file_handle.name, parquet_schema)
	return parquet_writer

airflow/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py

Lines 234 to 255 in 26a9d3b

	for group_name, df in self._partition_dataframe(df=data_df):
	buf = io.BytesIO()
	self.log.info("Writing data to in-memory buffer")
	clean_key = self._strip_suffixes(self.s3_key)
	object_key = (
	f"{clean_key}_{group_name}{file_options.suffix}"
	if group_name
	else f"{clean_key}{file_options.suffix}"
	)
	if self.file_format != FILE_FORMAT.PARQUET and self.df_kwargs.get("compression") == "gzip":
	object_key += ".gz"
	df_kwargs = {k: v for k, v in self.df_kwargs.items() if k != "compression"}
	with gzip.GzipFile(fileobj=buf, mode="wb", filename=object_key) as gz:
	getattr(df, file_options.function)(gz, **df_kwargs)
	else:
	if self.file_format == FILE_FORMAT.PARQUET:
	getattr(df, file_options.function)(buf, **self.df_kwargs)
	else:
	text_buf = io.TextIOWrapper(buf, encoding="utf-8", write_through=True)
	getattr(df, file_options.function)(text_buf, **self.df_kwargs)
	text_buf.flush()

[haseebmalik18]

I'm wondering we could move this kind of common serialization logic to common.io package In the long term. There are similar utils for serializing different output file format in both GCS and S3 Operators:

airflow/providers/google/src/airflow/providers/google/cloud/transfers/sql_to_gcs.py

Lines 377 to 385 in 26a9d3b

	def _configure_csv_file(self, file_handle, schema):
	"""Configure a csv writer with the file_handle and write schema as headers for the new file."""
	csv_writer = csv.writer(file_handle, delimiter=self.field_delimiter)
	csv_writer.writerow(schema)
	return csv_writer
	def _configure_parquet_file(self, file_handle, parquet_schema) -> pq.ParquetWriter:
	parquet_writer = pq.ParquetWriter(file_handle.name, parquet_schema)
	return parquet_writer

airflow/providers/amazon/src/airflow/providers/amazon/aws/transfers/sql_to_s3.py

Lines 234 to 255 in 26a9d3b

	for group_name, df in self._partition_dataframe(df=data_df):
	buf = io.BytesIO()
	self.log.info("Writing data to in-memory buffer")
	clean_key = self._strip_suffixes(self.s3_key)
	object_key = (
	f"{clean_key}_{group_name}{file_options.suffix}"
	if group_name
	else f"{clean_key}{file_options.suffix}"
	)
	if self.file_format != FILE_FORMAT.PARQUET and self.df_kwargs.get("compression") == "gzip":
	object_key += ".gz"
	df_kwargs = {k: v for k, v in self.df_kwargs.items() if k != "compression"}
	with gzip.GzipFile(fileobj=buf, mode="wb", filename=object_key) as gz:
	getattr(df, file_options.function)(gz, **df_kwargs)
	else:
	if self.file_format == FILE_FORMAT.PARQUET:
	getattr(df, file_options.function)(buf, **self.df_kwargs)
	else:
	text_buf = io.TextIOWrapper(buf, encoding="utf-8", write_through=True)
	getattr(df, file_options.function)(text_buf, **self.df_kwargs)
	text_buf.flush()

Yes, I agree I think this can be possibly be opened up as another issue?

[haseebmalik18]

After this PR, the Databricks provider will depend on the GCP provider. Eventually, the Databricks provider will depend on all three cloud providers (AWS, Azure, and GCP), right?

https://www.databricks.com/company/partners/cloud-partners

Yep, that's right

[haseebmalik18]

I'm not quite sure why the Docker build test is failing. The errors show Microsoft's apt repository returning 403 Forbidden during apt-get update, which appears unrelated to the Databricks provider changes

[jscheffl]

Restarted, was a problem in the backend. Assuming CI will turn green in a moment.

[jason810496]

Yes, I agree I think this can be possibly be opened up as another issue?

Yes, it's non-blocking. We could just create issue to track it as follow-up.

[eladkal]

I'm getting build failure on #60719
related to this PR

  Resolved 940 packages in 4.00s
     Building apache-airflow-providers-apache-hdfs @ file:///opt/airflow/providers/apache/hdfs
     Building fastavro==1.9.4
        Built apache-airflow-providers-apache-hdfs @ file:///opt/airflow/providers/apache/hdfs
    × Failed to build `fastavro==1.9.4`
    
    help: `fastavro` (v1.9.4) was included because
          `apache-airflow-providers-apache-hdfs` (v4.11.2) depends on
          `hdfs[dataframe]` (v2.7.3) which depends on `fastavro`

https://github.com/apache/airflow/actions/runs/21098174150/job/60678759100?pr=60719#step:8:1265

[haseebmalik18]

@jason810496 For the build issue Elad mentioned, the error shows fastavro 1.9.4 uses deprecated C APIs that were removed in Python 3.13. Would bumping fastavro up to >=1.10.0 work?

[jscheffl]

Attempted fix PR: #60732