Files

289 lines
24 KiB
HTML

<!DOCTYPE html>
<html class="writer-html5" lang="en" data-content_root="../../">
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
<title>clickhouse_pool.pool &mdash; clickhouse-pool 0.6.1 documentation</title>
<link rel="stylesheet" type="text/css" href="../../_static/pygments.css?v=80d5e7a1" />
<link rel="stylesheet" type="text/css" href="../../_static/css/theme.css?v=19f00094" />
<!--[if lt IE 9]>
<script src="../../_static/js/html5shiv.min.js"></script>
<![endif]-->
<script src="../../_static/documentation_options.js?v=2bd037f0"></script>
<script src="../../_static/doctools.js?v=9a2dae69"></script>
<script src="../../_static/sphinx_highlight.js?v=dc90522c"></script>
<script src="../../_static/js/theme.js"></script>
<link rel="index" title="Index" href="../../genindex.html" />
<link rel="search" title="Search" href="../../search.html" />
</head>
<body class="wy-body-for-nav">
<div class="wy-grid-for-nav">
<nav data-toggle="wy-nav-shift" class="wy-nav-side">
<div class="wy-side-scroll">
<div class="wy-side-nav-search" >
<a href="../../index.html" class="icon icon-home">
clickhouse-pool
</a>
<div role="search">
<form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
<input type="text" name="q" placeholder="Search docs" aria-label="Search docs" />
<input type="hidden" name="check_keywords" value="yes" />
<input type="hidden" name="area" value="default" />
</form>
</div>
</div><div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="Navigation menu">
<p class="caption" role="heading"><span class="caption-text">Contents:</span></p>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../introduction.html">Introduction</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../api.html">API</a></li>
</ul>
</div>
</div>
</nav>
<section data-toggle="wy-nav-shift" class="wy-nav-content-wrap"><nav class="wy-nav-top" aria-label="Mobile navigation menu" >
<i data-toggle="wy-nav-top" class="fa fa-bars"></i>
<a href="../../index.html">clickhouse-pool</a>
</nav>
<div class="wy-nav-content">
<div class="rst-content">
<div role="navigation" aria-label="Page navigation">
<ul class="wy-breadcrumbs">
<li><a href="../../index.html" class="icon icon-home" aria-label="Home"></a></li>
<li class="breadcrumb-item"><a href="../index.html">Module code</a></li>
<li class="breadcrumb-item active">clickhouse_pool.pool</li>
<li class="wy-breadcrumbs-aside">
</li>
</ul>
<hr/>
</div>
<div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
<div itemprop="articleBody">
<h1>Source code for clickhouse_pool.pool</h1><div class="highlight"><pre>
<span></span><span class="sd">&quot;&quot;&quot;Connection pool for clickhouse_driver</span>
<span class="sd">Heavily inspired by psycopg2/lib/pool.py, this module implements a thread-safe</span>
<span class="sd">connection pool. Each connection is an instance of a clickhouse_driver client.</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="kn">import</span> <span class="nn">threading</span>
<span class="kn">from</span> <span class="nn">contextlib</span> <span class="kn">import</span> <span class="n">contextmanager</span>
<span class="kn">from</span> <span class="nn">clickhouse_driver</span> <span class="kn">import</span> <span class="n">Client</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Generator</span>
<span class="k">class</span> <span class="nc">ChPoolError</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A generic exception that may be raised by ChPool&quot;&quot;&quot;</span>
<span class="k">class</span> <span class="nc">TooManyConnections</span><span class="p">(</span><span class="n">ChPoolError</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Raised when attempting to use more than connections_max clients.&quot;&quot;&quot;</span>
<div class="viewcode-block" id="ChPool">
<a class="viewcode-back" href="../../api.html#clickhouse_pool.ChPool">[docs]</a>
<span class="k">class</span> <span class="nc">ChPool</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A ChPool is a pool of connections (Clients) to a ClickhouseServer.</span>
<span class="sd"> Attributes:</span>
<span class="sd"> connections_min (int): minimum number of connections to keep open</span>
<span class="sd"> connections_max (int): maximum number of connections allowed open</span>
<span class="sd"> closed (bool): if closed the pool has no connections and cannot be used</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="c1"># pylint: disable=too-many-instance-attributes</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Initialize the pool of clickhouse clients.</span>
<span class="sd"> Args:</span>
<span class="sd"> **kwargs: similar to clickhouse-driver, all settings available are</span>
<span class="sd"> documented `here</span>
<span class="sd"> &lt;https://clickhouse.tech/docs/en/single/#settings&gt;`_.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="bp">self</span><span class="o">.</span><span class="n">connections_min</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;connections_min&quot;</span><span class="p">,</span> <span class="mi">10</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">connections_max</span> <span class="o">=</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;connections_max&quot;</span><span class="p">,</span> <span class="mi">20</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">connection_args</span> <span class="o">=</span> <span class="p">{</span><span class="s2">&quot;host&quot;</span><span class="p">:</span> <span class="n">kwargs</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;host&quot;</span><span class="p">,</span> <span class="s2">&quot;localhost&quot;</span><span class="p">),</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">}</span>
<span class="bp">self</span><span class="o">.</span><span class="n">closed</span> <span class="o">=</span> <span class="kc">False</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_pool</span> <span class="o">=</span> <span class="p">[]</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_used</span> <span class="o">=</span> <span class="p">{}</span>
<span class="c1"># similar to psycopg2 pools, _rused is used for mapping instances of conn</span>
<span class="c1"># to their respective keys in _used</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_rused</span> <span class="o">=</span> <span class="p">{}</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_keys</span> <span class="o">=</span> <span class="mi">0</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_lock</span> <span class="o">=</span> <span class="n">threading</span><span class="o">.</span><span class="n">Lock</span><span class="p">()</span>
<span class="k">for</span> <span class="n">_</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">connections_min</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_connect</span><span class="p">()</span>
<span class="k">def</span> <span class="nf">_connect</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Client</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Create a new client and assign to a key.&quot;&quot;&quot;</span>
<span class="n">client</span> <span class="o">=</span> <span class="n">Client</span><span class="p">(</span><span class="o">**</span><span class="bp">self</span><span class="o">.</span><span class="n">connection_args</span><span class="p">)</span>
<span class="k">if</span> <span class="n">key</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">client</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_rused</span><span class="p">[</span><span class="nb">id</span><span class="p">(</span><span class="n">client</span><span class="p">)]</span> <span class="o">=</span> <span class="n">key</span>
<span class="k">else</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_pool</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">client</span><span class="p">)</span>
<span class="k">return</span> <span class="n">client</span>
<span class="k">def</span> <span class="nf">_get_key</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Get an unused key.&quot;&quot;&quot;</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_keys</span> <span class="o">+=</span> <span class="mi">1</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_keys</span>
<div class="viewcode-block" id="ChPool.pull">
<a class="viewcode-back" href="../../api.html#clickhouse_pool.ChPool.pull">[docs]</a>
<span class="k">def</span> <span class="nf">pull</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Client</span><span class="p">:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Get an available client from the pool.</span>
<span class="sd"> Args:</span>
<span class="sd"> key: If known, the key of the client you would like.</span>
<span class="sd"> Returns:</span>
<span class="sd"> A clickhouse-driver client.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_lock</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
<span class="k">try</span><span class="p">:</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">closed</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">ChPoolError</span><span class="p">(</span><span class="s2">&quot;pool closed&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="n">key</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
<span class="n">key</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_key</span><span class="p">()</span>
<span class="k">if</span> <span class="n">key</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="p">:</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="p">[</span><span class="n">key</span><span class="p">]</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_pool</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="p">[</span><span class="n">key</span><span class="p">]</span> <span class="o">=</span> <span class="n">client</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_pool</span><span class="o">.</span><span class="n">pop</span><span class="p">()</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_rused</span><span class="p">[</span><span class="nb">id</span><span class="p">(</span><span class="n">client</span><span class="p">)]</span> <span class="o">=</span> <span class="n">key</span>
<span class="k">return</span> <span class="n">client</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="p">)</span> <span class="o">&gt;=</span> <span class="bp">self</span><span class="o">.</span><span class="n">connections_max</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">TooManyConnections</span><span class="p">(</span><span class="s2">&quot;too many connections&quot;</span><span class="p">)</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_connect</span><span class="p">(</span><span class="n">key</span><span class="p">)</span>
<span class="k">finally</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_lock</span><span class="o">.</span><span class="n">release</span><span class="p">()</span></div>
<div class="viewcode-block" id="ChPool.push">
<a class="viewcode-back" href="../../api.html#clickhouse_pool.ChPool.push">[docs]</a>
<span class="k">def</span> <span class="nf">push</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">client</span><span class="p">:</span> <span class="n">Client</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">key</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">,</span> <span class="n">close</span><span class="p">:</span> <span class="nb">bool</span> <span class="o">=</span> <span class="kc">False</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Return a client to the pool for reuse.</span>
<span class="sd"> Args:</span>
<span class="sd"> client: The client to return.</span>
<span class="sd"> key: If known, the key of the client.</span>
<span class="sd"> close: Close the client instead of adding back to pool.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_lock</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
<span class="k">try</span><span class="p">:</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">closed</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">ChPoolError</span><span class="p">(</span><span class="s2">&quot;pool closed&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="n">key</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
<span class="n">key</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_rused</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="nb">id</span><span class="p">(</span><span class="n">client</span><span class="p">))</span>
<span class="k">if</span> <span class="n">key</span> <span class="ow">is</span> <span class="kc">None</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">ChPoolError</span><span class="p">(</span><span class="s2">&quot;trying to put unkeyed client&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_pool</span><span class="p">)</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">connections_min</span> <span class="ow">and</span> <span class="ow">not</span> <span class="n">close</span><span class="p">:</span>
<span class="c1"># TODO: verify connection still valid</span>
<span class="k">if</span> <span class="n">client</span><span class="o">.</span><span class="n">connection</span><span class="o">.</span><span class="n">connected</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_pool</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">client</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">client</span><span class="o">.</span><span class="n">disconnect</span><span class="p">()</span>
<span class="c1"># ensure thread doesn&#39;t put connection back once the pool is closed</span>
<span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">closed</span> <span class="ow">or</span> <span class="n">key</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="p">:</span>
<span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="p">[</span><span class="n">key</span><span class="p">]</span>
<span class="k">del</span> <span class="bp">self</span><span class="o">.</span><span class="n">_rused</span><span class="p">[</span><span class="nb">id</span><span class="p">(</span><span class="n">client</span><span class="p">)]</span>
<span class="k">finally</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_lock</span><span class="o">.</span><span class="n">release</span><span class="p">()</span></div>
<div class="viewcode-block" id="ChPool.cleanup">
<a class="viewcode-back" href="../../api.html#clickhouse_pool.ChPool.cleanup">[docs]</a>
<span class="k">def</span> <span class="nf">cleanup</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Close all open connections in the pool.</span>
<span class="sd"> This method loops through eveery client and calls disconnect.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_lock</span><span class="o">.</span><span class="n">acquire</span><span class="p">()</span>
<span class="k">try</span><span class="p">:</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">closed</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">ChPoolError</span><span class="p">(</span><span class="s2">&quot;pool closed&quot;</span><span class="p">)</span>
<span class="k">for</span> <span class="n">client</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">_pool</span> <span class="o">+</span> <span class="nb">list</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">_used</span><span class="o">.</span><span class="n">values</span><span class="p">()):</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">client</span><span class="o">.</span><span class="n">disconnect</span><span class="p">()</span>
<span class="c1"># TODO: handle problems with disconnect</span>
<span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span>
<span class="k">pass</span>
<span class="bp">self</span><span class="o">.</span><span class="n">closed</span> <span class="o">=</span> <span class="kc">True</span>
<span class="k">finally</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_lock</span><span class="o">.</span><span class="n">release</span><span class="p">()</span></div>
<div class="viewcode-block" id="ChPool.get_client">
<a class="viewcode-back" href="../../api.html#clickhouse_pool.ChPool.get_client">[docs]</a>
<span class="nd">@contextmanager</span>
<span class="k">def</span> <span class="nf">get_client</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">key</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="kc">None</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="n">Generator</span><span class="p">[</span><span class="n">Client</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">]:</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A clean way to grab a client via a contextmanager.</span>
<span class="sd"> Args:</span>
<span class="sd"> key: If known, the key of the client to grab.</span>
<span class="sd"> Yields:</span>
<span class="sd"> Client: a clickhouse-driver client</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="n">client</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">pull</span><span class="p">(</span><span class="n">key</span><span class="p">)</span>
<span class="k">try</span><span class="p">:</span>
<span class="k">yield</span> <span class="n">client</span>
<span class="k">finally</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">push</span><span class="p">(</span><span class="n">client</span><span class="o">=</span><span class="n">client</span><span class="p">)</span></div>
</div>
</pre></div>
</div>
</div>
<footer>
<hr/>
<div role="contentinfo">
<p>&#169; Copyright 2020, Eric McCarthy.</p>
</div>
Built with <a href="https://www.sphinx-doc.org/">Sphinx</a> using a
<a href="https://github.com/readthedocs/sphinx_rtd_theme">theme</a>
provided by <a href="https://readthedocs.org">Read the Docs</a>.
</footer>
</div>
</div>
</section>
</div>
<script>
jQuery(function () {
SphinxRtdTheme.Navigation.enable(true);
});
</script>
</body>
</html>