Docker module basic functionality
This commit is contained in:
Niels Geens
@@ -2,6 +2,7 @@ from docker.models.containers import Container
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
import docker
|
||||
import logging
|
||||
import requests
|
||||
import subprocess
|
||||
import time
|
||||
|
||||
@@ -21,10 +22,54 @@ class DockerModule:
|
||||
def __init__(self):
|
||||
DockerModule._start_docker_engine()
|
||||
self.client = docker.DockerClient(base_url='unix://var/run/docker.sock')
|
||||
self.containers: Dict[str, Container] = {}
|
||||
self.container_sockets: Dict[str, docker.APIClient.AttachSocket] = {}
|
||||
self.containers: Dict[str, (Container, docker.APIClient.AttachSocket)] = {}
|
||||
self.logger = logging.getLogger(__name__)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exec_type, exec_value, traceback):
|
||||
del self
|
||||
|
||||
def __del__(self):
|
||||
"""Cleanup all resources before shutdown."""
|
||||
try:
|
||||
# Stop and remove all containers
|
||||
for (container, socket) in self.containers.values():
|
||||
self._wait_and_stop(container, socket, 0)
|
||||
# Close the Docker client
|
||||
if hasattr(self, 'client'):
|
||||
self.client.close()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error during cleanup: {str(e)}")
|
||||
|
||||
def _wait_and_stop(self, container: Container, socket, timeout: int) -> Tuple[int, str]:
|
||||
"""Helper method for waiting on long- and short-lived containers."""
|
||||
status_code = None
|
||||
if timeout > 0:
|
||||
try:
|
||||
result = container.wait(timeout=timeout/1000)
|
||||
status_code = result.get('StatusCode')
|
||||
except requests.exceptions.ReadTimeout as e:
|
||||
container.stop()
|
||||
status_code = -1
|
||||
except requests.exceptions.ConnectionError as e:
|
||||
container.stop()
|
||||
status_code = -1
|
||||
else:
|
||||
container.stop()
|
||||
status_code = -1
|
||||
try:
|
||||
socket.close()
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error closing socket: {str(e)}")
|
||||
output = container.logs(stdout=True, stderr=True).decode('utf-8')
|
||||
try:
|
||||
container.remove(force=True)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error removing container: {str(e)}")
|
||||
return status_code, output
|
||||
|
||||
def start_container(
|
||||
self,
|
||||
image: str,
|
||||
@@ -59,100 +104,46 @@ class DockerModule:
|
||||
# Validate inputs
|
||||
if name is None and timeout < 0:
|
||||
raise ValueError("Either name or timeout must be provided")
|
||||
|
||||
# Prepare command with arguments
|
||||
cmd = []
|
||||
if command:
|
||||
cmd.append(command)
|
||||
if arguments:
|
||||
cmd.extend(arguments)
|
||||
|
||||
# Prepare volume bindings
|
||||
volume_binds = {}
|
||||
if volumes:
|
||||
for host_path, container_path in volumes.items():
|
||||
volume_binds[host_path] = {'bind': container_path, 'mode': 'rw'}
|
||||
|
||||
# Prepare port bindings
|
||||
port_bindings = {}
|
||||
if ports:
|
||||
for host_port, container_port in ports.items():
|
||||
port_bindings[container_port] = host_port
|
||||
|
||||
try:
|
||||
self.client.images.pull('busybox:latest')
|
||||
# Create container with tty and open stdin
|
||||
container = self.client.containers.create(
|
||||
image,
|
||||
command=cmd if cmd else None,
|
||||
name=name,
|
||||
detach=True,
|
||||
volumes=volume_binds,
|
||||
ports=port_bindings,
|
||||
environment=environment,
|
||||
tty=True,
|
||||
stdin_open=True
|
||||
)
|
||||
|
||||
# Attach socket before starting container
|
||||
socket = container.attach_socket(params={
|
||||
'stdin': 1,
|
||||
'stdout': 1,
|
||||
'stderr': 1,
|
||||
'stream': 1
|
||||
})
|
||||
|
||||
# Start the container
|
||||
container.start()
|
||||
|
||||
if timeout >= 0:
|
||||
# For short-lived containers, wait for completion
|
||||
try:
|
||||
container.wait(timeout=timeout/1000) # Convert ms to seconds
|
||||
output = container.logs().decode('utf-8')
|
||||
socket.close()
|
||||
container.remove(force=True)
|
||||
return output
|
||||
except docker.errors.NotFound:
|
||||
self.logger.warning(f"Container was removed before timeout")
|
||||
return ""
|
||||
else:
|
||||
# For long-running containers, store references
|
||||
self.containers[name] = container
|
||||
self.container_sockets[name] = socket
|
||||
return name
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error starting container: {str(e)}")
|
||||
raise
|
||||
|
||||
def stop_container(self, name: str) -> None:
|
||||
"""Stop and remove a Docker container by name.
|
||||
|
||||
Args:
|
||||
name: Name of the container to stop
|
||||
|
||||
Raises:
|
||||
KeyError: If container name not found
|
||||
docker.errors.APIError: If container stop/removal fails
|
||||
"""
|
||||
if name not in self.containers:
|
||||
raise KeyError(f"Container {name} not found")
|
||||
|
||||
try:
|
||||
# Close socket if it exists
|
||||
if name in self.container_sockets:
|
||||
self.container_sockets[name].close()
|
||||
del self.container_sockets[name]
|
||||
|
||||
# Stop and remove container
|
||||
container = self.containers[name]
|
||||
container.stop()
|
||||
container.remove()
|
||||
del self.containers[name]
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error stopping container {name}: {str(e)}")
|
||||
raise
|
||||
self.client.images.pull(image)
|
||||
container = self.client.containers.create(
|
||||
image,
|
||||
command=cmd if cmd else None,
|
||||
name=name,
|
||||
detach=True,
|
||||
volumes=volume_binds,
|
||||
ports=port_bindings,
|
||||
environment=environment,
|
||||
tty=False,
|
||||
stdin_open=True
|
||||
)
|
||||
socket = container.attach_socket(params={
|
||||
'stdin': 1,
|
||||
'stdout': 1,
|
||||
'stderr': 1,
|
||||
'stream': 1,
|
||||
'demux': True
|
||||
})
|
||||
container.start()
|
||||
if timeout >= 0:
|
||||
# For short-lived containers
|
||||
return self._wait_and_stop(container, socket, timeout)[1]
|
||||
else:
|
||||
# For long-running containers
|
||||
self.containers[name] = (container, socket)
|
||||
return name
|
||||
|
||||
def write_container_stdin(self, name: str, data: str) -> None:
|
||||
"""Write data to a container's standard input.
|
||||
@@ -165,11 +156,11 @@ class DockerModule:
|
||||
KeyError: If container name not found
|
||||
docker.errors.APIError: If write operation fails
|
||||
"""
|
||||
if name not in self.containers or name not in self.container_sockets:
|
||||
raise KeyError(f"Container {name} not found or socket not initialized")
|
||||
if name not in self.containers:
|
||||
raise KeyError(f"Container {name} not found")
|
||||
|
||||
try:
|
||||
socket = self.container_sockets[name]
|
||||
_, socket = self.containers[name]
|
||||
socket._sock.send(data.encode('utf-8'))
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error writing to container {name} stdin: {str(e)}")
|
||||
@@ -193,7 +184,7 @@ class DockerModule:
|
||||
raise KeyError(f"Container {name} not found")
|
||||
|
||||
try:
|
||||
container = self.containers[name]
|
||||
container, socket = self.containers[name]
|
||||
logs = container.logs(stdout=True, stderr=False, tail=n if n > 0 else 'all')
|
||||
return logs.decode('utf-8')
|
||||
except Exception as e:
|
||||
@@ -218,7 +209,7 @@ class DockerModule:
|
||||
raise KeyError(f"Container {name} not found")
|
||||
|
||||
try:
|
||||
container = self.containers[name]
|
||||
container, socket = self.containers[name]
|
||||
logs = container.logs(stdout=False, stderr=True, tail=n if n > 0 else 'all')
|
||||
return logs.decode('utf-8')
|
||||
except Exception as e:
|
||||
@@ -243,14 +234,8 @@ class DockerModule:
|
||||
if name not in self.containers:
|
||||
raise KeyError(f"Container {name} not found")
|
||||
|
||||
try:
|
||||
container = self.containers[name]
|
||||
result = container.wait(timeout=timeout/1000) # Convert ms to seconds
|
||||
logs = container.logs().decode('utf-8')
|
||||
return (result['StatusCode'], logs)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error waiting for container {name}: {str(e)}")
|
||||
raise
|
||||
container, socket = self.containers.pop(name)
|
||||
return self._wait_and_stop(container, socket, timeout)
|
||||
|
||||
def get_container_status(self, name: str) -> Dict[str, any]:
|
||||
"""Get current status information for a container.
|
||||
@@ -269,7 +254,7 @@ class DockerModule:
|
||||
raise KeyError(f"Container {name} not found")
|
||||
|
||||
try:
|
||||
container = self.containers[name]
|
||||
container, socket = self.containers[name]
|
||||
container.reload() # Refresh container info
|
||||
|
||||
return {
|
||||
|
||||
Reference in New Issue
Block a user