Move cloud_sql_proxy installation before code copy

caliban/docker/build.py

@@ -20,14 +20,11 @@
 
 from __future__ import absolute_import, division, print_function
 
-import datetime
 import json
 import os
-import re
 import subprocess
 from enum import Enum
 from pathlib import Path
-from pkg_resources import resource_filename
 from typing import Any, Dict, List, NamedTuple, NewType, Optional, Union
 
 from absl import logging
@@ -36,7 +33,6 @@
 import caliban.config as c
 import caliban.util as u
 import caliban.util.fs as ufs
-import caliban.history.types as ht
 import caliban.util.metrics as um
 
 t = Terminal()
@@ -281,62 +277,60 @@ def copy(from_path, to_path):
   return ret
 
 
-def _cloud_sql_proxy_entry(user_id: int, user_group: int) -> str:
-  '''returns dockerfile entry to fetch cloud_sql_proxy, installing in /usr/bin'''
+def _cloud_sql_proxy_entry(
+    user_id: int,
+    user_group: int,
+    caliban_config: Optional[Dict[str, Any]] = None,
+) -> str:
+  """returns dockerfile entry to fetch cloud_sql_proxy, installing in /usr/bin.
 
-  cmd = (
-      'RUN wget -q https://dl.google.com/cloudsql/cloud_sql_proxy.linux.amd64 '
-      '-O /usr/bin/cloud_sql_proxy && '
-      'chmod 755 /usr/bin/cloud_sql_proxy')
+  Args:
+  user_id: id of non-root user
+  user_group: id of non-root group for user
+  caliban_config: dictionary of caliban configuration options
 
-  return f'''
+  Returns:
+  string with Dockerfile directives to install cloud_sql_proxy as root
+  and reset the user to the specified user_id:user_group.
+
+  """
+  caliban_config = caliban_config or {}
+  mlflow_cfg = caliban_config.get('mlflow_config')
+
+  if mlflow_cfg is None:
+    return ""
+
+  return f"""
 USER root
 
-{cmd}
+RUN wget \
+  -q https://dl.google.com/cloudsql/cloud_sql_proxy.linux.amd64 \
+  -O /usr/bin/cloud_sql_proxy \
+  && chmod 755 /usr/bin/cloud_sql_proxy
 
 USER {user_id}:{user_group}
-'''
+"""
 
 
-def _generate_entrypoint(
-    executable: str,
-    user_id: int,
-    user_group: int,
-    caliban_config: Optional[Dict[str, Any]] = None,
-) -> str:
-  '''generates dockerfile entry to set the container entrypoint, along with
-    any additional necessary entries, such as for services like cloud_sql_proxy
+def _generate_entrypoint(executable: str) -> str:
+  """generates dockerfile entry to set the container entrypoint.
 
   Args:
   executable: string of main executable
-  caliban_config: dictionary of caliban configuration options
-  user_id: id of non-root user
-  user_group: id of non-root group for user
 
   Returns:
   string with Dockerfile directives to set ENTRYPOINT
-  '''
-
-  caliban_config = caliban_config or {}
-
-  mlflow_cfg = caliban_config.get('mlflow_config')
-  if mlflow_cfg is not None:
-    cloud_sql_proxy_code = _cloud_sql_proxy_entry(user_id=user_id,
-                                                  user_group=user_group)
-  else:
-    cloud_sql_proxy_code = ''
 
+  """
   launcher_cmd = json.dumps([
       'python',
       os.path.join(RESOURCE_DIR, um.LAUNCHER_SCRIPT), '--caliban_command',
       executable
   ])
 
-  return f'''
-{cloud_sql_proxy_code}
-
+  return f"""
 ENTRYPOINT {launcher_cmd}
-'''
+"""
 
 
 def _package_entries(
@@ -360,6 +354,10 @@ def _package_entries(
   arg = package.main_module or package.script_path
   package_path = package.package_path
 
+  sql_proxy_code = _cloud_sql_proxy_entry(user_id,
+                                          user_group,
+                                          caliban_config=caliban_config)
+
   copy_code = copy_command(
       user_id,
       user_group,
@@ -370,19 +368,15 @@ def _package_entries(
   # This needs to use json so that quotes print as double quotes, not single
   # quotes.
   executable_s = json.dumps(package.executable + [arg])
+  entrypoint_code = _generate_entrypoint(executable_s)
 
-  entrypoint_code = _generate_entrypoint(
-      executable=executable_s,
-      user_id=user_id,
-      user_group=user_group,
-      caliban_config=caliban_config,
-  )
+  return f"""
+  {sql_proxy_code}
 
-  return f'''
   {copy_code}
 
   {entrypoint_code}
-'''
+"""
 
 
 def _service_account_entry(user_id: int, user_group: int, credentials_path: str,