Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix cloudsql-query system tests #41092

Merged
merged 1 commit into from
Jul 30, 2024
Merged

Fix cloudsql-query system tests #41092

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion airflow/providers/google/cloud/hooks/cloud_sql.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -931,7 +931,9 @@ def _set_temporary_ssl_file(
self.log.info("Neither cert path and cert value provided. Nothing to save.")
return None

_temp_file = NamedTemporaryFile(mode="w+b", prefix="/tmp/certs/")
certs_folder = "/tmp/certs/"
Path(certs_folder).mkdir(parents=True, exist_ok=True)
_temp_file = NamedTemporaryFile(mode="w+b", prefix=certs_folder)
if cert_path:
with open(cert_path, "rb") as cert_file:
_temp_file.write(cert_file.read())
Expand Down
66 changes: 23 additions & 43 deletions tests/system/providers/google/cloud/cloud_sql/example_cloud_sql_query.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,13 @@
from airflow.decorators import task, task_group
from airflow.models.connection import Connection
from airflow.models.dag import DAG
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.operators.cloud_sql import (
CloudSQLCreateInstanceDatabaseOperator,
CloudSQLCreateInstanceOperator,
CloudSQLDeleteInstanceOperator,
CloudSQLExecuteQueryOperator,
)
from airflow.settings import Session
from airflow.utils.trigger_rule import TriggerRule
from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID

Expand All @@ -52,18 +52,6 @@
HOME_DIR = Path.home()

COMPOSER_ENVIRONMENT = os.environ.get("COMPOSER_ENVIRONMENT", "")
if COMPOSER_ENVIRONMENT:
# We assume that the test is launched in Cloud Composer environment because the reserved environment
# variable is assigned (https://cloud.google.com/composer/docs/composer-2/set-environment-variables)
GET_COMPOSER_NETWORK_COMMAND = """
gcloud composer environments describe $COMPOSER_ENVIRONMENT \
--location=$COMPOSER_LOCATION \
--project=$GCP_PROJECT \
--format="value(config.nodeConfig.network)"
"""
else:
# The test is launched locally
GET_COMPOSER_NETWORK_COMMAND = "echo"


def run_in_composer():
Expand Down Expand Up @@ -115,7 +103,7 @@ def ip_configuration() -> dict[str, Any]:
"ipv4Enabled": True,
"requireSsl": False,
"enablePrivatePathForGoogleCloudServices": True,
"privateNetwork": """{{ task_instance.xcom_pull('get_composer_network')}}""",
"privateNetwork": f"projects/{PROJECT_ID}/global/networks/default",
}
else:
# Use connection to Cloud SQL instance via Public IP from anywhere (mask 0.0.0.0/0).
Expand Down Expand Up @@ -395,12 +383,6 @@ def cloud_sql_database_create_body(instance: str) -> dict[str, Any]:
catchup=False,
tags=["example", "cloudsql", "postgres"],
) as dag:
get_composer_network = BashOperator(
task_id="get_composer_network",
bash_command=GET_COMPOSER_NETWORK_COMMAND,
do_xcom_push=True,
)

for db_provider in DB_PROVIDERS:
database_type: str = db_provider["database_type"]
cloud_sql_instance_name: str = db_provider["cloud_sql_instance_name"]
Expand Down Expand Up @@ -467,9 +449,9 @@ def create_connection(
kwargs: dict[str, Any],
) -> str | None:
session = settings.Session()
if session.query(Connection).filter(Connection.conn_id == connection_id).first():
log.warning("Connection '%s' already exists", connection_id)
return connection_id
log.info("Removing connection %s if it exists", connection_id)
query = session.query(Connection).filter(Connection.conn_id == connection_id)
query.delete()

connection: dict[str, Any] = deepcopy(kwargs)
connection["extra"]["instance"] = instance
Expand Down Expand Up @@ -523,30 +505,28 @@ def execute_queries(db_type: str):

execute_queries_task = execute_queries(db_type=database_type)

@task_group(group_id=f"teardown_{database_type}")
def teardown(instance: str, db_type: str):
task_id = f"delete_cloud_sql_instance_{db_type}"
CloudSQLDeleteInstanceOperator(
task_id=task_id,
project_id=PROJECT_ID,
instance=instance,
trigger_rule=TriggerRule.ALL_DONE,
)
@task()
def delete_connection(connection_id: str) -> None:
session = Session()
log.info("Removing connection %s", connection_id)
query = session.query(Connection).filter(Connection.conn_id == connection_id)
query.delete()
session.commit()

for conn in CONNECTIONS:
connection_id = f"{conn.id}_{db_type}"
BashOperator(
task_id=f"delete_connection_{connection_id}",
bash_command=DELETE_CONNECTION_COMMAND.format(connection_id),
trigger_rule=TriggerRule.ALL_DONE,
)
delete_connections_task = delete_connection.expand(
connection_id=[f"{conn.id}_{database_type}" for conn in CONNECTIONS]
)

teardown_task = teardown(instance=cloud_sql_instance_name, db_type=database_type)
delete_instance = CloudSQLDeleteInstanceOperator(
task_id=f"delete_cloud_sql_instance_{database_type}",
project_id=PROJECT_ID,
instance=cloud_sql_instance_name,
trigger_rule=TriggerRule.ALL_DONE,
)

(
# TEST SETUP
get_composer_network
>> create_cloud_sql_instance
create_cloud_sql_instance
>> [
create_database,
create_user_task,
Expand All @@ -556,7 +536,7 @@ def teardown(instance: str, db_type: str):
# TEST BODY
>> execute_queries_task
# TEST TEARDOWN
>> teardown_task
>> [delete_instance, delete_connections_task]
)

# ### Everything below this line is not part of example ###
Expand Down
51 changes: 18 additions & 33 deletions tests/system/providers/google/cloud/cloud_sql/example_cloud_sql_query_ssl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,7 +38,6 @@
from airflow.decorators import task
from airflow.models.connection import Connection
from airflow.models.dag import DAG
from airflow.operators.bash import BashOperator
from airflow.providers.google.cloud.hooks.cloud_sql import CloudSQLHook
from airflow.providers.google.cloud.hooks.secret_manager import GoogleCloudSecretManagerHook
from airflow.providers.google.cloud.operators.cloud_sql import (
Expand All @@ -47,27 +46,17 @@
CloudSQLDeleteInstanceOperator,
CloudSQLExecuteQueryOperator,
)
from airflow.settings import Session
from airflow.utils.trigger_rule import TriggerRule
from tests.system.providers.google import DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID

ENV_ID = os.environ.get("SYSTEM_TESTS_ENV_ID")
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT", "Not found")
PROJECT_ID = os.environ.get("SYSTEM_TESTS_GCP_PROJECT") or DEFAULT_GCP_SYSTEM_TEST_PROJECT_ID
DAG_ID = "cloudsql-query-ssl"
REGION = "us-central1"
HOME_DIR = Path.home()

COMPOSER_ENVIRONMENT = os.environ.get("COMPOSER_ENVIRONMENT", "")
if COMPOSER_ENVIRONMENT:
# We assume that the test is launched in Cloud Composer environment because the reserved environment
# variable is assigned (https://cloud.google.com/composer/docs/composer-2/set-environment-variables)
GET_COMPOSER_NETWORK_COMMAND = """
gcloud composer environments describe $COMPOSER_ENVIRONMENT \
--location=$COMPOSER_LOCATION \
--project=$GCP_PROJECT \
--format="value(config.nodeConfig.network)"
"""
else:
# The test is launched locally
GET_COMPOSER_NETWORK_COMMAND = "echo"


def run_in_composer():
Expand Down Expand Up @@ -121,7 +110,7 @@ def ip_configuration() -> dict[str, Any]:
"requireSsl": False,
"sslMode": "ENCRYPTED_ONLY",
"enablePrivatePathForGoogleCloudServices": True,
"privateNetwork": """{{ task_instance.xcom_pull('get_composer_network')}}""",
"privateNetwork": f"projects/{PROJECT_ID}/global/networks/default",
}
else:
# Use connection to Cloud SQL instance via Public IP from anywhere (mask 0.0.0.0/0).
Expand Down Expand Up @@ -273,12 +262,6 @@ def cloud_sql_database_create_body(instance: str) -> dict[str, Any]:
catchup=False,
tags=["example", "cloudsql", "postgres"],
) as dag:
get_composer_network = BashOperator(
task_id="get_composer_network",
bash_command=GET_COMPOSER_NETWORK_COMMAND,
do_xcom_push=True,
)

for db_provider in DB_PROVIDERS:
database_type: str = db_provider["database_type"]
cloud_sql_instance_name: str = db_provider["cloud_sql_instance_name"]
Expand Down Expand Up @@ -342,9 +325,9 @@ def create_connection(
connection_id: str, instance: str, db_type: str, ip_address: str, port: str
) -> str | None:
session = settings.Session()
if session.query(Connection).filter(Connection.conn_id == connection_id).first():
log.warning("Connection '%s' already exists", connection_id)
return connection_id
log.info("Removing connection %s if it exists", connection_id)
query = session.query(Connection).filter(Connection.conn_id == connection_id)
query.delete()

connection: dict[str, Any] = deepcopy(CONNECTION_PUBLIC_TCP_SSL_KWARGS)
connection["extra"]["instance"] = instance
Expand Down Expand Up @@ -472,12 +455,15 @@ def save_ssl_cert_to_secret_manager(ssl_cert: dict[str, Any], db_type: str) -> s
trigger_rule=TriggerRule.ALL_DONE,
)

delete_connection = BashOperator(
task_id=f"delete_connection_{conn_id}",
bash_command=DELETE_CONNECTION_COMMAND.format(conn_id),
trigger_rule=TriggerRule.ALL_DONE,
skip_on_exit_code=1,
)
@task(task_id=f"delete_connection_{database_type}")
def delete_connection(connection_id: str) -> None:
session = Session()
log.info("Removing connection %s", connection_id)
query = session.query(Connection).filter(Connection.conn_id == connection_id)
query.delete()
session.commit()

delete_connection_task = delete_connection(connection_id=conn_id)

@task(task_id=f"delete_secret_{database_type}")
def delete_secret(ssl_secret_id, db_type: str) -> None:
Expand All @@ -491,8 +477,7 @@ def delete_secret(ssl_secret_id, db_type: str) -> None:

(
# TEST SETUP
get_composer_network
>> create_cloud_sql_instance
create_cloud_sql_instance
>> [create_database, create_user_task, get_ip_address_task]
>> create_connection_task
>> create_ssl_certificate_task
Expand All @@ -501,7 +486,7 @@ def delete_secret(ssl_secret_id, db_type: str) -> None:
>> query_task
>> query_task_secret
# TEST TEARDOWN
>> [delete_instance, delete_connection, delete_secret_task]
>> [delete_instance, delete_connection_task, delete_secret_task]
)

# ### Everything below this line is not part of example ###
Expand Down