from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.utils.context import Context
from config import Config
from util import Util
[docs]
class ExportQueryOperator(BaseOperator):
"""
Takes a query and exports the result to a specified flat file format (CSV, Parquet, etc.)
and uploads it to a given destination (S3, GDrive, etc.).
:param sql_query: The SQL query to execute in string form.
If you're executing against Redshift or Postgres, ensure that the query has schema-qualified table names.
Ex: 'SELECT * FROM california.pit_address_metadata' is valid,
while 'SELECT * FROM pit_address_metadata' is not since the schema is not being passed as a parameter.
:param database: Which database to execute the query in.
Valid options:
- shoreline
- readydatalake
- redshift_dw_(dev/prod)
- postgres_bossplatform_grapi_(dev/stage/prod)
- postgres_bossplatform_game_(dev/stage/prod)
:param filename: The name of the file to create (without extension).
:param file_format: The format of the file to export.
Valid types: csv, parquet, json, xlsx, geojson, shp, kml, kmz
:param destination: Where to send the file once it has been created.
Valid options: default, gdrive, public_bucket_dev, public_bucket_stage, public_bucket_prod
The default option looks up the preferred destination for the given state (requires that parameter)
Choosing an unreachable environment bucket from the executing environment will result in an error.
:param gdrive_url: The Google Drive folder URL to upload the file to (only necessary for GDrive).
By default, uses the flat file exporter folder in the File Transfer GDrive.
:param state: The 2 letter state abbreviation of where the file is going (only necessary for GDrive).
:param extra_tenants: Optional non-state tenant codes (e.g. ["PUBLIC"]) to accept as a valid `state`
in addition to Config.STATE_NAMES. Widens the allow-list only; use for tenants that intentionally
live outside the shared state map.
:param json_orient: If exporting to JSON, the orient parameter to pass to pandas.DataFrame.to_json().
Valid options: split, records, index, columns, values, table
:param geometry_config: Configuration dictionary for geometry export formats (geojson, shp, kml, kmz).
There are three types of geometry config supported:
1. Point geometry: {'type': 'point', 'lon_col': 'lon_column_name', 'lat_col': 'lat_column_name'}
2. WKT geometry: {'type': 'wkt', 'geom_col': 'wkt_column_name'}
3. Binary geometry: {'type': 'binary', 'geom_col': 'geometry_column_name'}
Also honored for `parquet`: when set, the file is written as GeoParquet.
:param parquet_compression: Compression algorithm for parquet files.
Valid options: snappy, gzip, brotli, zstd, None
Default is 'snappy'. Use 'zstd' for better compression ratio.
"""
ui_color = "#420a69"
template_fields = (
"sql_query",
"filename",
"file_format",
"destination",
"gdrive_url",
"state",
"json_orient",
"parquet_compression",
"partition_column",
"partition_drop_column",
)
@apply_defaults
def __init__(
self,
sql_query: str,
database: str,
filename: str,
file_format: str = "csv",
destination: str = "default",
gdrive_url: str = Config.GDRIVE_FOLDER_BASE_URI + Config.GDRIVE_FOLDER_FLAT_EXPORTER,
state: str = "XY",
extra_tenants: list[str] | None = None,
json_orient: str = "records",
geometry_config: dict = {},
parquet_compression: str = "snappy",
partition_column: str | None = None,
partition_drop_column: bool = True,
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.sql_query = sql_query
self.database = database
self.filename = filename
self.file_format = file_format
self.destination = destination
self.gdrive_url = gdrive_url
self.state = state
self.extra_tenants = extra_tenants or [] # e.g. ["PUBLIC"]
self.json_orient = json_orient
self.geometry_config = geometry_config
self.parquet_compression = parquet_compression
self.partition_drop_column = partition_drop_column
self.partition_column = partition_column
self.PUBLIC_BUCKET_STATES = ["CA", "TX"]
self.GEOMETERY_TYPES = ["shp", "kml", "kmz", "geojson"]
self.PUBLIC_BUCKETS = {
"public_bucket_dev": Config.BOSS_PUBLIC_BUCKET,
"public_bucket_stage": Config.BOSS_PUBLIC_BUCKET + "-stage",
"public_bucket_prod": Config.BOSS_PUBLIC_BUCKET + "-prod",
}
[docs]
def get_valid_destinations(self) -> list[str]:
"""
Returns a list of valid destinations the export can go based on the environment.
"""
if Util.is_prod_instance():
return ["default", "gdrive", "public_bucket_prod", "public_bucket_stage"]
return ["default", "gdrive", "public_bucket_dev"]
[docs]
def get_valid_databases(self) -> list[str]:
"""
Returns a list of valid databases that can be used for executing against.
"""
if Util.is_prod_instance():
return [
Config.DATA_LAKE_DB,
Config.REDSHIFT,
Config.GRAPI_PG,
Config.GAME_PG,
Config.GRAPI_PG_STAGE,
Config.GAME_PG_STAGE,
]
return [
Config.DATA_LAKE_DB,
Config.REDSHIFT_DEV,
Config.GRAPI_PG,
Config.GAME_PG,
]
[docs]
def validate_geometry_config(self) -> None:
"""
Validates the geometry configuration dictionary. Throws an exception if invalid.
"""
wants_geoparquet = self.file_format == "parquet" and bool(self.geometry_config)
if self.file_format not in self.GEOMETERY_TYPES and not wants_geoparquet:
return
if "type" not in self.geometry_config:
raise Exception("Geometry config must have a 'type' key.")
if self.geometry_config["type"] == "point":
if (
"lon_col" not in self.geometry_config
or "lat_col" not in self.geometry_config
):
raise Exception(
"Point geometry config must have 'lon_col' and 'lat_col' keys."
)
if not isinstance(self.geometry_config["lon_col"], str) or not isinstance(
self.geometry_config["lat_col"], str
):
raise Exception(
"'lon_col' and 'lat_col' must be strings representing column names."
)
elif self.geometry_config["type"] == "wkt":
if "geom_col" not in self.geometry_config:
raise Exception("WKT geometry config must have 'geom_col' key.")
if not isinstance(self.geometry_config["geom_col"], str):
raise Exception(
"'geom_col' must be a string representing the column name."
)
elif self.geometry_config["type"] == "binary":
if "geom_col" not in self.geometry_config:
raise Exception("Binary geometry config must have 'geom_col' key.")
if not isinstance(self.geometry_config["geom_col"], str):
raise Exception(
"'geom_col' must be a string representing the column name."
)
else:
raise Exception(
"Geometry config 'type' must be one of: 'point', 'wkt', 'binary'."
)
[docs]
def execute_via_athena(self, query: str, database: str):
"""
Executes the given SQL query against the specified Athena database and returns the result as a DataFrame.
"""
import awswrangler as wr
print(f"Executing query in Athena against database {database}")
return wr.athena.read_sql_query(sql=query, database=database)
[docs]
def execute_via_postgres(self, query: str, database: str):
"""
Executes the given SQL query against the specified Postgres database and returns the result as a DataFrame.
"""
from airflow.providers.postgres.hooks.postgres import PostgresHook
print(f"Executing query in Postgres against database {database}")
src_hook = PostgresHook(postgres_conn_id=database)
return src_hook.get_pandas_df(query)
[docs]
def execute_via_redshift(self, query: str, database: str):
"""
Executes the given SQL query against the specified Redshift database and returns the result as a DataFrame.
"""
from airflow.providers.amazon.aws.hooks.redshift_sql import RedshiftSQLHook
print(f"Executing query in Redshift against database {database}")
src_hook = RedshiftSQLHook(
redshift_conn_id=database, database=database.split("_")[-1]
)
return src_hook.get_pandas_df(query)
[docs]
def create_geo_dataframe(self, df):
"""
Converts the given DataFrame to a GeoDataFrame based on the geometry configuration and returns it.
"""
import geopandas as gpd
if self.geometry_config["type"] == "wkt":
geom_col = gpd.GeoSeries.from_wkt(
df[self.geometry_config["geom_col"]].astype(str), crs="EPSG:4326"
)
return gpd.GeoDataFrame(df, geometry=geom_col, crs="EPSG:4326")
if self.geometry_config["type"] == "binary":
return gpd.GeoDataFrame(
df, geometry=self.geometry_config["geom_col"], crs="EPSG:4326"
)
if self.geometry_config["type"] == "point":
lon_col = self.geometry_config["lon_col"]
lat_col = self.geometry_config["lat_col"]
return gpd.GeoDataFrame(
df,
geometry=gpd.points_from_xy(df[lon_col], df[lat_col]),
crs="EPSG:4326",
)
raise Exception("Unsupported geometry type for GeoDataFrame creation.")
[docs]
def convert_csv(self, df):
"""
Converts the given DataFrame to CSV format and returns the temporary file object.
"""
from tempfile import NamedTemporaryFile
tmp = NamedTemporaryFile(delete=False, suffix=".csv")
df.to_csv(tmp.name, index=False)
return tmp
[docs]
def convert_parquet(self, df):
"""
Converts the given DataFrame to Parquet format and returns the temporary file object.
Uses configurable compression (default: snappy, recommended: zstd for better ratio).
When `geometry_config` is set, the frame is converted to a GeoDataFrame first and
written as GeoParquet (geometry column + geo metadata, EPSG:4326).
"""
from tempfile import NamedTemporaryFile
tmp = NamedTemporaryFile(delete=False, suffix=".parquet")
if self.geometry_config:
gdf = self.create_geo_dataframe(df)
gdf.to_parquet(tmp.name, index=False, compression=self.parquet_compression)
return tmp
df.to_parquet(
tmp.name,
index=False,
compression=self.parquet_compression,
use_dictionary=True,
)
return tmp
[docs]
def convert_json(self, df, orient):
"""
Converts the given DataFrame to JSON format and returns the temporary file object.
"""
from tempfile import NamedTemporaryFile
tmp = NamedTemporaryFile(delete=False, suffix=".json")
df.to_json(tmp.name, orient=orient)
return tmp
[docs]
def convert_xlsx(self, df, sheet_name: str):
"""
Converts the given DataFrame to XLSX format and returns the temporary file object.
"""
from tempfile import NamedTemporaryFile
tmp = NamedTemporaryFile(delete=False, suffix=".xlsx")
df.to_excel(tmp.name, sheet_name=sheet_name, header=True, index=False)
return tmp
[docs]
def convert_geojson(self, df):
"""
Converts the given DataFrame to GeoJSON format and returns the temporary file object.
"""
from tempfile import NamedTemporaryFile
gdf = self.create_geo_dataframe(df)
tmp = NamedTemporaryFile(delete=False, suffix=".geojson")
gdf.to_file(tmp.name, index=False, driver="GeoJSON")
return tmp
[docs]
def convert_shp(self, df):
"""
Converts the given DataFrame to Shapefile format and returns the temporary file object.
"""
from tempfile import TemporaryDirectory, NamedTemporaryFile
import os
import zipfile
gdf = self.create_geo_dataframe(df)
with TemporaryDirectory() as tmpdir:
shp_path = os.path.join(tmpdir, "data.shp")
gdf.to_file(shp_path, index=False)
zip_path = os.path.join(tmpdir, "data_shapefile.zip")
with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zf:
for root, _, files in os.walk(tmpdir):
for file in files:
if file.startswith("data.") and not file.endswith(".zip"):
zf.write(os.path.join(root, file), arcname=file)
final_tmp = NamedTemporaryFile(delete=False, suffix=".zip")
os.rename(zip_path, final_tmp.name)
return final_tmp
[docs]
def convert_kml(self, df):
"""
Converts the given DataFrame to KML format and returns the temporary file object.
"""
from tempfile import NamedTemporaryFile
import fiona
# https://github.com/geopandas/geopandas/issues/2481
fiona.drvsupport.supported_drivers["KML"] = "rw"
gdf = self.create_geo_dataframe(df)
tmp = NamedTemporaryFile(delete=False, suffix=".kml")
gdf.to_file(tmp.name, index=False, driver="KML")
return tmp
[docs]
def convert_kmz(self, df):
"""
Converts the given DataFrame to KMZ format and returns the temporary file object.
"""
from tempfile import NamedTemporaryFile
import zipfile
import os
temp_kml = self.convert_kml(df)
kmz_tmp = NamedTemporaryFile(delete=False, suffix=".kmz")
zipfile.ZipFile(kmz_tmp.name, mode="w").write(temp_kml.name)
os.remove(temp_kml.name)
return kmz_tmp
[docs]
def get_export_extension(self) -> str:
"""
Returns the file extension for the export. Shapefiles are zipped, so return 'zip' for shp.
"""
return "zip" if self.file_format == "shp" else self.file_format
[docs]
def export_gdrive(self, local_file_path: str) -> str:
"""
Exports the given local file to Google Drive and returns the file URL.
"""
enviro = "prod" if Util.is_prod_instance() else "dev"
gdrive_service = Util.construct_gdrive_service()
parent_folder_id = self.gdrive_url.split("/")[-1]
if parent_folder_id == Config.GDRIVE_FOLDER_FLAT_EXPORTER:
parent_folder_id = Util.gdrive_create_folder(
drive_service=gdrive_service,
folder_name=str(self.state.upper() + "_" + enviro),
parent_folder_id=self.gdrive_url.split("/")[-1],
)
url = Util.gdrive_upload_file(
drive_service=gdrive_service,
file_name=self.filename + "." + self.get_export_extension(),
file_path=local_file_path,
parent_folder_id=parent_folder_id,
overwrite=True,
)
return "https://drive.google.com/file/d/" + url
[docs]
def export_s3(self, local_file_path: str) -> str:
"""
Exports the given local file to S3 and returns a presigned URL.
"""
export_key = self.filename + "." + self.get_export_extension()
Util.store_s3_file(
bucket=self.PUBLIC_BUCKETS[self.destination],
key=export_key,
file_path=local_file_path,
overwrite=True,
)
url = Util.get_presigned_url(
bucket=self.PUBLIC_BUCKETS[self.destination],
key=export_key,
)
return url
def _resolve_destination(self):
"""
Resolves the 'default' destination to the actual destination based on state preferences.
"""
if self.destination != "default":
return
VALID_DESTINATIONS = self.get_valid_destinations()
self.destination = (
VALID_DESTINATIONS[2]
if self.state in self.PUBLIC_BUCKET_STATES
else VALID_DESTINATIONS[1]
)
def _execute_single(self, df, filename: str) -> dict:
"""
Converts a DataFrame to the configured file format, exports it to the destination,
and returns metadata about the export. This is the core convert-export pipeline
used for both single-file and partitioned exports.
"""
import os
original_filename = self.filename
self.filename = filename
try:
if self.file_format == "csv":
tmp_file = self.convert_csv(df)
elif self.file_format == "parquet":
tmp_file = self.convert_parquet(df)
elif self.file_format == "json":
tmp_file = self.convert_json(df, self.json_orient)
elif self.file_format == "xlsx":
tmp_file = self.convert_xlsx(
df, sheet_name=filename[:31]
)
elif self.file_format == "geojson":
tmp_file = self.convert_geojson(df)
elif self.file_format == "shp":
tmp_file = self.convert_shp(df)
elif self.file_format == "kml":
tmp_file = self.convert_kml(df)
elif self.file_format == "kmz":
tmp_file = self.convert_kmz(df)
else:
raise Exception("Unsupported file format for conversion.")
if self.destination == "gdrive":
file_url = self.export_gdrive(tmp_file.name)
else:
file_url = self.export_s3(tmp_file.name)
print(f"File successfully exported to {file_url}")
os.remove(tmp_file.name)
return {
"filename": filename,
"url": file_url,
"extension": self.get_export_extension(),
}
finally:
self.filename = original_filename
def _execute_partitioned(self, df) -> list[dict]:
"""
Splits the DataFrame by distinct values of self.partition_column and exports
one file per distinct value. Returns a list of export result dicts.
"""
import re
import pandas as pd
if self.partition_column not in df.columns:
raise Exception(
f"partition_column '{self.partition_column}' not found in query results. "
f"Available columns: {list(df.columns)}"
)
distinct_values = df[self.partition_column].dropna().unique()
has_nulls = df[self.partition_column].isna().any()
n_partitions = len(distinct_values) + (1 if has_nulls else 0)
if n_partitions > 150:
raise Exception(
f"partition_column '{self.partition_column}' has {n_partitions} distinct values, "
f"which exceeds the safety limit of 150. Choose a lower-cardinality column."
)
print(f"Partitioning on column '{self.partition_column}': {n_partitions} partitions")
results = []
for value in sorted(distinct_values, key=str):
partition_df = df[df[self.partition_column] == value]
if self.partition_drop_column:
partition_df = partition_df.drop(columns=[self.partition_column])
if partition_df.empty:
print(f"Skipping empty partition for {self.partition_column}={value}")
continue
safe_value = re.sub(r'[^\x20-\x7E]', '', str(value))[:200].strip()
if not safe_value:
safe_value = 'empty'
partition_filename = f"{self.filename}_{safe_value}"
print(f"Exporting partition {self.partition_column}={value} ({len(partition_df)} rows)")
result = self._execute_single(partition_df, partition_filename)
results.append(result)
if has_nulls:
null_df = df[df[self.partition_column].isna()]
if self.partition_drop_column:
null_df = null_df.drop(columns=[self.partition_column])
if not null_df.empty:
print(f"Exporting partition {self.partition_column}=NULL ({len(null_df)} rows)")
result = self._execute_single(null_df, f"{self.filename}_NULL")
results.append(result)
print(f"Partitioning complete: {len(results)} files exported")
return results
[docs]
def execute(self, context: Context):
VALID_DATABASES = self.get_valid_databases()
VALID_FILE_FORMATS = [
"csv",
"parquet",
"json",
"xlsx",
"geojson",
"shp",
"kml",
"kmz",
]
VALID_DESTINATIONS = self.get_valid_destinations()
VALID_STATES = list(Config.STATE_NAMES.keys()) + self.extra_tenants
VALID_JSON_ORIENTS = ["split", "records", "index", "columns", "values", "table"]
# Validate parameters
if self.database not in VALID_DATABASES:
raise Exception(
f"Unknown database {self.database}. Valid options: "
+ ", ".join(VALID_DATABASES)
)
if self.file_format not in VALID_FILE_FORMATS:
raise Exception(
f"Unknown file format {self.file_format}. Valid options: "
+ ", ".join(VALID_FILE_FORMATS)
)
if self.destination not in VALID_DESTINATIONS:
raise Exception(
f"Unknown destination {self.destination}. Valid options: "
+ ", ".join(VALID_DESTINATIONS)
)
if self.state not in VALID_STATES:
raise Exception(
f"Unknown state {self.state}. Valid options: " + ", ".join(VALID_STATES)
)
if self.file_format == "json" and self.json_orient not in VALID_JSON_ORIENTS:
raise Exception(
f"Unknown json orient {self.json_orient}. Valid options: "
+ ", ".join(VALID_JSON_ORIENTS)
)
self.validate_geometry_config()
# Execute the query and get results as a DataFrame
if self.database == Config.DATA_LAKE_DB:
df = self.execute_via_athena(self.sql_query, self.database)
elif "redshift" in self.database:
df = self.execute_via_redshift(self.sql_query, self.database)
else:
df = self.execute_via_postgres(self.sql_query, self.database)
print(f"Query executed successfully, retrieved {len(df)} rows.")
# Resolve default destination before export
self._resolve_destination()
if self.partition_column:
return self._execute_partitioned(df)
# return so the export result ({filename, url, extension}) lands in XCom
# — downstream tasks (e.g. bead_sar_export.notify_export_urls) read the
# URL from here. The partitioned branch above already returns its list.
return self._execute_single(df, self.filename)