3535"""
3636
3737import logging
38+ import re
3839import subprocess
3940import time
4041from typing import Optional , Literal , Union
@@ -100,6 +101,50 @@ def __init__(
100101 self .db_type ,
101102 )
102103
104+ def wait_for_database (
105+ self ,
106+ container_id : str ,
107+ command : str ,
108+ max_attempts : int = 10 ,
109+ sleep_time : int = 3 ,
110+ ) -> bool :
111+ """
112+ Wait for the database to be ready.
113+ Args:
114+ container_id: Container ID or name
115+ command: Command to execute to test if the database is ready
116+ max_attempts: Maximum number of attempts to wait for the database to be ready
117+ sleep_time: Time to sleep between attempts
118+ Returns:
119+ True if database is ready, False otherwise
120+ """
121+ logger .debug ("Waiting for database to be ready..." )
122+ logger .debug ("Container ID: %s" , container_id )
123+ logger .debug ("Command: %s" , command )
124+ logger .debug ("Max attempts: %s" , max_attempts )
125+ logger .debug ("Sleep time: %s" , sleep_time )
126+ for attempt in range (1 , max_attempts + 1 ):
127+ try :
128+ output = PodmanCLIWrapper .podman_exec_shell_command (
129+ cid_file_name = container_id , cmd = command , not_shell = True
130+ )
131+ if isinstance (output , bool ) and not output :
132+ logger .debug (
133+ "Database not ready, attempt %s, retrying... (output: '%s')" ,
134+ attempt ,
135+ output ,
136+ )
137+ time .sleep (sleep_time )
138+ continue
139+ if isinstance (output , str ) and output .strip () == "" :
140+ logger .info ("Database is ready (output: %s)" , output )
141+ return True
142+ except subprocess .CalledProcessError as cpe :
143+ logger .error ("Error waiting for database: %s" , cpe )
144+ time .sleep (sleep_time )
145+ logger .error ("Database not ready after %s attempts" , max_attempts )
146+ return False
147+
103148 def assert_login_success (
104149 self ,
105150 container_ip : str ,
@@ -369,6 +414,7 @@ def postgresql_cmd(
369414 extra_args : str = "" ,
370415 sql_command : Optional [str ] = None ,
371416 podman_run_command : Optional [str ] = "run --rm" ,
417+ docker_args : str = "" ,
372418 ) -> str :
373419 """
374420 Execute a PostgreSQL command against a container.
@@ -389,7 +435,8 @@ def postgresql_cmd(
389435 port: Port number (default: 5432)
390436 extra_args: Additional arguments to pass to psql command
391437 sql_command: SQL command to execute (e.g., "-c 'SELECT 1;'")
392-
438+ podman_run_command: Podman run command to use (default: "run --rm")
439+ docker_args: Docker arguments to pass to podman run command
393440 Returns:
394441 Command output as string
395442
@@ -406,11 +453,12 @@ def postgresql_cmd(
406453 container_id = self .image_name
407454 cmd_parts = [
408455 podman_run_command ,
456+ docker_args ,
409457 f"-e PGPASSWORD={ password } " ,
410- self . image_name ,
458+ container_id ,
411459 "psql" ,
412460 "-v ON_ERROR_STOP=1" ,
413- f"' { connection_string } '" ,
461+ connection_string ,
414462 ]
415463
416464 if extra_args :
@@ -550,11 +598,14 @@ def run_sql_command(
550598 port : int = 3306 ,
551599 sql_cmd : Optional [Union [list [str ], str ]] = None ,
552600 database : str = "db" ,
553- max_attempts : int = 60 ,
601+ max_attempts : int = 10 ,
554602 sleep_time : int = 3 ,
555603 container_id : Optional [str ] = None ,
604+ docker_args : Optional [Union [list [str ], str ]] = "" ,
556605 podman_run_command : Optional [str ] = "run --rm" ,
557606 ignore_error : bool = False ,
607+ expected_output : Optional [str ] = None ,
608+ use_bash : bool = False ,
558609 ) -> str | bool :
559610 """
560611 Run a database command inside the container.
@@ -575,7 +626,8 @@ def run_sql_command(
575626 sleep_time: Time to sleep between attempts (default: 3)
576627 container_id: Container ID or name
577628 podman_run_command: Podman run command to use (default: "run --rm")
578-
629+ ignore_error: Ignore error and return output (default: False)
630+ expected_output: Expected output of the command (default: None)
579631 Returns:
580632 Command output as string or False if command failed
581633 """
@@ -585,10 +637,13 @@ def run_sql_command(
585637 sql_cmd = "SELECT 1;"
586638 if isinstance (sql_cmd , str ):
587639 sql_cmd = [sql_cmd ]
640+ if isinstance (docker_args , list ):
641+ docker_args = " " .join (docker_args )
588642 logger .debug (
589643 "Podman run command: %s with image: %s" , podman_run_command , container_id
590644 )
591645 logger .debug ("Database type: %s" , self .db_type )
646+ logger .debug ("Docker arguments: %s" , docker_args )
592647 logger .debug ("SQL command: %s" , sql_cmd )
593648 logger .debug ("Database: %s" , database )
594649 logger .debug ("Username: %s" , username )
@@ -599,17 +654,34 @@ def run_sql_command(
599654 logger .debug ("Sleep time: %s" , sleep_time )
600655 return_output = None
601656 for cmd in sql_cmd :
657+ if use_bash :
658+ cmd = f'bash -c "{ cmd } "'
602659 for attempt in range (1 , max_attempts + 1 ):
603660 if self .db_type in ["postgresql" , "postgres" ]:
604- return_output = self .postgresql_cmd (
605- container_ip = container_ip ,
606- username = username ,
607- password = password ,
608- database = database ,
609- sql_command = f"-e '{ cmd } '" ,
610- container_id = container_id ,
611- podman_run_command = podman_run_command ,
612- )
661+ try :
662+ return_output = self .postgresql_cmd (
663+ container_ip = container_ip ,
664+ username = username ,
665+ password = password ,
666+ database = database ,
667+ sql_command = cmd ,
668+ container_id = container_id ,
669+ podman_run_command = podman_run_command ,
670+ docker_args = docker_args ,
671+ )
672+ logger .info ("PostgreSQL return output: %s" , return_output )
673+ except subprocess .CalledProcessError as cpe :
674+ # In case of ignore_error, we return the output
675+ # This is useful for commands that are expected to fail, like wrong login
676+ if ignore_error :
677+ return_output = cpe .output
678+ else :
679+ logger .error (
680+ "Failed to execute command, output: %s, error: %s" ,
681+ cpe .output ,
682+ cpe .stderr ,
683+ )
684+ return False
613685 else :
614686 try :
615687 return_output = self .mysql_cmd (
@@ -621,6 +693,7 @@ def run_sql_command(
621693 container_id = container_id ,
622694 podman_run_command = podman_run_command ,
623695 )
696+ logger .info ("MySQL return output: %s" , return_output )
624697 except subprocess .CalledProcessError as cpe :
625698 # In case of ignore_error, we return the output
626699 # This is useful for commands that are expected to fail, like wrong login
@@ -633,27 +706,25 @@ def run_sql_command(
633706 cpe .stderr ,
634707 )
635708 return False
636- if return_output or return_output == "" :
709+ if (
710+ return_output
711+ and expected_output
712+ and re .search (expected_output , return_output )
713+ ):
637714 logger .info ("Command executed successfully on attempt %s" , attempt )
638- # Let's break out of the loop and return the output
639715 break
716+ if return_output and not expected_output :
717+ logger .info (
718+ "Command executed successfully without checking for expected output on attempt %s"
719+ % attempt
720+ )
721+ break
722+ if attempt < max_attempts :
723+ time .sleep (sleep_time )
640724 else :
641- if attempt < max_attempts :
642- logger .debug (
643- "Attempt %s failed, output: '%s', retrying..." ,
644- attempt ,
645- return_output ,
646- )
647- time .sleep (sleep_time )
648- else :
649- logger .error (
650- "Failed to execute command after %s attempts, output: %s" ,
651- max_attempts ,
652- return_output ,
653- )
654- return False
725+ return False
655726 if return_output :
656727 logger .info ("All commands executed successfully" )
657- logger .debug ("Output:\n %s " , return_output )
728+ logger .debug ("Output:\n '%s' " , return_output )
658729 return return_output
659730 return False
0 commit comments