claudevdm commented on code in PR #34782:
URL: https://github.com/apache/beam/pull/34782#discussion_r2080206254
##########
sdks/python/apache_beam/yaml/integration_tests.py:
##########
@@ -76,6 +135,265 @@ def temp_bigquery_table(project, prefix='yaml_bq_it_'):
bigquery_client.client.datasets.Delete(request)
+@contextlib.contextmanager
+def temp_sqlite_database(prefix='yaml_jdbc_it_'):
+ """Context manager to provide a temporary SQLite database via JDBC for
+ testing.
+
+ This function creates a temporary SQLite database file on the local
+ filesystem. It establishes a connection using 'sqlite3', creates a predefined
+ 'tmp_table', and then yields a JDBC connection string suitable for use in
+ tests that require a generic JDBC connection (specifically configured for
+ SQLite in this case).
+
+ The SQLite database file is automatically cleaned up (closed and deleted)
+ when the context manager exits.
+
+ Args:
+ prefix (str): A prefix to use for the temporary database file name.
+
+ Yields:
+ str: A JDBC connection string for the temporary SQLite database.
+ Example format: "jdbc:sqlite:<path_to_db_file>"
+
+ Raises:
+ sqlite3.Error: If there's an error connecting to or interacting with
+ the SQLite database during setup.
+ Exception: Any other exception encountered during the setup or cleanup
+ process.
+ """
+ conn = cursor = None
+ try:
+ # Establish connection to the temp file
+ db_name = f'{prefix}{uuid.uuid4().hex}.db'
+ conn = sqlite3.connect(db_name)
+ cursor = conn.cursor()
+
+ # Create a temp table for tests
+ cursor.execute(
+ '''
+ CREATE TABLE tmp_table (
+ value INTEGER PRIMARY KEY,
+ rank INTEGER
+ )
+ ''')
+ conn.commit()
+ yield f'jdbc:sqlite:{db_name}'
+ except (sqlite3.Error, Exception) as err:
+ logging.error("Error interacting with temporary SQLite DB: %s", err)
+ raise err
+ finally:
+ # Close connections
+ if cursor:
+ cursor.close()
+ if conn:
+ conn.close()
+ try:
+ if os.path.exists(db_name):
+ os.remove(db_name)
+ except Exception as err:
+ logging.error("Error deleting temporary SQLite DB: %s", err)
+ raise err
+
+
+@contextlib.contextmanager
+def temp_mysql_database():
+ """Context manager to provide a temporary MySQL database for testing.
+
+ This function utilizes the 'testcontainers' library to spin up a
+ MySQL instance within a Docker container. It then connects
+ to this temporary database using 'mysql.connector', creates a predefined
+ 'tmp_table', and yields a JDBC connection string suitable for use in tests.
+
+ The Docker container and the database instance are automatically managed
+ and torn down when the context manager exits.
+
+ Yields:
+ str: A JDBC connection string for the temporary MySQL database.
+ Example format:
+ "jdbc:mysql://<host>:<port>/<db_name>?
Review Comment:
Theres quite a lot of repetition for setting up all the jdbc compatible
databases. Can we simplify it by using SqlAlcemy like in jdbcio tests
https://github.com/apache/beam/blob/d7fdefe20afc534a353ffda7e62e350add16e769/sdks/python/apache_beam/io/external/xlang_jdbcio_it_test.py#L139?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@beam.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
