def to_s3(self,
sql: str,
path: str,
connection: Any,
engine: str = "mysql") -> str:
"""
Write a query result on S3.
:param sql: SQL Query
:param path: AWS S3 path to write the data (e.g. s3://...)
:param connection: A PEP 249 compatible connection (Can be generated with Redshift.generate_connection())
:param engine: Only "mysql" by now
:return: Manifest S3 path
"""
if "mysql" not in engine.lower():
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql'!")
path = path[-1] if path[-1] == "/" else path
self._session.s3.delete_objects(path=path)
sql = f"{sql}\n" \
f"INTO OUTFILE S3 '{path}'\n" \
"FIELDS TERMINATED BY ',' OPTIONALLY ENCLOSED BY '\"' ESCAPED BY '\\\\'\n" \
"LINES TERMINATED BY '\\n'\n" \
"MANIFEST ON\n" \
"OVERWRITE ON"
with connection.cursor() as cursor:
logger.debug(sql)
cursor.execute(sql)
connection.commit()
return path + ".manifest"
def _get_schema(dataframe,
dataframe_type: str,
preserve_index: bool,
engine: str = "mysql") -> List[Tuple[str, str]]:
schema_built: List[Tuple[str, str]] = []
if "postgres" in engine.lower():
convert_func = data_types.pyarrow2postgres
elif "mysql" in engine.lower():
convert_func = data_types.pyarrow2mysql
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
if dataframe_type.lower() == "pandas":
pyarrow_schema: List[Tuple[
str, str]] = data_types.extract_pyarrow_schema_from_pandas(
dataframe=dataframe,
preserve_index=preserve_index,
indexes_position="right")
for name, dtype in pyarrow_schema:
aurora_type: str = convert_func(dtype)
schema_built.append((name, aurora_type))
else:
raise InvalidDataframeType(
f"{dataframe_type} is not a valid DataFrame type. Please use 'pandas'!"
)
return schema_built
def _validate_connection(
database: str,
host: str,
port: Union[str, int],
user: str,
password: str,
engine: str = "mysql",
tcp_keepalive: bool = True,
application_name: str = "aws-data-wrangler-validation",
validation_timeout: int = 10) -> None:
if "postgres" in engine.lower():
conn = pg8000.connect(database=database,
host=host,
port=int(port),
user=user,
password=password,
ssl=True,
application_name=application_name,
tcp_keepalive=tcp_keepalive,
timeout=validation_timeout)
elif "mysql" in engine.lower():
conn = pymysql.connect(database=database,
host=host,
port=int(port),
user=user,
password=password,
program_name=application_name,
connect_timeout=validation_timeout)
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
conn.close()
def _get_load_sql(path: str,
schema_name: str,
table_name: str,
engine: str,
region: str = "us-east-1") -> str:
if "postgres" in engine.lower():
bucket, key = Aurora._parse_path(path=path)
sql: str = (
"-- AWS DATA WRANGLER\n"
"SELECT aws_s3.table_import_from_s3(\n"
f"'{schema_name}.{table_name}',\n"
"'',\n"
"'(FORMAT CSV, DELIMITER '','', QUOTE ''\"'', ESCAPE ''\"'')',\n"
f"'({bucket},{key},{region})')")
elif "mysql" in engine.lower():
sql = (
"-- AWS DATA WRANGLER\n"
f"LOAD DATA FROM S3 MANIFEST '{path}'\n"
"REPLACE\n"
f"INTO TABLE {schema_name}.{table_name}\n"
"FIELDS TERMINATED BY ',' OPTIONALLY ENCLOSED BY '\"' ESCAPED BY '\"'\n"
"LINES TERMINATED BY '\\n'")
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
return sql
def generate_connection(database: str,
host: str,
port: Union[str, int],
user: str,
password: str,
engine: str = "mysql",
tcp_keepalive: bool = True,
application_name: str = "aws-data-wrangler",
connection_timeout: Optional[int] = None,
validation_timeout: int = 10):
"""
Generate a valid connection object.
:param database: The name of the database instance to connect with.
:param host: The hostname of the Aurora server to connect with.
:param port: The TCP/IP port of the Aurora server instance.
:param user: The username to connect to the Aurora database with.
:param password: The user password to connect to the server with.
:param engine: "mysql" or "postgres"
:param tcp_keepalive: If True then use TCP keepalive
:param application_name: Application name
:param connection_timeout: Connection Timeout
:param validation_timeout: Timeout to try to validate the connection
:return: PEP 249 compatible connection
"""
Aurora._validate_connection(database=database,
host=host,
port=port,
user=user,
password=password,
engine=engine,
tcp_keepalive=tcp_keepalive,
application_name=application_name,
validation_timeout=validation_timeout)
if "postgres" in engine.lower():
conn = pg8000.connect(database=database,
host=host,
port=int(port),
user=user,
password=password,
ssl=True,
application_name=application_name,
tcp_keepalive=tcp_keepalive,
timeout=connection_timeout)
elif "mysql" in engine.lower():
conn = pymysql.connect(database=database,
host=host,
port=int(port),
user=user,
password=password,
program_name=application_name,
connect_timeout=validation_timeout)
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
return conn
def load_table(dataframe: pd.DataFrame,
dataframe_type: str,
load_paths: List[str],
schema_name: str,
table_name: str,
connection: Any,
num_files: int,
columns: Optional[List[str]] = None,
mode: str = "append",
preserve_index: bool = False,
engine: str = "mysql",
region: str = "us-east-1"):
"""
Load text/CSV files into a Aurora table using a manifest file.
Creates the table if necessary.
:param dataframe: Pandas or Spark Dataframe
:param dataframe_type: "pandas" or "spark"
:param load_paths: S3 paths to be loaded (E.g. S3://...)
:param schema_name: Aurora schema
:param table_name: Aurora table name
:param connection: A PEP 249 compatible connection (Can be generated with Aurora.generate_connection())
:param num_files: Number of files to be loaded
:param columns: List of columns to load
:param mode: append or overwrite
:param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
:param engine: "mysql" or "postgres"
:param region: AWS S3 bucket region (Required only for postgres engine)
:return: None
"""
if "postgres" in engine.lower():
Aurora.load_table_postgres(dataframe=dataframe,
dataframe_type=dataframe_type,
load_paths=load_paths,
schema_name=schema_name,
table_name=table_name,
connection=connection,
mode=mode,
preserve_index=preserve_index,
region=region,
columns=columns)
elif "mysql" in engine.lower():
Aurora.load_table_mysql(dataframe=dataframe,
dataframe_type=dataframe_type,
manifest_path=load_paths[0],
schema_name=schema_name,
table_name=table_name,
connection=connection,
mode=mode,
preserve_index=preserve_index,
num_files=num_files,
columns=columns)
else:
raise InvalidEngine(f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!")
def _create_table(
cursor,
dataframe,
dataframe_type,
schema_name,
table_name,
preserve_index=False,
engine: str = "mysql",
columns: Optional[List[str]] = None,
varchar_default_length: int = 256,
varchar_lengths: Optional[Dict[str, int]] = None) -> None:
"""
Create Aurora table.
:param cursor: A PEP 249 compatible cursor
:param dataframe: Pandas or Spark Dataframe
:param dataframe_type: "pandas" or "spark"
:param schema_name: Redshift schema
:param table_name: Redshift table name
:param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
:param engine: "mysql" or "postgres"
:param columns: List of columns to load
:param varchar_default_length: The size that will be set for all VARCHAR columns not specified with varchar_lengths
:param varchar_lengths: Dict of VARCHAR length by columns. (e.g. {"col1": 10, "col5": 200})
:return: None
"""
sql: str = f"-- AWS DATA WRANGLER\n" \
f"DROP TABLE IF EXISTS {schema_name}.{table_name}"
logger.debug(f"Drop table query:\n{sql}")
if "postgres" in engine.lower():
cursor.execute(sql)
elif "mysql" in engine.lower():
with warnings.catch_warnings():
warnings.filterwarnings(action="ignore",
message=".*Unknown table.*")
cursor.execute(sql)
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
schema = Aurora._get_schema(
dataframe=dataframe,
dataframe_type=dataframe_type,
preserve_index=preserve_index,
engine=engine,
columns=columns,
varchar_default_length=varchar_default_length,
varchar_lengths=varchar_lengths)
cols_str: str = "".join([f"{col[0]} {col[1]},\n"
for col in schema])[:-2]
sql = f"-- AWS DATA WRANGLER\n" f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (\n" f"{cols_str})"
logger.debug(f"Create table query:\n{sql}")
cursor.execute(sql)
def _create_table(cursor,
dataframe,
dataframe_type,
schema_name,
table_name,
preserve_index=False,
engine: str = "mysql"):
"""
Creates Aurora table.
:param cursor: A PEP 249 compatible cursor
:param dataframe: Pandas or Spark Dataframe
:param dataframe_type: "pandas" or "spark"
:param schema_name: Redshift schema
:param table_name: Redshift table name
:param preserve_index: Should we preserve the Dataframe index? (ONLY for Pandas Dataframe)
:param engine: "mysql" or "postgres"
:return: None
"""
sql: str = f"-- AWS DATA WRANGLER\n" \
f"DROP TABLE IF EXISTS {schema_name}.{table_name}"
logger.debug(f"Drop table query:\n{sql}")
if "postgres" in engine.lower():
cursor.execute(sql)
elif "mysql" in engine.lower():
with warnings.catch_warnings():
warnings.filterwarnings(action="ignore",
message=".*Unknown table.*")
cursor.execute(sql)
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
schema = Aurora._get_schema(dataframe=dataframe,
dataframe_type=dataframe_type,
preserve_index=preserve_index,
engine=engine)
cols_str: str = "".join([f"{col[0]} {col[1]},\n"
for col in schema])[:-2]
sql = f"-- AWS DATA WRANGLER\n" f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} (\n" f"{cols_str})"
logger.debug(f"Create table query:\n{sql}")
cursor.execute(sql)
def _get_load_sql(path: str,
schema_name: str,
table_name: str,
engine: str,
region: str = "us-east-1",
columns: Optional[List[str]] = None) -> str:
if "postgres" in engine.lower():
bucket, key = Aurora._parse_path(path=path)
if columns is None:
cols_str: str = ""
else:
cols_str = ",".join(columns)
sql: str = (
"-- AWS DATA WRANGLER\n"
"SELECT aws_s3.table_import_from_s3(\n"
f"'{schema_name}.{table_name}',\n"
f"'{cols_str}',\n"
"'(FORMAT CSV, DELIMITER '','', QUOTE ''\"'', ESCAPE ''\"'')',\n"
f"'({bucket},{key},{region})')")
elif "mysql" in engine.lower():
if columns is None:
cols_str = ""
else:
# building something like: (@col1,@col2) set col1=@col1,col2=@col2
col_str = [f"@{x}" for x in columns]
set_str = [f"{x}=@{x}" for x in columns]
cols_str = f"({','.join(col_str)}) SET {','.join(set_str)}"
logger.debug(f"cols_str: {cols_str}")
sql = (
"-- AWS DATA WRANGLER\n"
f"LOAD DATA FROM S3 MANIFEST '{path}'\n"
"REPLACE\n"
f"INTO TABLE {schema_name}.{table_name}\n"
"FIELDS TERMINATED BY ',' OPTIONALLY ENCLOSED BY '\"' ESCAPED BY '\"'\n"
"LINES TERMINATED BY '\\n'"
f"{cols_str}")
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
return sql
def _get_schema(
dataframe,
dataframe_type: str,
preserve_index: bool,
engine: str = "mysql",
columns: Optional[List[str]] = None,
varchar_default_length: int = 256,
varchar_lengths: Optional[Dict[str,
int]] = None) -> List[Tuple[str, str]]:
varchar_lengths = {} if varchar_lengths is None else varchar_lengths
schema_built: List[Tuple[str, str]] = []
if "postgres" in engine.lower():
convert_func = data_types.pyarrow2postgres
elif "mysql" in engine.lower():
convert_func = data_types.pyarrow2mysql
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
if dataframe_type.lower() == "pandas":
pyarrow_schema: List[Tuple[
str, str]] = data_types.extract_pyarrow_schema_from_pandas(
dataframe=dataframe,
preserve_index=preserve_index,
indexes_position="right")
for name, dtype in pyarrow_schema:
if columns is None or name in columns:
varchar_len = varchar_lengths.get(name,
varchar_default_length)
aurora_type: str = convert_func(dtype=dtype,
varchar_length=varchar_len)
schema_built.append((name, aurora_type))
else:
raise InvalidDataframeType(
f"{dataframe_type} is not a valid DataFrame type. Please use 'pandas'!"
)
return schema_built
password=password,
ssl=True,
application_name=application_name,
tcp_keepalive=tcp_keepalive,
timeout=connection_timeout)
elif "mysql" in engine.lower():
conn = pymysql.connect(database=database,
host=host,
port=int(port),
user=user,
password=password,
program_name=application_name,
connect_timeout=connection_timeout)
else:
raise InvalidEngine(
f"{engine} is not a valid engine. Please use 'mysql' or 'postgres'!"
)
return conn
def write_load_manifest(
self, manifest_path: str, objects_paths: List[str]
) -> Dict[str, List[Dict[str, Union[str, bool]]]]:
manifest: Dict[str, List[Dict[str, Union[str, bool]]]] = {
"entries": []
}
path: str
for path in objects_paths:
entry: Dict[str, Union[str, bool]] = {
"url": path,
"mandatory": True
}