def custom_execute(statement, *args, **kwargs):
if not isinstance(statement, SQLALCHEMY_BASES) and strip(
statement).lower().startswith("copy"):
return execute_mock_s3_copy_command(statement, engine)
if not isinstance(statement, SQLALCHEMY_BASES) and strip(
statement).lower().startswith("unload"):
return execute_mock_s3_unload_command(statement, engine)
return default_execute(statement, *args, **kwargs)
def execute(self, sql, args=None):
if isinstance(sql, Executable):
return super().execute(sql, args)
if strip(sql).lower().startswith("copy"):
mock_s3_copy_command(sql, self)
sql = "commit"
if strip(sql).lower().startswith("unload"):
mock_s3_unload_command(sql, self)
sql = "commit"
return super().execute(sql, args)
def execute(self, sql, args=None):
dsn_params = self.connection.get_dsn_parameters()
engine = create_engine(
"postgresql+psycopg2://{user}:{password}@{hostname}:{port}/{dbname}"
.format(
user=dsn_params["user"],
password=config["password"],
hostname=dsn_params["host"],
port=dsn_params["port"],
dbname=dsn_params["dbname"],
))
if not isinstance(
sql, TextClause) and strip(sql).lower().startswith("copy"):
return execute_mock_s3_copy_command(sql, engine)
if not isinstance(
sql, TextClause) and strip(sql).lower().startswith("unload"):
return execute_mock_s3_unload_command(sql, engine)
return super(CustomCursor, self).execute(sql, args)
def receive_before_cursor_execute(_, cursor, statement: str, parameters, context, executemany):
"""Handle the `before_cursor_execute` event.
This is where we add support for custom features such as redshift COPY/UNLOAD because
the query has already been rendered (into a string) at this point.
Notably, COPY/UNLOAD need to perform extra non-sql behavior and potentially execute
more than a single query and the interface requires that we return a query. Thus,
we return a no-op query to be executed by sqlalchemy for certain kinds of supported
extra features.
"""
normalized_statement = strip(statement).lower()
if normalized_statement.startswith("unload"):
mock_s3_unload_command(statement, cursor)
return "SELECT 1", {}
if normalized_statement.startswith("copy"):
mock_s3_copy_command(statement, cursor)
context.should_autocommit = True
return "SELECT 1", {}
return statement, parameters
def _parse_s3_command(statement):
"""Format and Parse 'UNLOAD' command."""
statement = strip(statement)
params = dict()
# deleting 'unload'
tokens = statement.split()[1:]
# Fetching select statement
select_statement = ""
error_flag = False
for index, token in enumerate(tokens):
if token.lower() == "to":
tokens = tokens[index:]
break
select_statement += " " + token
params["select_statement"] = select_statement
if error_flag:
raise ValueError((
"Possibly malformed SELECT Statement. "
"Statement = {statement}"
"Redshift fixture only supports S3 Unload statments with the following syntax: "
"UNLOAD ('select-statement') TO 's3://object-path/name-prefix'"
"authorization 'aws_access_key_id=<aws_access_key_id>;"
"aws_secret_access_key=<aws_secret_access_key>'"
"[GZIP] [DELIMITER [ AS ] 'delimiter-char']").format(
statement=statement))
# Fetching s3_uri
if tokens.pop(0).lower() != "to":
raise ValueError((
"Possibly malformed S3 URI Format. "
"Statement = {statement}"
"Redshift fixture only supports S3 Unload statments with the following syntax: "
"UNLOAD ('select-statement') TO 's3://object-path/name-prefix'"
"authorization 'aws_access_key_id=<aws_access_key_id>;"
"aws_secret_access_key=<aws_secret_access_key>'"
"[GZIP] [DELIMITER [ AS ] 'delimiter-char']").format(
statement=statement))
params["s3_uri"] = strip(tokens.pop(0))
# Fetching authorization
for token in tokens:
if "aws_access_key_id" in token.lower(
) or "aws_secret_access_key" in token.lower():
# This is because of the following possibiliteis:
# ... [with ]authorization[ AS] 'aws_access_key_id=x;aws_secret_access_key=y'
# OR
# ... [with ]authorization[ AS] 'aws_secret_access_key=y;aws_access_key_id=x'
# OR
# ... [with ]authorization[ AS] 'aws_secret_access_key=y;\naws_access_key_id=x'
# OR
# ... [with ]authorization[ AS] 'aws_secret_access_key=y; aws_access_key_id=x'
# Supportred AWS authorization format:
# [with ]authorization[ AS] 'aws_secret_access_key=y; aws_access_key_id=x'
# No Support for additional credential formats, eg IAM roles, etc, yet.
credentials_list = token.split(";")
for credentials in credentials_list:
if "aws_access_key_id" in credentials:
params["aws_access_key_id"] = credentials.split("=")[-1]
elif "aws_secret_access_key" in credentials:
params["aws_secret_access_key"] = credentials.split(
"=")[-1]
else:
raise ValueError((
"Possibly malformed AWS Credentials Format. "
"Statement = {statement}"
"Redshift fixture only supports S3 Copy statments with the following "
"syntax: COPY <table_name> FROM [(column 1, [column2, [..]])] '"
"<file path on S3 bucket>' "
"credentials 'aws_access_key_id=<aws_access_key_id>;"
"aws_secret_access_key=<aws_secret_access_key>' "
"Supportred AWS credentials format: "
"[with ]credentials[ AS] 'aws_secret_access_key=y; aws_access_key_id=x'"
" No Support for additional credential formats, eg IAM roles, etc, yet."
).format(statement=statement))
# Fetching GZIP Flag
params["gzip"] = False
for token in tokens:
if strip(token.lower()) == "gzip":
params["gzip"] = True
# Fetching delimiter
for index, token in enumerate(tokens):
if token.lower() == "delimiter":
try:
if tokens[index + 1].lower() != "as":
params["delimiter"] = strip(tokens[index + 1])
else:
params["delimiter"] = strip(tokens[index + 2])
except IndexError:
raise ValueError((
"Possibly malformed Delimiter Format. "
"Statement = {statement}"
"Redshift fixture only supports S3 Unload statments with the following"
"syntax: UNLOAD ('select-statement') TO 's3://object-path/name-prefix'"
"authorization 'aws_access_key_id=<aws_access_key_id>;"
"aws_secret_access_key=<aws_secret_access_key>'"
"[GZIP] [DELIMITER [ AS ] 'delimiter-char']").format(
statement=statement))
return params