fix: Clean up Python sample at cloud-sql/mysql/client-side-encryption (GoogleCloudPlatform#10028)

m-strzelczyk · web-flow · commit 7222e36a7532 · 2023-05-25T12:54:12.000Z
## Description Fixes b/280880978 Note: Before submitting a pull request, please open an issue for discussion if you are not associated with Google. ## Checklist - [ ] I have followed [Sample Guidelines from AUTHORING_GUIDE.MD](https://togithub.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md) - [ ] README is updated to include [all relevant information](https://togithub.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#readme-file) - [ ] **Tests** pass: `nox -s py-3.9` (see [Test Environment Setup](https://togithub.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#test-environment-setup)) - [ ] **Lint** pass: `nox -s lint` (see [Test Environment Setup](https://togithub.com/GoogleCloudPlatform/python-docs-samples/blob/main/AUTHORING_GUIDE.md#test-environment-setup)) - [ ] These samples need a new **API enabled** in testing projects to pass (let us know which ones) - [ ] These samples need a new/updated **env vars** in testing projects set to pass (let us know which ones) - [ ] This sample adds a new sample directory, and I updated the [CODEOWNERS file](https://togithub.com/GoogleCloudPlatform/python-docs-samples/blob/main/.github/CODEOWNERS) with the codeowners for this sample - [ ] This sample adds a new **Product API**, and I updated the [Blunderbuss issue/PR auto-assigner](https://togithub.com/GoogleCloudPlatform/python-docs-samples/blob/main/.github/blunderbuss.yml) with the codeowners for this sample - [ ] Please **merge** this PR for me once it is approved
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/cloud_kms_env_aead.py b/cloud-sql/mysql/client-side-encryption/snippets/cloud_kms_env_aead.py
@@ -22,9 +22,10 @@
 logger = logging.getLogger(__name__)
 
 
-def init_tink_env_aead(
-        key_uri: str,
-        credentials: str) -> tink.aead.KmsEnvelopeAead:
+def init_tink_env_aead(key_uri: str, credentials: str) -> tink.aead.KmsEnvelopeAead:
+    """
+    Initiates the Envelope AEAD object using the KMS credentials.
+    """
     aead.register()
 
     try:
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/cloud_kms_env_aead_test.py b/cloud-sql/mysql/client-side-encryption/snippets/cloud_kms_env_aead_test.py
@@ -12,23 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2021 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import os
 
 import pytest
+from tink.aead import KmsEnvelopeAead
 
 from snippets.cloud_kms_env_aead import init_tink_env_aead
 
@@ -40,12 +27,11 @@ def setup() -> str:
     yield kms_uri
 
 
-def test_cloud_kms_env_aead(
-        capsys: pytest.CaptureFixture, kms_uri: str) -> None:
+def test_cloud_kms_env_aead(kms_uri: str) -> None:
     credentials = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", "")
 
     # Create env_aead primitive
-    init_tink_env_aead(kms_uri, credentials)
+    envelope = init_tink_env_aead(kms_uri, credentials)
 
-    captured = capsys.readouterr().out
-    assert f"Created envelope AEAD Primitive using KMS URI: {kms_uri}" in captured
+    assert isinstance(envelope, KmsEnvelopeAead)
+    assert envelope.key_template == kms_uri
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/cloud_sql_connection_pool.py b/cloud-sql/mysql/client-side-encryption/snippets/cloud_sql_connection_pool.py
@@ -19,6 +19,9 @@
 def init_tcp_connection_engine(
     db_user: str, db_pass: str, db_name: str, db_host: str
 ) -> sqlalchemy.engine.base.Engine:
+    """
+    Creates a connection to the database using tcp socket.
+    """
     # Remember - storing secrets in plaintext is potentially unsafe. Consider using
     # something like https://cloud.google.com/secret-manager/docs/overview to help keep
     # secrets secret.
@@ -50,6 +53,9 @@ def init_unix_connection_engine(
     instance_connection_name: str,
     db_socket_dir: str,
 ) -> sqlalchemy.engine.base.Engine:
+    """
+    Creates a connection to the database using unix socket.
+    """
     # Remember - storing secrets in plaintext is potentially unsafe. Consider using
     # something like https://cloud.google.com/secret-manager/docs/overview to help keep
     # secrets secret.
@@ -78,7 +84,7 @@ def init_db(
     db_socket_dir: str = None,
     db_host: str = None,
 ) -> sqlalchemy.engine.base.Engine:
-
+    """Starts a connection to the database and creates voting table if it doesn't exist."""
     if db_host:
         db = init_tcp_connection_engine(db_user, db_pass, db_name, db_host)
     else:
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/cloud_sql_connection_pool_test.py b/cloud-sql/mysql/client-side-encryption/snippets/cloud_sql_connection_pool_test.py
@@ -12,32 +12,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2021 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 from __future__ import annotations
 
 import os
 import uuid
 
 import pytest
+import sqlalchemy
 
-from snippets.cloud_sql_connection_pool import (
-    init_db,
-    init_tcp_connection_engine,
-    init_unix_connection_engine
-)
+from snippets.cloud_sql_connection_pool import (init_db,
+                                                init_tcp_connection_engine,
+                                                init_unix_connection_engine)
 
 
 @pytest.fixture(name="conn_vars")
@@ -60,50 +45,46 @@ def setup() -> dict[str, str]:
         yield conn_vars
 
 
-def test_init_tcp_connection_engine(
-        capsys: pytest.CaptureFixture,
-        conn_vars: dict[str, str]) -> None:
-
-    init_tcp_connection_engine(
+def test_init_tcp_connection_engine(conn_vars: dict[str, str]) -> None:
+    engine = init_tcp_connection_engine(
         db_user=conn_vars["db_user"],
         db_name=conn_vars["db_name"],
         db_pass=conn_vars["db_pass"],
         db_host=conn_vars["db_host"],
     )
 
-    captured = capsys.readouterr().out
-    assert "Created TCP connection pool" in captured
+    assert isinstance(engine, sqlalchemy.engine.base.Engine)
+    assert conn_vars["db_name"] in engine.url
 
 
-def test_init_unix_connection_engine(
-        capsys: pytest.CaptureFixture,
-        conn_vars: dict[str, str]) -> None:
-
-    init_unix_connection_engine(
+def test_init_unix_connection_engine(conn_vars: dict[str, str]) -> None:
+    engine = init_unix_connection_engine(
         db_user=conn_vars["db_user"],
         db_name=conn_vars["db_name"],
         db_pass=conn_vars["db_pass"],
         instance_connection_name=conn_vars["instance_conn_name"],
         db_socket_dir=conn_vars["db_socket_dir"],
     )
 
-    captured = capsys.readouterr().out
-    assert "Created Unix socket connection pool" in captured
-
+    assert isinstance(engine, sqlalchemy.engine.base.Engine)
+    assert conn_vars["db_name"] in engine.url
 
-def test_init_db(
-        capsys: pytest.CaptureFixture,
-        conn_vars: dict[str, str]) -> None:
 
+def test_init_db(conn_vars: dict[str, str]) -> None:
     table_name = f"votes_{uuid.uuid4().hex}"
 
-    init_db(
+    engine = init_db(
         db_user=conn_vars["db_user"],
         db_name=conn_vars["db_name"],
         db_pass=conn_vars["db_pass"],
         table_name=table_name,
         db_host=conn_vars["db_host"],
     )
 
-    captured = capsys.readouterr().out
-    assert f"Created table {table_name} in db {conn_vars['db_name']}" in captured
+    assert isinstance(engine, sqlalchemy.engine.base.Engine)
+
+    try:
+        with engine.connect() as conn:
+            conn.execute(f"SELECT count(*) FROM {table_name}").all()
+    except Exception as error:
+        pytest.fail(f"Database wasn't initialized properly: {error}")
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/encrypt_and_insert_data.py b/cloud-sql/mysql/client-side-encryption/snippets/encrypt_and_insert_data.py
@@ -23,11 +23,13 @@
 from .cloud_kms_env_aead import init_tink_env_aead
 from .cloud_sql_connection_pool import init_db
 
-
 logger = logging.getLogger(__name__)
 
 
 def main() -> None:
+    """
+    Connects to the database, encrypts and inserts some data.
+    """
     db_user = os.environ["DB_USER"]  # e.g. "root", "mysql"
     db_pass = os.environ["DB_PASS"]  # e.g. "mysupersecretpassword"
     db_name = os.environ["DB_NAME"]  # e.g. "votes_db"
@@ -73,6 +75,10 @@ def encrypt_and_insert_data(
     team: str,
     email: str,
 ) -> None:
+    """
+    Inserts a vote into the database with email address previously encrypted using
+    a KmsEnvelopeAead object.
+    """
     time_cast = datetime.datetime.now(tz=datetime.timezone.utc)
     # Use the envelope AEAD primitive to encrypt the email, using the team name as
     # associated data. Encryption with associated data ensures authenticity
@@ -93,11 +99,7 @@ def encrypt_and_insert_data(
     # Using a with statement ensures that the connection is always released
     # back into the pool at the end of statement (even if an error occurs)
     with db.connect() as conn:
-        conn.execute(
-            stmt,
-            time_cast=time_cast,
-            team=team,
-            voter_email=encrypted_email)
+        conn.execute(stmt, time_cast=time_cast, team=team, voter_email=encrypted_email)
     print(f"Vote successfully cast for '{team}' at time {time_cast}!")
 
 
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/encrypt_and_insert_data_test.py b/cloud-sql/mysql/client-side-encryption/snippets/encrypt_and_insert_data_test.py
@@ -23,7 +23,6 @@
 from snippets.cloud_sql_connection_pool import init_db
 from snippets.encrypt_and_insert_data import encrypt_and_insert_data
 
-
 table_name = f"votes_{uuid.uuid4().hex}"
 
 
@@ -65,17 +64,10 @@ def setup_key() -> tink.aead.KmsEnvelopeAead:
 
 
 def test_encrypt_and_insert_data(
-    capsys: pytest.CaptureFixture,
     pool: sqlalchemy.engine.Engine,
-    env_aead: tink.aead.KmsEnvelopeAead
+    env_aead: tink.aead.KmsEnvelopeAead,
 ) -> None:
-    encrypt_and_insert_data(
-        pool,
-        env_aead,
-        table_name,
-        "SPACES",
-        "hello@example.com")
-    captured = capsys.readouterr()
+    encrypt_and_insert_data(pool, env_aead, table_name, "SPACES", "hello@example.com")
 
     decrypted_emails = []
     with pool.connect() as conn:
@@ -89,5 +81,4 @@ def test_encrypt_and_insert_data(
             email = env_aead.decrypt(row[2], team.encode()).decode()
             decrypted_emails.append(email)
 
-    assert "Vote successfully cast for 'SPACES'" in captured.out
     assert "hello@example.com" in decrypted_emails
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/query_and_decrypt_data.py b/cloud-sql/mysql/client-side-encryption/snippets/query_and_decrypt_data.py
@@ -24,6 +24,9 @@
 
 
 def main() -> None:
+    """
+    Connects to the database, inserts encrypted data and retrieves encrypted data.
+    """
     db_user = os.environ["DB_USER"]  # e.g. "root", "mysql"
     db_pass = os.environ["DB_PASS"]  # e.g. "mysupersecretpassword"
     db_name = os.environ["DB_NAME"]  # e.g. "votes_db"
@@ -67,7 +70,10 @@ def query_and_decrypt_data(
     db: sqlalchemy.engine.base.Engine,
     env_aead: tink.aead.KmsEnvelopeAead,
     table_name: str,
-) -> None:
+) -> list[tuple[str]]:
+    """
+    Retrieves data from the database and decrypts it using the KmsEnvelopeAead object.
+    """
     with db.connect() as conn:
         # Execute the query and fetch all results
         recent_votes = conn.execute(
@@ -76,6 +82,7 @@ def query_and_decrypt_data(
         ).fetchall()
 
         print("Team\tEmail\tTime Cast")
+        output = []
 
         for row in recent_votes:
             team = row[0]
@@ -88,6 +95,8 @@ def query_and_decrypt_data(
 
             # Print recent votes
             print(f"{team}\t{email}\t{time_cast}")
+            output.append((team, email, time_cast))
+    return output
 
 
 # [END cloud_sql_mysql_cse_query]
diff --git a/cloud-sql/mysql/client-side-encryption/snippets/query_and_decrypt_data_test.py b/cloud-sql/mysql/client-side-encryption/snippets/query_and_decrypt_data_test.py
@@ -12,20 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Copyright 2021 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
 import os
 import uuid
 
@@ -79,21 +65,16 @@ def setup_key() -> tink.aead.KmsEnvelopeAead:
 
 
 def test_query_and_decrypt_data(
-    capsys: pytest.CaptureFixture,
     pool: sqlalchemy.engine.Engine,
-    env_aead: tink.aead.KmsEnvelopeAead
+    env_aead: tink.aead.KmsEnvelopeAead,
 ) -> None:
-
     # Insert data into table before testing
-    encrypt_and_insert_data(
-        pool,
-        env_aead,
-        table_name,
-        "SPACES",
-        "hello@example.com")
-
-    query_and_decrypt_data(pool, env_aead, table_name)
-
-    captured = capsys.readouterr()
-    assert "Team\tEmail\tTime Cast" in captured.out
-    assert "hello@example.com" in captured.out
+    encrypt_and_insert_data(pool, env_aead, table_name, "SPACES", "hello@example.com")
+
+    output = query_and_decrypt_data(pool, env_aead, table_name)
+
+    for row in output:
+        if row[1] == "hello@example.com":
+            break
+    else:
+        pytest.fail("Failed to find vote in the decrypted data.")
