Implement retention policy

Author: andrewrjones

customers.contract.yml

@@ -3,6 +3,9 @@ owner:
  email: product-team@data-contracts.com
 description: All active customers of our product.
 version: 1
+retention:
+ policy: customer
+ timestamp: created
 
 columns:
  - name: id
@@ -28,5 +31,5 @@ columns:
 checks:
 - type: metric_expression
  metric: retention_policy
- expression_sql: COUNT(CASE WHEN created < NOW() - interval '3 year' THEN 1 END)
+ expression_sql: COUNT(CASE WHEN created < NOW() - interval '3 years' THEN 1 END)
  must_be: 0

retention.py

@@ -0,0 +1,31 @@
+#!/usr/bin/env python
+
+import yaml
+import psycopg2
+
+# Static lookup for the retention policy
+retention_policy_lookup = {
+ 'customer': '3 years',
+ 'employee': '5 years',
+}
+
+# Open and validate the data contract
+with open("./customers.contract.yml", "r") as stream:
+ contents = stream.read()
+ contract = yaml.safe_load(contents)
+
+if 'retention' not in contract or 'policy' not in contract['retention'] or 'timestamp' not in contract['retention']:
+ raise Exception(f"Retention policy not defined in the `{contract['dataset']}` contract")
+
+# Generate some SQL to apply the retention policy
+conn = psycopg2.connect("dbname=postgres user=postgres host='localhost' password='secret' port='5432'")
+cur = conn.cursor()
+cur.execute(f"""
+DELETE FROM {contract['dataset']}
+WHERE
+ {contract['retention']['timestamp']} < NOW() - interval '{retention_policy_lookup[contract['retention']['policy']]}'
+""")
+conn.commit()
+
+# Print if successful
+print(f"Retention policy applied to `{contract['dataset']}`")