Support use_pscopg2 flag for postgres integrations

Author: Seth Cleveland

manifests/integrations/postgres.pp

@@ -13,6 +13,9 @@
 #       The postgres port number
 #   $username
 #       The username for the datadog user
+#   $use_psycopg2
+#       Boolean to flag connecting to postgres with psycopg2 instead of pg8000.
+#       Warning, psycopg2 doesn't support ssl mode.
 #   $tags
 #       Optional array of tags
 #   $tables
@@ -54,6 +57,7 @@
   $dbname = 'postgres',
   $port   = '5432',
   $username = 'datadog',
+  $use_psycopg2 = false,
   $tags = [],
   $tables = [],
   $custom_metrics = {},
@@ -62,6 +66,7 @@
 
   validate_array($tags)
   validate_array($tables)
+  validate_bool($use_psycopg2)
 
   file { "${datadog_agent::params::conf_dir}/postgres.yaml":
     ensure  => file,

spec/classes/datadog_agent_integrations_postgres_spec.rb

@@ -35,10 +35,19 @@
       it { should contain_file(conf_file).with_content(%r{dbname: postgres}) }
       it { should contain_file(conf_file).with_content(%r{port: 5432}) }
       it { should contain_file(conf_file).with_content(%r{username: datadog}) }
+      it { should contain_file(conf_file).without_content(%r{^\s*use_psycopg2: }) }
       it { should contain_file(conf_file).without_content(%r{tags: })}
       it { should contain_file(conf_file).without_content(%r{^[^#]*relations: }) }
     end
 
+    context 'with use_psycopg2 enabled' do
+      let(:params) {{
+        use_psycopg2: true,
+        password: 'abc123',
+      }}
+      it { should contain_file(conf_file).with_content(%r{use_psycopg2: true}) }
+    end
+
     context 'with parameters set' do
       let(:params) {{
         host: 'postgres1',

templates/agent-conf.d/postgres.yaml.erb

@@ -6,6 +6,7 @@ instances:
 #       username: my_username
 #       password: my_password
 #       dbname: db_name
+#       use_psycopg2: false # Force using psycogp2 instead of pg8000 to connect. WARNING: psycopg2 doesn't support ssl mode.
 #       tags:
 #            - optional_tag1
 #            - optional_tag2
@@ -15,6 +16,9 @@ instances:
     username: <%= @username %>
     password: <%= @password %>
     dbname: <%= @dbname %>
+<% if @use_psycopg2 == true -%>
+    use_psycopg2: true
+<% end -%>
 <% if @tags and ! @tags.empty? -%>
     tags:
   <%- Array(@tags).each do |tag| -%>