Set up db connections to use SOCKS proxy in Development environment

Pipfile

@@ -30,6 +30,7 @@ attrs = "==18.2.0"
 pytest-cov = "*"
 supervisor = "~=4.2"
 yoyo-migrations = "==7.3.2"
+pysocks = "*"
 
 # Versions required for dependabot
 sqlparse = "~=0.5.0"

wp1/credentials.py.example

@@ -29,18 +29,18 @@ CONF_LANG = 'en'
 CREDENTIALS = {
     Environment.DEVELOPMENT: {
         # Database credentials for the wikipedia replica database.
-        # For development, you can use your toolforge credentials and connect
-        # to a live replica via an SSH tunnel, as outlined here:
-        # https://wikitech.wikimedia.org/wiki/Help:Toolforge/Database#Connecting_to_the_database_replicas_from_your_own_computer
-        # The 'local port' you use, 4711 in the example, will correspond to the
-        # port you set below.
+        # For development, start a SOCKS5 proxy connection to login.toolforge.org
+        # with the command:
+        # $ ssh -D 1080 login.toolforge.org
+        # (This is assuming you have set up your SSH credentials for Toolforge)
+        # Database traffic will be tunneled through this proxy, so that URLs
+        # inside of *.eqiad.wmflabs can be resolved.
         'WIKIDB': {
-            'user': 'yourtoolforgeuser', # EDIT this line
-            'password': 'yourtoolforgepass', # EDIT this line
-            'host': 'localhost',
-            'port': 4711,
-            'db': 'enwiki_p',
-        },
+         'user': 'someuser', # EDIT this line for production.
+         'password': 'somepass', # EDIT this line for production.
+         'host': 'enwiki.analytics.db.svc.eqiad.wmflabs',
+         'db': 'enwiki_p',
+       },
 
         # Database credentials for the enwp10 project/application database.
         # For development, use the docker-compose-dev.yml file and spin up a

wp1/db.py

@@ -5,6 +5,7 @@
 import pymysql
 import pymysql.cursors
 import pymysql.err
+import socks
 
 logger = logging.getLogger(__name__)
 
@@ -19,22 +20,35 @@
 RETRY_TIME_SECONDS = 5
 
 
-def connect(db_name):
+def connect(db_name, **overrides):
   creds = CREDENTIALS[ENV].get(db_name)
   if creds is None:
-    raise ValueError('db credentials for %r in ENV=%s are None')
+    raise ValueError('db credentials for %r in ENV=%s are None' %
+                     (db_name, ENV))
 
   kwargs = {
       'charset': None,
       'use_unicode': False,
       'cursorclass': pymysql.cursors.SSDictCursor,
-      **creds
+      **creds,
+      **overrides
   }
 
   tries = 4
   while True:
     try:
-      return pymysql.connect(**kwargs)
+      if db_name == 'WIKIDB' and ENV == ENV.DEVELOPMENT:
+        # In development, connect through a SOCKS5 proxy so that hosts on
+        # *.eqiad.wmflabs can be reached.
+        s = socks.socksocket()
+        s.set_proxy(socks.SOCKS5, 'localhost')
+        s.connect((kwargs['host'], kwargs.get('port', 3306)))
+        conn = pymysql.connect(**kwargs, defer_connect=True)
+        conn.connect(sock=s)
+      else:
+        conn = pymysql.connect(**kwargs)
+
+      return conn
     except pymysql.err.InternalError:
       if tries > 0:
         logging.warning('Could not connect to database, retrying in %s seconds',

wp1/db_test.py

@@ -1,7 +1,8 @@
 import unittest
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pymysql.err
+import socks
 
 from wp1.environment import Environment
 
@@ -30,3 +31,31 @@ def test_retries_four_times_failure(self, patched_sleep, patched_pymysql):
     with self.assertRaises(pymysql.err.InternalError):
       connect('WP10DB')
     self.assertEqual(5, patched_pymysql.call_count)
+
+  @patch('wp1.db.pymysql.connect')
+  @patch('wp1.db.socks.socksocket')
+  @patch('wp1.db.ENV', Environment.DEVELOPMENT)
+  def test_socks_proxy(self, mock_socket, mock_connect):
+    from wp1.db import connect
+
+    socket = MagicMock()
+    mock_socket.return_value = socket
+    conn = MagicMock()
+    mock_connect.return_value = conn
+
+    connect('WIKIDB', host='foo.wikimedia.cloud', port=6000)
+
+    socket.set_proxy.assert_called_once_with(socks.SOCKS5, 'localhost')
+    socket.connect.assert_called_once_with(('foo.wikimedia.cloud', 6000))
+    defer = mock_connect.call_args.kwargs.get('defer_connect')
+    self.assertTrue(defer)
+    conn.connect.assert_called_once_with(sock=socket)
+
+  @patch('wp1.db.pymysql.connect')
+  @patch('wp1.db.socks.socksocket')
+  @patch('wp1.db.ENV', Environment.DEVELOPMENT)
+  def test_socks_proxy_not_used(self, mock_socket, mock_connect):
+    from wp1.db import connect
+    connect('WP10DB')
+    mock_socket.assert_not_called()
+    mock_connect.assert_called_once()

wp1/wiki_db.py

@@ -1,5 +1,5 @@
 from functools import partial
 
-from wp1.db import connect
+from wp1.db import connect as base_connect
 
-connect = partial(connect, 'WIKIDB')
+connect = partial(base_connect, 'WIKIDB')

wp1/wikilang_db.py

@@ -0,0 +1,10 @@
+from functools import partial
+
+from wp1.db import connect as db_connect
+
+
+def connect(lang):
+  wiki = f'{lang}wiki'
+  db = f'{wiki}_p'
+  host = f'{wiki}.analytics.db.svc.eqiad.wmflabs'
+  return db_connect('WIKIDB', host=host, db=db)

wp1/wikilang_db_test.py

@@ -0,0 +1,14 @@
+import unittest
+from unittest.mock import patch
+
+import wp1.wikilang_db
+
+
+class WikiLangDbTest(unittest.TestCase):
+
+  @patch('wp1.wikilang_db.db_connect')
+  def test_connect(self, mock_connect):
+    wp1.wikilang_db.connect('fr')
+
+    mock_connect.assert_called_once_with(
+        'WIKIDB', host='frwiki.analytics.db.svc.eqiad.wmflabs', db='frwiki_p')