def get_connections(cls, conn_id: str) -> List[Connection]:
"""
Get all connections as an iterable.
:param conn_id: connection id
:return: array of connections
"""
return secrets.get_connections(conn_id)
def _record_statistics(ts, postgres_conn_id, ssh_conn_id, **context):
# Can't use PostgresHook because our TCP port is likely to be dynamic
# because we are connecting through an SSH Tunnel
(pg_secret, ) = secrets.get_connections(postgres_conn_id)
(gh_token, ) = secrets.get_connections("github_metrics_token")
ssh_conn = SSHHook(ssh_conn_id=ssh_conn_id)
tunnel = ssh_conn.get_tunnel(remote_port=pg_secret.port,
remote_host=pg_secret.host)
tunnel.start()
with tunnel:
log(f"Connected SSH Tunnel: {tunnel}")
# Airflow conflates dbname and schema, even though they are very different in PG
constr = (
f"host=localhost user={pg_secret.login} dbname={pg_secret.schema} "
f"port={tunnel.local_bind_port} password={pg_secret.password}")
# It's important to wrap this connection in a try/finally block, otherwise
# we can cause a deadlock with the SSHTunnel
conn = psycopg2.connect(constr)
try:
log(f"Connected to Postgres: {conn}")
cur = conn.cursor()
# ensure_pg_table(cur)
stats_retriever = GitHubStatsRetriever(gh_token.password)
for owner, repo in TRACKED_REPOS:
log(f"Recording repo stats for {owner}/{repo}")
stats = stats_retriever.get_repo_stats(owner, repo)
log(stats)
save_record_to_pg(cur, ts, f"{owner}/{repo}", stats)
conn.commit()
log("Transaction committed")
cur.close()
finally:
conn.close()
def test_backend_fallback_to_env_var(self, mock_get_uri):
mock_get_uri.return_value = None
backends = ensure_secrets_loaded()
backend_classes = [backend.__class__.__name__ for backend in backends]
self.assertIn('AwsSsmSecretsBackend', backend_classes)
uri = get_connections(conn_id="test_mysql")
# Assert that AwsSsmSecretsBackend.get_conn_uri was called
mock_get_uri.assert_called_once_with(conn_id='test_mysql')
self.assertEqual('mysql://*****:*****@host:5432/airflow', uri[0].get_uri())
def test_get_conn_uri_non_existent_key(self):
"""
Test that if the key with connection ID is not present in SSM,
AwsSsmSecretsBackend.get_connections should return None and fallback to the
environment variable if it is set
"""
conn_id = "test_mysql"
test_client = AwsSsmSecretsBackend()
self.assertIsNone(test_client.get_conn_uri(conn_id=conn_id))
self.assertEqual([], test_client.get_connections(conn_id=conn_id))
# Fallback to connection defined in Environment Variable
self.assertEqual(
"mysql://*****:*****@host:5432/airflow",
get_connections(conn_id="test_mysql")[0].get_uri())
def test_get_connections_first_try(self, mock_env_get, mock_meta_get):
mock_env_get.side_effect = [["something"]] # returns nonempty list
get_connections("fake_conn_id")
mock_env_get.assert_called_once_with(conn_id="fake_conn_id")
mock_meta_get.not_called()
def test_get_connections_second_try(self, mock_env_get, mock_meta_get):
mock_env_get.side_effect = [[]] # return empty list
get_connections("fake_conn_id")
mock_meta_get.assert_called_once_with(conn_id="fake_conn_id")
mock_env_get.assert_called_once_with(conn_id="fake_conn_id")
